# api/v1/chat.py
import httpx
import json
import asyncio
from fastapi import APIRouter, HTTPException, Query
from fastapi.responses import StreamingResponse
from typing import List, Dict, Any, Optional
from models.chat import ChatRequest, ChatResponse
from services import redis_service  # now using updated service with session support
from utils.logging import logger
from config import settings

router = APIRouter()

#MAX_HISTORY_LENGTH = 50
MAX_HISTORY_LENGTH = 20


@router.post("/chat", response_model=ChatResponse)
async def chat_endpoint(payload: ChatRequest):
    try:
        # Create a new session if session_id not provided
        session_id = payload.session_id
        if not session_id:
            meta = redis_service.create_session(payload.user_id, payload.message)
            session_id = meta["session_id"]

        # Save user message
        redis_service.save_chat(payload.user_id, session_id, {"role": "user", "content": payload.message})
        history = redis_service.get_chat(payload.user_id, session_id, limit=MAX_HISTORY_LENGTH)

        async with httpx.AsyncClient(timeout=settings.REQUEST_TIMEOUT) as client:
            resp = await client.post(
                settings.LM_STUDIO_URL,
                json={"model": settings.MODEL_NAME, "messages": history},
            )
            resp.raise_for_status()
            data = resp.json()

        reply = data["choices"][0]["message"]["content"]

        # Save assistant message
        redis_service.save_chat(payload.user_id, session_id, {"role": "assistant", "content": reply})

        # Return normal ChatResponse, but could also include session_id if needed
        return ChatResponse(response=reply, session_id=session_id)

    except Exception:
        logger.exception("Error in /chat endpoint")
        raise HTTPException(status_code=500, detail="Internal server error")


@router.post("/chat-stream")
async def chat_stream_endpoint(payload: ChatRequest):
    """
    Streams model output token-by-token using SSE.
    """
    session_id = payload.session_id
    if not session_id:
        meta = redis_service.create_session(payload.user_id, payload.message)
        session_id = meta["session_id"]

    redis_service.save_chat(payload.user_id, session_id, {"role": "user", "content": payload.message})
    history = redis_service.get_chat(payload.user_id, session_id, limit=MAX_HISTORY_LENGTH)

    async def event_generator():
        assistant_text = ""
        try:
            async with httpx.AsyncClient(timeout=None) as client:
                async with client.stream(
                    "POST",
                    settings.LM_STUDIO_URL,
                    json={
                        "model": settings.MODEL_NAME,
                        "messages": history,
                        "stream": True
                    }
                ) as r:
                    async for raw_line in r.aiter_lines():
                        if not raw_line:
                            continue

                        line = raw_line if raw_line.startswith("data:") else f"data: {raw_line}"
                        payload_str = line[len("data: "):].strip()

                        if payload_str == "[DONE]":
                            yield "data: [DONE]\n\n"
                            break

                        yield f"data: {payload_str}\n\n"

                        try:
                            obj = json.loads(payload_str)
                            choice = obj.get("choices", [{}])[0]
                            delta = choice.get("delta", {})
                            piece = delta.get("content") or choice.get("text")
                            if piece:
                                assistant_text += piece
                        except json.JSONDecodeError:
                            pass

                        await asyncio.sleep(0)
        except Exception as e:
            logger.exception("Streaming error in /chat-stream")
            yield f"event: error\ndata: {str(e)}\n\n"
        finally:
            if assistant_text:
                redis_service.save_chat(payload.user_id, session_id, {"role": "assistant", "content": assistant_text})

    headers = {
        "Cache-Control": "no-cache",
        "Connection": "keep-alive",
        "X-Accel-Buffering": "no"
    }
    return StreamingResponse(event_generator(), media_type="text/event-stream", headers=headers)


@router.get("/history")
async def get_history(
    user_id: str = Query(..., description="User ID"),
    session_id: str = Query(..., description="Session ID"),
    limit: int = Query(MAX_HISTORY_LENGTH, description="Max number of messages to return")
) -> List[Dict[str, Any]]:
    """
    Return all history saved for a given user/session.
    """
    logger.info(f"[GET /history] user_id={user_id}, session_id={session_id}, limit={limit}")
    history = redis_service.get_chat(user_id, session_id, limit=limit)
    return history or []


@router.delete("/history")
async def delete_history(
    user_id: str = Query(..., description="User ID"),
    session_id: str = Query(..., description="Session ID")
):
    """
    Clears history for a given user/session.
    """
    logger.info(f"[DELETE /history] user_id={user_id}, session_id={session_id}")
    redis_service.clear_chat(user_id, session_id)
    return {"status": "cleared"}