import asyncio
import json
import math
import traceback
import uuid
import dirtyjson
from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from sqlalchemy.ext.asyncio import AsyncSession

from models.presentation_outline_model import PresentationOutlineModel
from models.sql.presentation import PresentationModel
from models.sse_response import (
    SSECompleteResponse,
    SSEErrorResponse,
    SSEResponse,
    SSEStatusResponse,
)
from services.temp_file_service import TEMP_FILE_SERVICE
from services.database import get_async_session, async_session_maker
from services.documents_loader import DocumentsLoader
from utils.llm_calls.generate_presentation_outlines import generate_ppt_outline
from utils.ppt_utils import get_presentation_title_from_outlines

OUTLINES_ROUTER = APIRouter(prefix="/outlines", tags=["Outlines"])


@OUTLINES_ROUTER.get("/stream/{id}")
async def stream_outlines(
    id: uuid.UUID, sql_session: AsyncSession = Depends(get_async_session)
):
    presentation = await sql_session.get(PresentationModel, id)

    if not presentation:
        raise HTTPException(status_code=404, detail="Presentation not found")

    # Capture all data from the presentation before the endpoint returns
    # (Depends session closes when endpoint function returns, before StreamingResponse runs)
    presentation_id = presentation.id
    presentation_content = presentation.content
    presentation_file_paths = presentation.file_paths
    presentation_n_slides = presentation.n_slides
    presentation_language = presentation.language
    presentation_tone = presentation.tone
    presentation_verbosity = presentation.verbosity
    presentation_instructions = presentation.instructions
    presentation_include_title_slide = presentation.include_title_slide
    presentation_include_table_of_contents = presentation.include_table_of_contents
    presentation_web_search = presentation.web_search

    temp_dir = TEMP_FILE_SERVICE.create_temp_dir()

    async def inner():
        yield SSEStatusResponse(
            status="Generating presentation outlines..."
        ).to_string()

        additional_context = ""
        if presentation_file_paths:
            documents_loader = DocumentsLoader(file_paths=presentation_file_paths)
            await documents_loader.load_documents(temp_dir)
            documents = documents_loader.documents
            if documents:
                additional_context = "\n\n".join(documents)

        presentation_outlines_text = ""

        n_slides_to_generate = presentation_n_slides
        if presentation_include_table_of_contents:
            needed_toc_count = math.ceil((presentation_n_slides - 1) / 10)
            n_slides_to_generate -= math.ceil(
                (presentation_n_slides - needed_toc_count) / 10
            )

        async for chunk in generate_ppt_outline(
            presentation_content,
            n_slides_to_generate,
            presentation_language,
            additional_context,
            presentation_tone,
            presentation_verbosity,
            presentation_instructions,
            presentation_include_title_slide,
            presentation_web_search,
        ):
            # Give control to the event loop
            await asyncio.sleep(0)

            if isinstance(chunk, HTTPException):
                yield SSEErrorResponse(detail=chunk.detail).to_string()
                return

            yield SSEResponse(
                event="response",
                data=json.dumps({"type": "chunk", "chunk": chunk}),
            ).to_string()

            presentation_outlines_text += chunk

        try:
            presentation_outlines_json = dict(
                dirtyjson.loads(presentation_outlines_text)
            )

            # Fix: LLM sometimes returns slides as JSON string instead of list
            if "slides" in presentation_outlines_json and isinstance(presentation_outlines_json["slides"], str):
                print("[OUTLINE] Warning: slides field is a string, parsing as JSON...")
                presentation_outlines_json["slides"] = dirtyjson.loads(presentation_outlines_json["slides"])

        except Exception as e:
            traceback.print_exc()
            yield SSEErrorResponse(
                detail=f"Failed to generate presentation outlines. Please try again. {str(e)}",
            ).to_string()
            return

        presentation_outlines = PresentationOutlineModel(**presentation_outlines_json)

        presentation_outlines.slides = presentation_outlines.slides[
            :n_slides_to_generate
        ]

        outlines_data = presentation_outlines.model_dump()
        title = get_presentation_title_from_outlines(presentation_outlines)

        # Use a fresh session for the DB write — Depends session is already closed
        async with async_session_maker() as db:
            pres = await db.get(PresentationModel, presentation_id)
            if pres:
                pres.outlines = outlines_data
                pres.title = title
                db.add(pres)
                await db.commit()
                await db.refresh(pres)
                yield SSECompleteResponse(
                    key="presentation", value=pres.model_dump(mode="json")
                ).to_string()
            else:
                yield SSEErrorResponse(detail="Presentation not found after outline generation").to_string()

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