app.py

import json
import logging
import os
from queue import Empty, Queue
from threading import Lock

import requests as http_requests
from dotenv import load_dotenv
from flask import Flask, Response, jsonify, render_template, request

load_dotenv()

app = Flask(__name__)
app.logger.setLevel(logging.INFO)

CREWAI_ENTERPRISE_URL = os.environ["CREWAI_ENTERPRISE_URL"]
CREWAI_ENTERPRISE_TOKEN = os.environ["CREWAI_ENTERPRISE_TOKEN"]

sessions: dict[str, dict] = {}
sessions_lock = Lock()

sse_clients: dict[str, list[Queue]] = {}
sse_clients_lock = Lock()

# ──────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────


def _get_or_create_session(flow_id: str) -> dict:
    with sessions_lock:
        if flow_id not in sessions:
            sessions[flow_id] = {
                "stage": "processing",
                "status_message": "Starting support ticket flow...",
                "original_email": None,
                "draft_reply": None,
                "keep_processing": True,
                "end_of_conversation": False,
                "pending_feedback": None,
                "result": None,
            }
        return sessions[flow_id]


def _notify_sse_clients(flow_id: str):
    with sse_clients_lock:
        for q in sse_clients.get(flow_id, []):
            q.put(True)


# ──────────────────────────────────────────────
# Pages
# ──────────────────────────────────────────────


@app.route("/")
def index():
    return render_template("index.html")


# ──────────────────────────────────────────────
# API routes (called by the browser)
# ──────────────────────────────────────────────


@app.route("/api/warmup", methods=["POST"])
def api_warmup():
    """Ping CrewAI Enterprise /inputs to warm up the deployment."""
    try:
        resp = http_requests.get(
            f"{CREWAI_ENTERPRISE_URL}/inputs",
            headers={"Authorization": f"Bearer {CREWAI_ENTERPRISE_TOKEN}"},
            timeout=45,
        )
        resp.raise_for_status()
        return jsonify(resp.json())
    except http_requests.RequestException as exc:
        app.logger.warning("Warmup request failed: %s", exc)
        return jsonify({"error": "warmup failed"}), 502


@app.route("/api/start", methods=["POST"])
def api_start():
    """Start the support ticket flow by calling CrewAI Enterprise /kickoff."""
    kickoff_payload = {"inputs": {}}

    try:
        resp = http_requests.post(
            f"{CREWAI_ENTERPRISE_URL}/kickoff",
            json=kickoff_payload,
            headers={
                "Authorization": f"Bearer {CREWAI_ENTERPRISE_TOKEN}",
                "Content-Type": "application/json",
            },
            timeout=30,
        )
        resp.raise_for_status()
    except http_requests.RequestException as exc:
        app.logger.error("Kickoff request failed: %s", exc)
        return jsonify({"error": "Failed to start support ticket flow"}), 502

    data = resp.json()
    kickoff_id = data.get("kickoff_id", data.get("id", ""))

    if kickoff_id:
        _get_or_create_session(kickoff_id)

    return jsonify({"kickoff_id": kickoff_id, "raw": data})


@app.route("/api/stream/<flow_id>")
def api_stream(flow_id: str):
    """SSE stream that pushes ticket state whenever it changes."""

    def _snapshot() -> str:
        session = sessions.get(flow_id)
        if session is None:
            return json.dumps(
                {
                    "stage": "unknown",
                    "status_message": "",
                    "original_email": None,
                    "draft_reply": None,
                    "keep_processing": False,
                    "end_of_conversation": False,
                    "pending_feedback": None,
                    "result": None,
                }
            )
        with sessions_lock:
            return json.dumps(
                {
                    "stage": session["stage"],
                    "status_message": session["status_message"],
                    "original_email": session["original_email"],
                    "draft_reply": session["draft_reply"],
                    "keep_processing": session["keep_processing"],
                    "end_of_conversation": session["end_of_conversation"],
                    "pending_feedback": session["pending_feedback"],
                    "result": session["result"],
                }
            )

    def generate():
        q: Queue = Queue()
        with sse_clients_lock:
            sse_clients.setdefault(flow_id, []).append(q)
        try:
            yield f"data: {_snapshot()}\n\n"
            while True:
                try:
                    q.get(timeout=30)
                except Empty:
                    yield ": keepalive\n\n"
                    continue
                yield f"data: {_snapshot()}\n\n"
        finally:
            with sse_clients_lock:
                clients = sse_clients.get(flow_id, [])
                if q in clients:
                    clients.remove(q)

    return Response(
        generate(),
        mimetype="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


@app.route("/api/feedback/<flow_id>", methods=["POST"])
def api_feedback(flow_id: str):
    """User submits approve/reject; Flask forwards it to the callback_url."""
    session = sessions.get(flow_id)
    if not session or not session.get("pending_feedback"):
        return jsonify({"error": "No pending feedback for this session"}), 404

    body = request.get_json(force=True)
    feedback_text = body.get("feedback", "")

    callback_url = session["pending_feedback"]["callback_url"]

    try:
        resp = http_requests.post(
            callback_url,
            json={"feedback": feedback_text, "source": "support_tickets_ui"},
            headers={"Content-Type": "application/json"},
            timeout=30,
        )
        resp.raise_for_status()
    except http_requests.RequestException as exc:
        app.logger.error("Feedback callback failed: %s", exc)
        return jsonify({"error": "Failed to submit feedback"}), 502

    with sessions_lock:
        session["pending_feedback"] = None
        session["stage"] = "processing"
        session["status_message"] = "Processing your decision..."
        session["keep_processing"] = True

    _notify_sse_clients(flow_id)

    return jsonify({"ok": True})


# ──────────────────────────────────────────────
# Webhook routes (called by CrewAI Enterprise)
# ──────────────────────────────────────────────


@app.route("/webhook/messages", methods=["POST"])
def webhook_messages():
    """Receive status/message events from the DispatcherEventBusService."""
    payload = request.get_json(force=True)
    app.logger.info(
        "Message webhook raw payload: %s", json.dumps(payload, default=str)[:2000]
    )

    result = payload.get("result") or {}
    if not isinstance(result, dict):
        return jsonify({"ok": True, "skipped": True})

    flow_id = (
        payload.get("source_fingerprint")
        or (payload.get("fingerprint_metadata") or {}).get("id", "")
        or payload.get("flow_id", "")
    )

    if not flow_id:
        app.logger.warning("Message webhook: could not extract flow_id")
        return jsonify({"ok": True, "skipped": True})

    session = _get_or_create_session(flow_id)

    # ── Handle structured stage events from emit_status() ──
    stage = result.get("stage")
    data = result.get("data") or {}

    if stage:
        keep_processing = result.get("keep_processing", True)
        end_of_conversation = result.get("end_of_conversation", False)

        with sessions_lock:
            session["stage"] = stage
            session["keep_processing"] = keep_processing
            if end_of_conversation:
                session["end_of_conversation"] = True

            # Update session fields based on stage
            if stage == "email_fetched":
                session["original_email"] = {
                    "subject": data.get("subject", ""),
                    "from": data.get("from", ""),
                    "body": data.get("body", ""),
                }
                session["status_message"] = "Email fetched. Analyzing and drafting response..."

            elif stage == "draft_ready":
                session["draft_reply"] = data.get("draft", "")
                # Also set original_email if provided and not yet set
                if not session["original_email"] and data.get("from"):
                    session["original_email"] = {
                        "subject": data.get("subject", ""),
                        "from": data.get("from", ""),
                        "body": data.get("body", ""),
                    }
                session["status_message"] = "Draft reply ready for your review."

            elif stage == "sent":
                session["result"] = "sent"
                session["status_message"] = data.get("step", "Email sent successfully.")

            elif stage == "rejected":
                session["result"] = "rejected"
                session["status_message"] = data.get("step", "Reply rejected. No email sent.")

            elif stage == "no_email":
                session["result"] = "no_email"
                session["status_message"] = data.get("step", "No support email found.")

            elif stage == "error":
                session["result"] = "error"
                session["status_message"] = data.get("step", "An error occurred.")

            elif stage == "processing":
                session["status_message"] = data.get("step", "Processing...")

        _notify_sse_clients(flow_id)

        app.logger.info(
            "Status webhook: flow=%s stage=%s keep_processing=%s end=%s",
            flow_id, stage, keep_processing, end_of_conversation,
        )
        return jsonify({"ok": True})

    # ── Handle plain message events (fallback, same as concierge) ──
    message_data = result.get("message")
    if message_data:
        if isinstance(message_data, str):
            content = message_data
        else:
            role = message_data.get("role", "assistant")
            content = message_data.get("content", "")
            if role == "user":
                return jsonify({"ok": True})

        keep_processing = result.get("keep_processing", True)
        end_of_conversation = result.get("end_of_conversation", False)

        with sessions_lock:
            session["status_message"] = content
            session["keep_processing"] = keep_processing
            if end_of_conversation:
                session["end_of_conversation"] = True

        _notify_sse_clients(flow_id)

        app.logger.info(
            "Message webhook: flow=%s content=%.80s keep_processing=%s",
            flow_id, content, keep_processing,
        )

    return jsonify({"ok": True})


@app.route("/webhook/feedback", methods=["POST"])
def webhook_feedback():
    """Receive human feedback requests from CrewAI Enterprise."""
    payload = request.get_json(force=True)
    app.logger.info(
        "Feedback webhook raw payload: %s", json.dumps(payload, default=str)[:2000]
    )

    flow_id = payload.get("flow_id", "")
    if not flow_id:
        return jsonify({"error": "Missing flow_id"}), 400

    callback_url = payload.get("callback_url", "")
    emit_options = payload.get("emit", [])
    output = payload.get("output", "")
    message = payload.get("message", "")

    session = _get_or_create_session(flow_id)
    with sessions_lock:
        session["stage"] = "waiting_for_review"
        session["status_message"] = message or "Please review the proposed reply."
        session["keep_processing"] = False
        session["pending_feedback"] = {
            "emit": emit_options,
            "callback_url": callback_url,
            "message": message,
            "output": output,
        }

    _notify_sse_clients(flow_id)

    app.logger.info("Feedback webhook: flow=%s emit=%s", flow_id, emit_options)
    return jsonify({"ok": True})


# ──────────────────────────────────────────────
# Entry point
# ──────────────────────────────────────────────

if __name__ == "__main__":
    app.run(debug=True, port=5002)