from __future__ import annotations
import base64
import uuid
from typing import Any

import anthropic
import asyncpg
import httpx
from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile
from fastapi.responses import JSONResponse

from ..ai.chat import chat_with_document
from ..ai.citation_mapper import map_citations
from ..ai.claude import extract_facts
from ..ai.post_process import assign_page_numbers
from ..db.documents import get_document_by_id, insert_document, list_documents
from ..db.extractions import (
    get_all_facts_by_document_id,
    insert_extraction_result,
)
from ..models import Document, DocumentContent
from ..parsers.pdf import parse_pdf
from ..parsers.text import parse_text
from ..parsers.web import parse_web
from ..parsers.word import parse_word

router = APIRouter()


def _pool(request: Request) -> asyncpg.Pool:
    return request.app.state.pool


def _client(request: Request) -> anthropic.AsyncAnthropic:
    return request.app.state.anthropic_client


# ---------------------------------------------------------------------------
# GET / — list all documents
# ---------------------------------------------------------------------------

@router.get("/")
async def list_docs(request: Request) -> dict:
    documents = await list_documents(_pool(request))
    return {"documents": [d.model_dump() for d in documents]}


# ---------------------------------------------------------------------------
# POST / — upload a document (file or URL)
# ---------------------------------------------------------------------------

@router.post("/", status_code=201)
async def upload_document(
    request: Request,
    file: UploadFile | None = File(default=None),
    url: str | None = Form(default=None),
) -> dict:
    doc_id = str(uuid.uuid4())
    doc: Document

    if file is not None:
        buffer = await file.read()
        original_name: str = file.filename or "upload"
        mime_type: str = file.content_type or ""

        if mime_type == "application/pdf" or original_name.endswith(".pdf"):
            doc_type = "pdf"
        elif (
            mime_type
            == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
            or original_name.endswith(".docx")
        ):
            doc_type = "word"
        elif original_name.endswith(".md") or original_name.endswith(".markdown"):
            doc_type = "markdown"
        else:
            doc_type = "text"

        if doc_type == "pdf":
            content = parse_pdf(buffer)
            pdf_data = base64.b64encode(buffer).decode()
            from datetime import datetime, timezone
            doc = Document(
                id=doc_id,
                title=original_name,
                type=doc_type,
                content=content,
                originalFileName=original_name,
                pdfData=pdf_data,
                createdAt=datetime.now(timezone.utc).isoformat(),
            )
        elif doc_type == "word":
            content = parse_word(buffer)
            from datetime import datetime, timezone
            doc = Document(
                id=doc_id,
                title=original_name,
                type=doc_type,
                content=content,
                originalFileName=original_name,
                createdAt=datetime.now(timezone.utc).isoformat(),
            )
        else:
            text_content = buffer.decode("utf-8", errors="replace")
            content = parse_text(text_content, doc_type)
            from datetime import datetime, timezone
            doc = Document(
                id=doc_id,
                title=original_name,
                type=doc_type,
                content=content,
                originalFileName=original_name,
                createdAt=datetime.now(timezone.utc).isoformat(),
            )

    elif url:
        is_pdf = url.lower().endswith(".pdf")
        if not is_pdf:
            try:
                async with httpx.AsyncClient(follow_redirects=True, timeout=10.0) as hx:
                    head_res = await hx.head(url)
                ct = head_res.headers.get("content-type", "")
                is_pdf = "application/pdf" in ct
            except Exception:
                pass

        if is_pdf:
            async with httpx.AsyncClient(follow_redirects=True, timeout=30.0) as hx:
                pdf_res = await hx.get(url)
            if pdf_res.status_code >= 400:
                raise HTTPException(
                    status_code=502,
                    detail=f"Failed to fetch PDF: {pdf_res.status_code}",
                )
            buffer = pdf_res.content
            content = parse_pdf(buffer)
            pdf_data = base64.b64encode(buffer).decode()
            title = url.split("/")[-1] or url
            from datetime import datetime, timezone
            doc = Document(
                id=doc_id,
                title=title,
                type="pdf",
                content=content,
                sourceUrl=url,
                pdfData=pdf_data,
                createdAt=datetime.now(timezone.utc).isoformat(),
            )
        else:
            title, content = await parse_web(url)
            from datetime import datetime, timezone
            doc = Document(
                id=doc_id,
                title=title,
                type="web",
                content=content,
                sourceUrl=url,
                createdAt=datetime.now(timezone.utc).isoformat(),
            )
    else:
        raise HTTPException(status_code=400, detail="No file or URL provided")

    await insert_document(_pool(request), doc)
    return {"document": doc.model_dump()}


# ---------------------------------------------------------------------------
# GET /:id — get document by ID
# ---------------------------------------------------------------------------

@router.get("/{doc_id}")
async def get_document(doc_id: str, request: Request) -> dict:
    doc = await get_document_by_id(_pool(request), doc_id)
    if doc is None:
        raise HTTPException(status_code=404, detail="Document not found")
    return {"document": doc.model_dump()}


# ---------------------------------------------------------------------------
# POST /:id/extract — run fact extraction
# ---------------------------------------------------------------------------

@router.post("/{doc_id}/extract")
async def extract(doc_id: str, request: Request) -> dict:
    body: dict[str, Any] = await request.json() if await request.body() else {}
    prompt: str | None = body.get("prompt") or None

    doc = await get_document_by_id(_pool(request), doc_id)
    if doc is None:
        raise HTTPException(status_code=404, detail="Document not found")

    extraction, usage = await extract_facts(_client(request), doc.content.rawText, prompt)
    facts = map_citations(doc.content.rawText, extraction.facts)
    assign_page_numbers(doc, facts)

    result = await insert_extraction_result(
        _pool(request),
        doc.id,
        facts,
        prompt,
        extraction.no_mention_found,
        extraction.absence_evidence,
    )

    result_dict = result.model_dump()
    result_dict["usage"] = usage.model_dump()
    return {"result": result_dict}


# ---------------------------------------------------------------------------
# GET /:id/facts — get all accumulated facts for this document
# ---------------------------------------------------------------------------

@router.get("/{doc_id}/facts")
async def get_facts(doc_id: str, request: Request) -> dict:
    result = await get_all_facts_by_document_id(_pool(request), doc_id)
    if result is None:
        raise HTTPException(
            status_code=404,
            detail="No extraction results found for this document",
        )
    return {"result": result.model_dump()}


# ---------------------------------------------------------------------------
# POST /:id/chat — chat with a document
# ---------------------------------------------------------------------------

@router.post("/{doc_id}/chat")
async def chat(doc_id: str, request: Request) -> dict:
    body: dict[str, Any] = await request.json()
    message: str = body.get("message", "").strip()
    history: list[dict] = body.get("history", [])

    if not message:
        raise HTTPException(status_code=400, detail="message is required")

    doc = await get_document_by_id(_pool(request), doc_id)
    if doc is None:
        raise HTTPException(status_code=404, detail="Document not found")

    result = await chat_with_document(_client(request), doc.content.rawText, history, message)
    assign_page_numbers(doc, result.facts)

    if result.facts:
        await insert_extraction_result(
            _pool(request),
            doc.id,
            result.facts,
            message,
            False,
            None,
            source="chat",
        )

    return {
        "reply": {
            "message": result.message,
            "facts": [f.model_dump() for f in result.facts],
            "notInDocument": result.notInDocument,
            "notInDocumentExplanation": result.notInDocumentExplanation,
            "usage": result.usage,
        }
    }
