import asyncio
import json
import logging
from typing import Optional

from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from app.service.rag_pipeline import get_answer_from_rag, get_answer_stream_from_rag
from app.service.session_manager import session_manager

logger = logging.getLogger(__name__)

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


class AskQuestionRequest(BaseModel):
    question: str
    context: Optional[str] = None
    session_id: Optional[str] = None
    model: Optional[str] = None
    stream: Optional[bool] = False
    use_rag: Optional[bool] = None


class StopAskRequest(BaseModel):
    session_id: Optional[str] = None


@router.post("/ask/stream")
async def ask_question_stream(request: AskQuestionRequest):
    """
    Stream tokens in real-time using Server-Sent Events (SSE).
    """
    session_id = (request.session_id or "default").strip()
    current_task = asyncio.current_task()
    if current_task:
        session_manager.register_task(session_id, current_task)

    logger.info(
        f"[ask-stream] session_id={session_id!r} "
        f"model={request.model!r} "
        f"use_rag={request.use_rag!r} "
        f"question={request.question[:80]!r}"
    )

    def event_stream():
        try:
            for token in get_answer_stream_from_rag(
                question=request.question,
                context=request.context,
                session_id=session_id,
                model=request.model,
                use_rag=request.use_rag,
            ):
                payload = json.dumps({"token": token})
                yield f"data: {payload}\n\n"
            yield "data: [DONE]\n\n"
        except asyncio.CancelledError:
            logger.info(f"[ask-stream] Cancelled for session '{session_id}'.")
            err_payload = json.dumps({"error": f"Generation stopped for session '{session_id}'."})
            yield f"data: {err_payload}\n\n"
        except Exception as e:
            logger.error(f"[ask-stream] error: {e}")
            err_payload = json.dumps({"error": str(e)})
            yield f"data: {err_payload}\n\n"
        finally:
            if current_task:
                session_manager.unregister_task(session_id, current_task)

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


@router.post("/ask")
async def ask_question(request: AskQuestionRequest):
    if request.stream:
        return await ask_question_stream(request)

    session_id = (request.session_id or "default").strip()
    current_task = asyncio.current_task()
    if current_task:
        session_manager.register_task(session_id, current_task)

    try:
        logger.info(
            f"[ask] session_id={session_id!r} "
            f"model={request.model!r} "
            f"use_rag={request.use_rag!r} "
            f"question={request.question[:80]!r}"
        )

        answer = get_answer_from_rag(
            question=request.question,
            context=request.context,
            session_id=session_id,
            model=request.model,
            use_rag=request.use_rag,
        )

        logger.info(
            f"[ask] answer preview: {str(answer)[:120]!r}"
        )

        return {
            "status": "success",
            "answer": answer,
        }

    except asyncio.CancelledError:
        logger.info(f"[ask] Generation stopped for session '{session_id}'.")
        raise HTTPException(
            status_code=499,
            detail=f"Generation terminated for session '{session_id}'."
        )
    except Exception as e:
        logger.error(f"[ask] error: {e}")
        raise HTTPException(
            status_code=500,
            detail=str(e)
        )
    finally:
        if current_task:
            session_manager.unregister_task(session_id, current_task)


@router.post("/ask/stop")
async def stop_ask(
    request: Optional[StopAskRequest] = None,
    session_id: Optional[str] = None,
):
    """
    Cancel active model generation task(s).
    If session_id is provided, only halts that specific session.
    """
    target_id = (request.session_id if request else None) or session_id
    if target_id:
        target_id = target_id.strip()
        terminated = session_manager.terminate_session(target_id)
        logger.info(f"[ask] Stop signal for session '{target_id}' (terminated={terminated}).")
        return {
            "status": "success",
            "session_id": target_id,
            "terminated": terminated,
            "message": (
                f"Generation for session '{target_id}' stopped."
                if terminated
                else f"No active generation running for session '{target_id}'."
            ),
        }

    count = session_manager.terminate_all()
    return {
        "status": "success",
        "message": f"Stopped {count} active generation(s).",
        "stopped_count": count,
    }