import asyncio
import json
import logging
import time

from fastapi import WebSocket, WebSocketDisconnect

from agent.pipeline.pipeline_streaming_text import CustomerPipeline
from app.services.session_context import extract_protocol

logger = logging.getLogger(__name__)


class TextSessionPipelineStream:
    """
    Pipeline para comunicação texto via WebSocket com streaming token a token.

    Protocolo sugerido (JSON):
      -> {"type":"start", "data": {...}}
      -> {"type":"text", "text":"..."}
      -> {"type":"stop"}

    Eventos enviados pelo servidor:
      <- {"type":"ready", "session_id": "..."}
      <- {"type":"agent_start", "text": "..."}
      <- {"type":"agent_reply_chunk", "delta": "token"}
      <- {"type":"agent_reply_end", "text": "resposta completa"}
      <- {"type":"bye"}
      <- {"type":"error","message":"..."}
    """

    def __init__(
        self,
        ws: WebSocket,
        start_payload: dict,
        session_id: str,
        *,
        session_data: dict | None = None,
        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._close = False

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

    # ============================================================
    #  STREAMING TOKEN-A-TOKEN
    # ============================================================

    async def _to_async_gen(self, generator):
        """Transforma generator síncrono em assíncrono."""
        for item in generator:
            yield item

    async def stream_reply(self, text: str):
        """
        Envia resposta token-a-token via WebSocket:
        - {"type":"agent_reply_chunk", "delta": token}
        - {"type":"agent_reply_end", "text": full_response}
        """
        try:
            agent_started = time.perf_counter()
            token_generator = self.agent.run(text)
        except Exception as e:
            await self.send_json({"type": "error", "message": f"stream_init_failed: {e}"})
            return

        full_text = ""

        try:
            async for token in self._to_async_gen(token_generator):
                full_text += token
                await self.send_json({"type": "agent_reply_chunk", "delta": token})
                await asyncio.sleep(0)
        except Exception as e:
            await self.send_json({"type": "error", "message": f"stream_failed: {e}"})
            return

        # Finaliza streaming
        agent_duration_ms = round((time.perf_counter() - agent_started) * 1000)
        logger.info(
            "TEXT_STREAM_AGENT_REPLY | session_id=%s | duration_ms=%s | text=%r",
            self.session_id,
            agent_duration_ms,
            full_text,
        )
        await self.send_json({
            "type": "agent_reply_end",
            "text": full_text
        })

    # ============================================================

    async def run(self):
        # 1) envia ready
        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_STREAM_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 WS
        while not self._close:
            try:
                msg = await asyncio.wait_for(self.ws.receive(), timeout=300)
            except asyncio.TimeoutError:
                # ping 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

            raw_text = msg.get("text")
            if not raw_text:
                continue

            # aceita JSON ou texto puro
            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_STREAM_USER | session_id=%s | protocol=%s | text=%r",
                    self.session_id,
                    self.protocol,
                    text,
                )

                # STREAMING TOKEN A TOKEN
                await self.stream_reply(text)

                # depois do streaming, atualiza stage
                try:
                    self._close = str(self.agent.stage) in ("Concluído", "DONE", "DONE_CONCLUIDO")
                except Exception as e:
                    await self.send_json({"type": "error", "message": f"run_failed: {e}"})

            else:
                await self.send_json({"type": "error", "message": f"unknown_type:{mtype}"})

        try:
            await self.ws.close()
        except Exception:
            pass
