import json
import logging

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from app.main import async_session
from app.models.database import ConversationSource, ConversationStatus
from app.services.conversation import ConversationManager

logger = logging.getLogger(__name__)
router = APIRouter()


@router.websocket("/ws/chat")
async def ws_chat(websocket: WebSocket):
    await websocket.accept()
    session_id = None

    try:
        async with async_session() as db:
            manager = ConversationManager(db)

            # Collect client info
            client_ip = websocket.client.host if websocket.client else ""
            forwarded = websocket.headers.get("x-forwarded-for", "")
            if forwarded:
                client_ip = forwarded.split(",")[0].strip()
            user_agent = websocket.headers.get("user-agent", "")
            client_info = {"ip": client_ip, "user_agent": user_agent}

            # First message: init or resume
            init_data = await websocket.receive_text()
            init_msg = json.loads(init_data)

            if "session_id" in init_msg:
                conv = await manager.get_by_session(init_msg["session_id"])
                if conv and conv.status == ConversationStatus.active:
                    session_id = conv.session_id
                    await websocket.send_text(json.dumps({
                        "type": "resumed",
                        "session_id": session_id,
                    }))
                else:
                    session_id = None

            # page_token: reuse active conversation from same page session
            if not session_id and "page_token" in init_msg and init_msg["page_token"]:
                conv = await manager.get_by_page_token(init_msg["page_token"])
                if conv:
                    session_id = conv.session_id
                    await websocket.send_text(json.dumps({
                        "type": "resumed",
                        "session_id": session_id,
                    }))

            if not session_id:
                source = ConversationSource(
                    init_msg.get("source", "standalone")
                )
                merged_info = {**client_info, **(init_msg.get("client_info") or {})}
                conv = await manager.create_conversation(
                    source=source,
                    client_info=merged_info,
                    page_token=init_msg.get("page_token"),
                )
                session_id = conv.session_id
                await websocket.send_text(json.dumps({
                    "type": "init",
                    "session_id": session_id,
                }))

            # Chat loop
            while True:
                data = await websocket.receive_text()
                msg = json.loads(data)
                content = msg.get("content", "").strip()
                if not content:
                    continue

                ended = False

                async for event in manager.stream_process_message(session_id, content):
                    if event["type"] == "error":
                        await websocket.send_text(json.dumps({
                            "type": "error",
                            "message": event["message"],
                        }))
                        break

                    await websocket.send_text(json.dumps(event, ensure_ascii=False))

                    if event["type"] == "evaluated" and event.get("ended"):
                        ended = True
                        break

                if ended:
                    break

    except WebSocketDisconnect:
        logger.info("Client disconnected: %s", session_id)
    except Exception:
        logger.exception("WebSocket error for session %s", session_id)