# app/api/chat.py
import httpx
import json
from fastapi import APIRouter, Depends
from fastapi.responses import StreamingResponse
from app.models.schemas import ChatRequest, ChatResponse, ResetRequest
from app.services import validator, ollama
from app.core.auth import require_api_key, ClientInfo
from app.core.config import settings
from app.db import database as db
from app.core.context import SYSTEM_PROMPT

router = APIRouter(prefix="/api", tags=["Chat"])


@router.post("/chat/stream")
async def chat_stream(
    req: ChatRequest,
    client: ClientInfo = Depends(require_api_key),
):
    history = db.get_history(req.session_id, client.key[:8])
    max_msgs = settings.MAX_HISTORY * 2
    if len(history) > max_msgs:
        history = history[-max_msgs:]

    ok, fallback = validator.validate(req.message)
    if not ok:
        async def fallback_gen():
            yield f"data: {json.dumps({'token': fallback, 'done': True})}\n\n"
        return StreamingResponse(fallback_gen(), media_type="text/event-stream")

    system_prompt = SYSTEM_PROMPT

    messages = (
        [{"role": "system", "content": system_prompt}]
        + history
        + [{"role": "user", "content": req.message}]
    )

    payload = {
        "model":      settings.OLLAMA_MODEL,
        "messages":   messages,
        "stream":     True,
        "keep_alive": "10m",
        "options": {
            "temperature":    0.1,
            "num_predict":    300,
            "num_ctx":        2048,
            "top_k":          10,
            "repeat_penalty": 1.1,
        }
    }

    async def stream_gen():
        full_reply = []
        try:
            async with httpx.AsyncClient(timeout=120.0) as client_http:
                async with client_http.stream(
                    "POST", f"{settings.OLLAMA_URL}/api/chat", json=payload
                ) as resp:
                    async for line in resp.aiter_lines():
                        if not line:
                            continue
                        try:
                            data  = json.loads(line)
                            token = data.get("message", {}).get("content", "")
                            done  = data.get("done", False)
                            if token:
                                full_reply.append(token)
                            yield f"data: {json.dumps({'token': token, 'done': done})}\n\n"
                            if done:
                                break
                        except json.JSONDecodeError:
                            continue
        except Exception as e:
            yield f"data: {json.dumps({'token': f'Error: {e}', 'done': True})}\n\n"

        if full_reply:
            reply_text = "".join(full_reply)
            db.save_message(req.session_id, client.key[:8], "user",      req.message)
            db.save_message(req.session_id, client.key[:8], "assistant", reply_text)

    return StreamingResponse(stream_gen(), media_type="text/event-stream")


@router.post("/chat", response_model=ChatResponse)
async def chat_endpoint(
    req: ChatRequest,
    client: ClientInfo = Depends(require_api_key),
):
    ok, fallback = validator.validate(req.message)
    if not ok:
        return ChatResponse(reply=fallback, source="validator",
                            session_id=req.session_id, client=client.name)

    scoped_session = f"{client.key[:8]}_{req.session_id}"
    reply, source = await ollama.chat(req.message, scoped_session)
    return ChatResponse(reply=reply, source=source,
                        session_id=req.session_id, client=client.name)


@router.post("/reset")
async def reset_endpoint(
    req: ResetRequest,
    client: ClientInfo = Depends(require_api_key),
):
    db.clear_history(req.session_id, client.key[:8])
    return {"status": "ok", "session_id": req.session_id}