"""Norm extraction API routes."""
from __future__ import annotations

from typing import List, Union
from uuid import UUID

from fastapi import (
    APIRouter,
    BackgroundTasks,
    Depends,
    File,
    HTTPException,
    Query,
    UploadFile,
    status,
)
from fastapi.responses import JSONResponse
from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from normbase.db.connection import get_db
from normbase.models.extraction import ExtractionDocument
from normbase.models.norm import utc_now
from normbase.schemas import (
    ExtractionBatchCreate,
    ExtractionBatchDetailResponse,
    ExtractionBatchResponse,
    ExtractionDocumentResponse,
    ExtractionExtractQueuedResponse,
    ExtractionFromTextRequest,
    ExtractionFromTextResponse,
    ExtractionImportIssue,
)
from normbase.services.extraction_background_jobs import run_batch_document_extraction_job
from normbase.services.extraction_batch_service import ExtractionBatchService
from normbase.services.norm_extraction_service import NormExtractionService

router = APIRouter(prefix="/extraction", tags=["extraction"])
_service = NormExtractionService()
_batch_service = ExtractionBatchService()


@router.post("/from-text", response_model=ExtractionFromTextResponse)
async def extraction_from_text(
    req: ExtractionFromTextRequest,
    db: AsyncSession = Depends(get_db),
):
    """Extract draft norms from natural language; write process JSON artifact; import valid rows."""
    session_id, rel, norm_ids, raw_issues = await _service.extract_from_text(
        db,
        text=req.text,
        default_lang=req.default_lang,
        field_id=req.field_id,
        created_by=req.created_by,
        max_norms=req.max_norms,
        default_authority_level=req.default_authority_level,
        default_formality_level=req.default_formality_level,
        submitter_user_id=req.submitter_user_id,
    )
    issues = [ExtractionImportIssue(**item) for item in raw_issues]
    return ExtractionFromTextResponse(
        session_id=session_id,
        artifact_relative_path=rel,
        norm_ids=norm_ids,
        imported_count=len(norm_ids),
        issues=issues,
    )


@router.post(
    "/batches",
    response_model=ExtractionBatchResponse,
    status_code=201,
)
async def create_extraction_batch(
    req: ExtractionBatchCreate,
    db: AsyncSession = Depends(get_db),
):
    """Create a batch container for document uploads."""
    batch = await _batch_service.create_batch(db, req)
    return ExtractionBatchResponse.model_validate(batch)


@router.get("/batches/{batch_id}", response_model=ExtractionBatchDetailResponse)
async def get_extraction_batch(
    batch_id: UUID,
    db: AsyncSession = Depends(get_db),
):
    """Get batch metadata and document rows."""
    batch = await _batch_service.get_batch(db, batch_id)
    data = ExtractionBatchResponse.model_validate(batch).model_dump()
    documents = [
        ExtractionDocumentResponse.model_validate(d) for d in batch.documents
    ]
    return ExtractionBatchDetailResponse(**data, documents=documents)


@router.post(
    "/batches/{batch_id}/documents",
    response_model=List[ExtractionDocumentResponse],
)
async def upload_extraction_documents(
    batch_id: UUID,
    files: List[UploadFile] = File(...),
    db: AsyncSession = Depends(get_db),
):
    """Upload one or more files under the batch (stored under NORMBASE_EXTRACTION_UPLOAD_DIR)."""
    docs = await _batch_service.add_documents(db, batch_id=batch_id, files=files)
    return [ExtractionDocumentResponse.model_validate(d) for d in docs]


@router.post(
    "/batches/{batch_id}/documents/{document_id}/extract",
    response_model=ExtractionFromTextResponse,
    responses={
        202: {
            "model": ExtractionExtractQueuedResponse,
            "description": "Background extraction scheduled",
        },
    },
)
async def extract_extraction_document(
    batch_id: UUID,
    document_id: UUID,
    background_tasks: BackgroundTasks,
    db: AsyncSession = Depends(get_db),
    default_lang: str = Query("zh-CN", min_length=2),
    max_norms: int = Query(30, ge=1, le=50),
    submitter_user_id: str | None = Query(None, max_length=256),
    background: bool = Query(
        False,
        description="If true, return 202 and run extraction after the response (poll batch for status).",
    ),
) -> Union[ExtractionFromTextResponse, JSONResponse]:
    """Run LLM extraction on parsed text (.txt / .md / .pdf) of an uploaded document."""
    if background:
        doc_result = await db.execute(
            select(ExtractionDocument).where(
                ExtractionDocument.id == document_id,
                ExtractionDocument.batch_id == batch_id,
            )
        )
        doc = doc_result.scalar_one_or_none()
        if doc is None:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="Extraction document not found",
            )
        if doc.extraction_status == "running":
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail="Extraction already running for this document.",
            )
        if doc.extraction_status == "queued":
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail="Extraction already queued for this document.",
            )
        if doc.extraction_status not in ("pending", "failed"):
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail="Background extract is only allowed when status is pending or failed.",
            )
        await db.execute(
            update(ExtractionDocument)
            .where(ExtractionDocument.id == document_id)
            .values(
                extraction_status="queued",
                error_message=None,
                updated_at=utc_now(),
            )
        )
        await db.commit()
        background_tasks.add_task(
            run_batch_document_extraction_job,
            batch_id,
            document_id,
            default_lang=default_lang,
            max_norms=max_norms,
            submitter_user_id=submitter_user_id,
        )
        body = ExtractionExtractQueuedResponse(
            batch_id=batch_id,
            document_id=document_id,
        )
        return JSONResponse(
            status_code=202,
            content=body.model_dump(mode="json"),
        )

    session_id, rel, norm_ids, raw_issues = await _batch_service.extract_document(
        db,
        batch_id=batch_id,
        document_id=document_id,
        default_lang=default_lang,
        max_norms=max_norms,
        submitter_user_id=submitter_user_id,
    )
    issues = [ExtractionImportIssue(**item) for item in raw_issues]
    return ExtractionFromTextResponse(
        session_id=session_id,
        artifact_relative_path=rel,
        norm_ids=norm_ids,
        imported_count=len(norm_ids),
        issues=issues,
    )
