import asyncio
import json
import logging
import time

from fastapi import WebSocket, WebSocketDisconnect

from agent.pipeline.customer_pipeline_langgraph import CustomerPipeline

from app.services.session_context import extract_protocol

logger = logging.getLogger(__name__)


class TextSessionPipeline:
    """
    Pipeline puramente texto:
      - envia "ready"
      - envia a abertura do agente via agent.start()
      - para cada mensagem de texto recebida, responde com agent.run(text)
    Protocolo sugerido (JSON):
      -> {"type":"start", "data": {...}}     # primeiro pacote
      -> {"type":"text", "text":"..."}      # mensagens do usuário
      -> {"type":"stop"}                     # encerra
    Eventos enviados pelo servidor:
      <- {"type":"ready", "session_id": "..."}
      <- {"type":"agent_start", "text": "..."}
      <- {"type":"agent_reply", "text": "..."}
      <- {"type":"bye"}
      <- {"type":"error","message":"..."}
    """
    def __init__(
        self,
        ws: WebSocket,
        start_payload: dict,
        session_id: str,
        *,
        session_data: dict,
        intro: str,
    ):
        self.ws = ws
        self.start_payload = start_payload or {}
        self.session_id = session_id
        self.session_data = dict(session_data or {})
        self.intro = (intro or "").strip()
        self.protocol = extract_protocol(self.start_payload, self.session_data)
        self.agent = CustomerPipeline(self.session_data)
        self._close = False

    async def send_json(self, payload: dict):
        await self.ws.send_text(json.dumps(payload, ensure_ascii=False))

    async def run(self):
        # 1) sinaliza pronto
        await self.send_json({"type": "ready", "session_id": self.session_id})

        # 2) mensagem inicial do agente
        stage = ""
        try:
            stage, opening = self.agent.start()
        except Exception as e:
            logger.exception("Erro no agent.start()")
            await self.send_json({"type": "error", "message": f"start_failed: {e}"})
            opening = None

        if not opening:
            opening = self.intro or None

        if opening:
            logger.info(
                "TEXT_AGENT_START | session_id=%s | stage=%s | text=%r",
                self.session_id,
                stage,
                opening,
            )
            await self.send_json({"type": "agent_start", "text": opening})

        self.agent.prepare(True, self.protocol)
        # 3) loop de interação
        while not self._close:
            try:
                msg = await asyncio.wait_for(self.ws.receive(), timeout=300)
            except asyncio.TimeoutError:
                # mantém conexão viva (opcional)
                try:
                    await self.ws.send_text("keepalive")
                except Exception:
                    pass
                continue
            except WebSocketDisconnect:
                logger.info(f"Cliente desconectou: {self.session_id}")
                break
            except Exception as e:
                logger.exception("Erro no loop text ws")
                await self.send_json({"type":"error", "message": str(e)})
                break

            if msg.get("type") == "websocket.disconnect":
                break

            if msg.get("type") != "websocket.receive":
                continue

            # Só tratamos conteúdo textual nessa rota
            raw_text = msg.get("text")
            if not raw_text:
                # ignora binários nesta rota
                continue

            # Aceita texto puro ou JSON
            try:
                data = json.loads(raw_text)
            except Exception:
                data = {"type": "text", "text": raw_text}

            mtype = data.get("type")
            if mtype in ("stop", "bye"):
                await self.send_json({"type": "bye"})
                break

            if mtype in ("text", "user_input", "message"):
                text = (data.get("text") or "").strip()
                if not text:
                    continue
                logger.info(
                    "TEXT_USER | session_id=%s | protocol=%s | text=%r",
                    self.session_id,
                    self.protocol,
                    text,
                )

                try:
                    agent_started = time.perf_counter()
                    stage, reply = self.agent.run(text)
                    agent_duration_ms = round((time.perf_counter() - agent_started) * 1000)
                    self.agent.set_interruption(was_interrupted=False)
                    #print(f"stage <> {stage}")
                    self._close = str(stage) in ("Concluído", "DONE", "DONE_CONCLUIDO")
                    

                except Exception as e:
                    logger.exception("Erro no agent.run()")
                    await self.send_json({"type": "error", "message": f"run_failed: {e}"})
                    continue

                logger.info(
                    "TEXT_AGENT_REPLY | session_id=%s | stage=%s | close=%s | duration_ms=%s | text=%r",
                    self.session_id,
                    stage,
                    self._close,
                    agent_duration_ms,
                    reply,
                )
                await self.send_json({"type": "agent_reply", "text": reply})
                if(self._close): break
            else:
                await self.send_json({"type": "error", "message": f"unknown_type:{mtype}"})
        try:
            await self.ws.close()
            output = await self.agent._end_service()
        except Exception:
            logger.exception("Erro ao encerrar pipeline de texto")
