import base64
import json

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from sqlalchemy.orm import Session

from agent.bailian_tts import BailianTTS
from agent.chat_agent import ChatAgent
from database import SessionLocal
from presets.preset_store import PresetStore

router = APIRouter()


@router.websocket("/ws/chat")
async def websocket_chat(websocket: WebSocket):
    await websocket.accept()
    db: Session = SessionLocal()
    try:
        while True:
            raw = await websocket.receive_text()
            try:
                data = json.loads(raw)
            except json.JSONDecodeError:
                await websocket.send_json(
                    {"type": "error", "code": "INVALID_JSON", "message": "无效 JSON"}
                )
                continue

            msg_type = data.get("type")
            session_id = data.get("session_id")
            if not session_id:
                await websocket.send_json(
                    {
                        "type": "error",
                        "code": "MISSING_SESSION",
                        "message": "缺少 session_id",
                    }
                )
                continue

            user_input: str | None = None
            model: str | None = None
            reply_mode = data.get("reply_mode", "text")
            if reply_mode not in ("text", "voice", "both"):
                reply_mode = "text"
            show_text_stream = reply_mode in ("text", "both")
            need_voice = reply_mode in ("voice", "both")

            if msg_type == "chat":
                user_input = (data.get("message") or "").strip()
                if not user_input:
                    await websocket.send_json(
                        {
                            "type": "error",
                            "code": "EMPTY_MESSAGE",
                            "message": "消息不能为空",
                        }
                    )
                    continue

            elif msg_type == "preset_chat":
                preset_id = data.get("preset_id")
                if not preset_id:
                    await websocket.send_json(
                        {
                            "type": "error",
                            "code": "MISSING_PRESET",
                            "message": "缺少 preset_id",
                        }
                    )
                    continue
                preset_store = PresetStore(db)
                try:
                    preset, user_input = preset_store.render_prompt(
                        preset_id, data.get("variables") or {}
                    )
                    model = preset.model
                except LookupError:
                    await websocket.send_json(
                        {
                            "type": "error",
                            "code": "PRESET_NOT_FOUND",
                            "message": f"预设不存在: {preset_id}",
                        }
                    )
                    continue
                except ValueError as exc:
                    await websocket.send_json(
                        {
                            "type": "error",
                            "code": "MISSING_VARIABLES",
                            "message": str(exc),
                        }
                    )
                    continue
            else:
                await websocket.send_json(
                    {
                        "type": "error",
                        "code": "UNKNOWN_TYPE",
                        "message": f"未知消息类型: {msg_type}",
                    }
                )
                continue

            agent = ChatAgent(db)
            full_content = ""
            message_id = ""

            async def notify_search(query: str) -> None:
                await websocket.send_json({"type": "searching", "query": query})

            try:
                async for chunk in agent.stream_reply(
                    session_id, user_input, model, on_search=notify_search
                ):
                    if chunk.startswith("\0"):
                        message_id = chunk[1:]
                        continue
                    if chunk.startswith("\1SEARCH:"):
                        continue
                    full_content += chunk
                    if show_text_stream:
                        await websocket.send_json({"type": "token", "content": chunk})
            except LookupError:
                await websocket.send_json(
                    {
                        "type": "error",
                        "code": "SESSION_NOT_FOUND",
                        "message": "会话不存在",
                    }
                )
                continue
            except RuntimeError as exc:
                await websocket.send_json(
                    {"type": "error", "code": "CONFIG_ERROR", "message": str(exc)}
                )
                continue
            except Exception as exc:
                await websocket.send_json(
                    {
                        "type": "error",
                        "code": "API_ERROR",
                        "message": f"DeepSeek 请求失败: {exc}",
                    }
                )
                continue

            await websocket.send_json(
                {
                    "type": "done",
                    "message_id": message_id,
                    "full_content": full_content,
                    "user_message": user_input,
                }
            )

            if need_voice and full_content.strip():
                tts = BailianTTS()
                tts_voice_override = (data.get("tts_voice") or "").strip() or None
                if not tts.is_configured():
                    await websocket.send_json(
                        {
                            "type": "tts_error",
                            "message": "未配置阿里云百炼 API Key，无法语音回复",
                        }
                    )
                else:
                    await websocket.send_json({"type": "tts_start"})
                    try:
                        audio_fmt, audio_bytes = await tts.synthesize(
                            full_content, voice=tts_voice_override
                        )
                        await websocket.send_json(
                            {
                                "type": "audio",
                                "format": audio_fmt,
                                "data": base64.b64encode(audio_bytes).decode("ascii"),
                            }
                        )
                    except Exception as exc:
                        await websocket.send_json(
                            {"type": "tts_error", "message": str(exc)}
                        )
    except WebSocketDisconnect:
        pass
    finally:
        db.close()