"""HTTP + WebSocket API server for gateway integrations."""

import asyncio
import json
import uuid
from typing import Any, Callable, Awaitable

import aiohttp
from aiohttp import web
from loguru import logger

from flowly.session.manager import SessionManager

# Maximum allowed request body size (1MB)
_MAX_BODY_SIZE = 1024 * 1024

# Type alias for the chat callback used by the /ws endpoint.
# Signature: (session_key, message, run_id, stream_callback) -> response_text
ChatCallback = Callable[
    [str, str, str, Callable[[str], Awaitable[None]] | None],
    Awaitable[str],
]


class GatewayServer:
    """
    HTTP + WebSocket API server for gateway integrations.

    Provides endpoints for:
    - Health check (GET /health)
    - Voice message handling (POST /api/voice/message)
    - Cron triggers (POST /api/cron/run, /api/cron/reload)
    - Desktop WebSocket chat (GET /ws)
    """

    def __init__(
        self,
        host: str = "127.0.0.1",
        port: int = 18790,
        on_voice_message: Callable[[str, str, str], Awaitable[str]] | None = None,
        on_cron_run: Callable[[str, bool], Awaitable[bool]] | None = None,
        on_cron_reload: Callable[[], Awaitable[int]] | None = None,
        on_chat_message: ChatCallback | None = None,
        sessions: SessionManager | None = None,
    ):
        self.host = host
        self.port = port
        self.on_voice_message = on_voice_message
        self.on_cron_run = on_cron_run
        self.on_cron_reload = on_cron_reload
        self.on_chat_message = on_chat_message
        self.sessions = sessions
        self._app: web.Application | None = None
        self._runner: web.AppRunner | None = None
        self._site: web.TCPSite | None = None
        # Track active WebSocket clients and their running tasks for abort support.
        self._ws_clients: dict[str, web.WebSocketResponse] = {}
        self._active_tasks: dict[str, asyncio.Task] = {}
        # Tick task for periodic health pings to connected clients.
        self._tick_task: asyncio.Task | None = None

    def _create_app(self) -> web.Application:
        """Create the aiohttp application."""
        app = web.Application(client_max_size=_MAX_BODY_SIZE)
        app.router.add_get("/health", self._handle_health)
        if self.on_voice_message:
            app.router.add_post("/api/voice/message", self._handle_voice_message)
        app.router.add_post("/api/cron/run", self._handle_cron_run)
        app.router.add_post("/api/cron/reload", self._handle_cron_reload)
        # Desktop WebSocket endpoint — active when chat callback is provided.
        if self.on_chat_message:
            app.router.add_get("/ws", self._handle_ws)
        return app

    # ------------------------------------------------------------------
    # HTTP handlers (unchanged)
    # ------------------------------------------------------------------

    async def _handle_health(self, request: web.Request) -> web.Response:
        return web.json_response({"status": "ok"})

    async def _handle_cron_run(self, request: web.Request) -> web.Response:
        try:
            data = await request.json()
            job_id = data.get("job_id", "")
            force = bool(data.get("force", False))
            if not job_id or len(job_id) > 256:
                return web.json_response({"error": "Invalid job_id"}, status=400)
            if not self.on_cron_run:
                return web.json_response({"error": "Cron handler not configured"}, status=500)
            success = await self.on_cron_run(job_id, force)
            if success:
                return web.json_response({"ok": True})
            else:
                return web.json_response({"error": f"Job '{job_id}' not found or disabled"}, status=404)
        except json.JSONDecodeError:
            return web.json_response({"error": "Invalid JSON"}, status=400)
        except Exception as e:
            logger.error(f"Error triggering cron job: {e}")
            return web.json_response({"error": "Internal server error"}, status=500)

    async def _handle_cron_reload(self, request: web.Request) -> web.Response:
        try:
            if not self.on_cron_reload:
                return web.json_response({"error": "Cron reload not configured"}, status=500)
            count = await self.on_cron_reload()
            return web.json_response({"ok": True, "jobs": count})
        except Exception as e:
            logger.error(f"Error reloading cron jobs: {e}")
            return web.json_response({"error": "Internal server error"}, status=500)

    async def _handle_voice_message(self, request: web.Request) -> web.Response:
        try:
            data = await request.json()
            call_sid = data.get("call_sid", "")
            from_number = data.get("from", "")
            text = data.get("text", "")
            if not call_sid or len(call_sid) > 128:
                return web.json_response({"error": "Invalid call_sid"}, status=400)
            if len(text) > 50000:
                return web.json_response({"error": "Message too large"}, status=413)
            logger.info(f"Voice message from {from_number}: {text[:50]}...")
            if not self.on_voice_message:
                return web.json_response({"error": "Voice handler not configured"}, status=500)
            response = await asyncio.wait_for(
                self.on_voice_message(call_sid, from_number, text),
                timeout=30.0,
            )
            return web.json_response({"response": response})
        except json.JSONDecodeError:
            return web.json_response({"error": "Invalid JSON"}, status=400)
        except asyncio.TimeoutError:
            logger.error("Voice message handler timeout")
            return web.json_response({"error": "Handler timeout"}, status=504)
        except Exception as e:
            logger.error(f"Error handling voice message: {e}")
            return web.json_response({"error": "Internal server error"}, status=500)

    # ------------------------------------------------------------------
    # WebSocket handler — desktop chat protocol
    # ------------------------------------------------------------------

    async def _handle_ws(self, request: web.Request) -> web.WebSocketResponse:
        """Handle a desktop WebSocket connection."""
        ws = web.WebSocketResponse(heartbeat=30.0)
        await ws.prepare(request)

        client_id = str(uuid.uuid4())
        self._ws_clients[client_id] = ws
        logger.info(f"[WS] Desktop client connected: {client_id}")

        try:
            async for raw_msg in ws:
                if raw_msg.type == aiohttp.WSMsgType.TEXT:
                    try:
                        data = json.loads(raw_msg.data)
                    except json.JSONDecodeError:
                        await self._ws_send(ws, {"type": "error", "message": "Invalid JSON"})
                        continue
                    msg_type = data.get("type")
                    if msg_type == "rpc":
                        await self._handle_ws_rpc(ws, client_id, data)
                    elif msg_type == "ping":
                        await self._ws_send(ws, {"type": "pong", "timestamp": data.get("timestamp")})
                elif raw_msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSE):
                    break
        except Exception as e:
            logger.error(f"[WS] Client {client_id} error: {e}")
        finally:
            self._ws_clients.pop(client_id, None)
            logger.info(f"[WS] Desktop client disconnected: {client_id}")

        return ws

    async def _handle_ws_rpc(self, ws: web.WebSocketResponse, client_id: str, data: dict) -> None:
        """Dispatch an RPC call to the appropriate handler."""
        method = data.get("method", "")
        rpc_id = data.get("id", "")
        params = data.get("params") or {}

        try:
            if method == "health":
                await self._ws_rpc_reply(ws, rpc_id, {"ok": True})

            elif method == "sessions.list":
                await self._ws_rpc_sessions_list(ws, rpc_id, params)

            elif method == "chat.history":
                await self._ws_rpc_chat_history(ws, rpc_id, params)

            elif method == "chat.send":
                await self._ws_rpc_chat_send(ws, client_id, rpc_id, params)

            elif method == "chat.abort":
                await self._ws_rpc_chat_abort(ws, rpc_id, params)

            else:
                await self._ws_rpc_error(ws, rpc_id, "INVALID_REQUEST", f"Unknown method: {method}")
        except Exception as e:
            logger.error(f"[WS] RPC {method} error: {e}")
            await self._ws_rpc_error(ws, rpc_id, "UNAVAILABLE", str(e))

    # --- RPC: sessions.list ---

    async def _ws_rpc_sessions_list(self, ws: web.WebSocketResponse, rpc_id: str, params: dict) -> None:
        if not self.sessions:
            await self._ws_rpc_reply(ws, rpc_id, {"sessions": []})
            return

        raw_sessions = self.sessions.list_sessions()
        sessions = []
        for s in raw_sessions:
            key = s.get("key", "")
            sessions.append({
                "key": key,
                "displayName": key.split(":", 1)[-1] if ":" in key else key,
                "updatedAt": s.get("updated_at"),
            })

        limit = params.get("limit")
        if limit and isinstance(limit, int):
            sessions = sessions[:limit]

        await self._ws_rpc_reply(ws, rpc_id, {"sessions": sessions})

    # --- RPC: chat.history ---

    async def _ws_rpc_chat_history(self, ws: web.WebSocketResponse, rpc_id: str, params: dict) -> None:
        session_key = params.get("sessionKey", "")
        if not session_key or not self.sessions:
            await self._ws_rpc_reply(ws, rpc_id, {"sessionKey": session_key, "messages": []})
            return

        session = self.sessions.get_or_create(session_key)
        # Return full messages (not just role/content) for richer UI rendering.
        messages = []
        for m in session.messages:
            msg: dict[str, Any] = {"role": m["role"]}
            # Normalise content into array format expected by the iOS/desktop protocol.
            content_raw = m.get("content", "")
            if isinstance(content_raw, str):
                msg["content"] = [{"type": "text", "text": content_raw}]
            elif isinstance(content_raw, list):
                msg["content"] = content_raw
            else:
                msg["content"] = [{"type": "text", "text": str(content_raw)}]
            if "timestamp" in m:
                msg["timestamp"] = m["timestamp"]
            if "usage" in m:
                msg["usage"] = m["usage"]
            messages.append(msg)

        await self._ws_rpc_reply(ws, rpc_id, {
            "sessionKey": session_key,
            "sessionId": session_key,
            "messages": messages,
            "thinkingLevel": session.metadata.get("thinkingLevel"),
        })

    # --- RPC: chat.send ---

    async def _ws_rpc_chat_send(
        self, ws: web.WebSocketResponse, client_id: str, rpc_id: str, params: dict,
    ) -> None:
        message = params.get("message", "")
        if not message:
            await self._ws_rpc_error(ws, rpc_id, "INVALID_REQUEST", "Empty message")
            return

        session_key = params.get("sessionKey") or f"desktop:{client_id}"
        idempotency_key = params.get("idempotencyKey") or str(uuid.uuid4())
        run_id = idempotency_key

        # ACK immediately with runId so the client can track the run.
        await self._ws_rpc_reply(ws, rpc_id, {"runId": run_id, "status": "accepted"})

        # Build streaming callback that pushes delta events to this client.
        async def stream_callback(delta: str) -> None:
            try:
                await self._ws_send(ws, {
                    "type": "event",
                    "event": "agent",
                    "data": {
                        "runId": run_id,
                        "stream": "assistant",
                        "data": {"text": delta},
                    },
                })
            except Exception:
                pass  # Client may have disconnected mid-stream.

        # Process in background so we don't block the recv loop.
        task = asyncio.create_task(
            self._run_chat(ws, client_id, session_key, message, run_id, stream_callback)
        )
        self._active_tasks[run_id] = task
        task.add_done_callback(lambda _: self._active_tasks.pop(run_id, None))

    async def _run_chat(
        self,
        ws: web.WebSocketResponse,
        client_id: str,
        session_key: str,
        message: str,
        run_id: str,
        stream_callback: Callable[[str], Awaitable[None]],
    ) -> None:
        """Execute the chat and send final/error events."""
        accumulated_text = ""

        # Wrap the stream callback to accumulate full text for the final event.
        async def tracking_callback(delta: str) -> None:
            nonlocal accumulated_text
            accumulated_text += delta
            await stream_callback(delta)

        try:
            assert self.on_chat_message is not None
            response = await self.on_chat_message(session_key, message, run_id, tracking_callback)

            # Send final chat event.
            await self._ws_send(ws, {
                "type": "event",
                "event": "chat",
                "data": {
                    "state": "final",
                    "runId": run_id,
                    "sessionKey": session_key,
                    "message": {
                        "role": "assistant",
                        "content": [{"type": "text", "text": response}],
                    },
                },
            })
        except asyncio.CancelledError:
            await self._ws_send(ws, {
                "type": "event",
                "event": "chat",
                "data": {"state": "aborted", "runId": run_id, "sessionKey": session_key},
            })
        except Exception as e:
            logger.error(f"[WS] chat.send run {run_id} failed: {e}")
            await self._ws_send(ws, {
                "type": "event",
                "event": "chat",
                "data": {
                    "state": "error",
                    "runId": run_id,
                    "sessionKey": session_key,
                    "errorMessage": str(e),
                },
            })

    # --- RPC: chat.abort ---

    async def _ws_rpc_chat_abort(self, ws: web.WebSocketResponse, rpc_id: str, params: dict) -> None:
        run_id = params.get("runId", "")
        task = self._active_tasks.get(run_id)
        if task and not task.done():
            task.cancel()
        await self._ws_rpc_reply(ws, rpc_id, {"ok": True})

    # ------------------------------------------------------------------
    # WebSocket helpers
    # ------------------------------------------------------------------

    async def _ws_send(self, ws: web.WebSocketResponse, data: dict) -> None:
        """Send JSON to a WebSocket client, silently ignoring closed connections."""
        try:
            if not ws.closed:
                await ws.send_json(data)
        except (ConnectionResetError, RuntimeError):
            pass

    async def _ws_rpc_reply(self, ws: web.WebSocketResponse, rpc_id: str, result: Any) -> None:
        await self._ws_send(ws, {"type": "rpc", "id": rpc_id, "result": result})

    async def _ws_rpc_error(self, ws: web.WebSocketResponse, rpc_id: str, code: str, message: str) -> None:
        await self._ws_send(ws, {
            "type": "rpc",
            "id": rpc_id,
            "error": {"code": code, "message": message},
        })

    # ------------------------------------------------------------------
    # Tick loop — periodic health pings to connected WS clients
    # ------------------------------------------------------------------

    async def _tick_loop(self) -> None:
        """Send periodic tick events to all connected desktop clients."""
        while True:
            try:
                await asyncio.sleep(10)
                if not self._ws_clients:
                    continue
                tick = {"type": "event", "event": "tick"}
                for ws in list(self._ws_clients.values()):
                    await self._ws_send(ws, tick)
            except asyncio.CancelledError:
                break
            except Exception:
                pass

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def start(self) -> None:
        """Start the HTTP + WebSocket server."""
        self._app = self._create_app()
        self._runner = web.AppRunner(self._app)
        await self._runner.setup()
        self._site = web.TCPSite(self._runner, self.host, self.port)
        await self._site.start()
        if self.on_chat_message:
            self._tick_task = asyncio.create_task(self._tick_loop())
        logger.info(f"Gateway API listening on http://{self.host}:{self.port}")
        if self.on_chat_message:
            logger.info(f"Desktop WebSocket available at ws://{self.host}:{self.port}/ws")

    async def stop(self) -> None:
        """Stop the server and clean up."""
        if self._tick_task:
            self._tick_task.cancel()
            try:
                await self._tick_task
            except asyncio.CancelledError:
                pass
        # Cancel any in-flight chat tasks.
        for task in self._active_tasks.values():
            task.cancel()
        # Close all WebSocket connections.
        for ws in list(self._ws_clients.values()):
            await ws.close()
        self._ws_clients.clear()
        if self._site:
            await self._site.stop()
        if self._runner:
            await self._runner.cleanup()
        logger.info("Gateway API stopped")
