import logging
import uuid
import json

from typing import Any, Dict
from fastapi import APIRouter, HTTPException, status
from fastapi.responses import StreamingResponse

from models import ChatRequest
from utils import convert_prompt_to_messages, standard_sse_generator
from config import GLM_API_KEY, GLM_UPSTREAM_URL

router = APIRouter()
logger = logging.getLogger("proxy")


@router.post("/glm/chat")
async def proxy_glm_chat(request: ChatRequest):

    request_id = f"glm-{uuid.uuid4()}"

    if not GLM_API_KEY:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="GLM API key is not configured.",
        )

    client_request_body = request.model_dump()
    logger.debug(
        f"[{request_id}] === GLM New Request ===\n"
        f"[{request_id}] Client Request Body:\n"
        f"{json.dumps(client_request_body, indent=2, ensure_ascii=False)}"
    )

    messages = convert_prompt_to_messages(request.prompt)
    logger.debug(
        f"[{request_id}] Agent Transformed Messages:\n"
        f"{json.dumps(messages, indent=2, ensure_ascii=False)}"
    )

    # 基础的请求体
    payload: Dict[str, Any] = {
        "model": "glm-4-flash",
        "messages": messages,
        "stream": True,
        "temperature": request.temperature,
        "top_p": request.top_p,
    }

    # 如果请求中指定了 max_tokens ,则添加到 payload
    if request.max_tokens:
        payload["max_tokens"] = request.max_tokens

    # 如果客户端请求 web_search=True，则为 GLM 启用网页搜索工具
    if request.web_search is True:
        payload["tools"] = [{"type": "web_search", "web_search": {"enable": True}}]
        logger.info(f"[{request_id}] Enabled web search for GLM.")

    logger.debug(
        f"[{request_id}] Upstream Payload:\n"
        f"{json.dumps(payload, indent=2, ensure_ascii=False)}"
    )

    headers = {"Authorization": f"Bearer {GLM_API_KEY}"}
    return StreamingResponse(
        standard_sse_generator(GLM_UPSTREAM_URL, payload, headers, "GLM", request_id),
        media_type="text/event-stream",
    )