"""Extraction batch lifecycle: create batch, upload files, run per-document LLM extraction."""
from __future__ import annotations

from pathlib import Path
from typing import List, Optional, Tuple
from uuid import UUID, uuid4

from fastapi import HTTPException, UploadFile, status
from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from normbase.config import settings
from normbase.models.extraction import ExtractionBatch, ExtractionDocument
from normbase.models.norm import utc_now
from normbase.schemas import ExtractionBatchCreate
from normbase.services.extraction_artifact_storage import resolve_under_cwd
from normbase.services.extraction_document_parser import extract_text_from_upload
from normbase.services.norm_extraction_service import NormExtractionService


def _safe_filename(name: str) -> str:
    base = Path(name).name
    if not base or base in (".", ".."):
        return "upload.bin"
    return base[:512] if len(base) > 512 else base


class ExtractionBatchService:
    async def create_batch(self, db: AsyncSession, req: ExtractionBatchCreate) -> ExtractionBatch:
        batch = ExtractionBatch(
            title=req.title,
            field_id=req.field_id,
            created_by=req.created_by,
            options=dict(req.options or {}),
        )
        db.add(batch)
        await db.commit()
        await db.refresh(batch)
        return batch

    async def get_batch(self, db: AsyncSession, batch_id: UUID) -> ExtractionBatch:
        result = await db.execute(
            select(ExtractionBatch)
            .options(selectinload(ExtractionBatch.documents))
            .where(ExtractionBatch.id == batch_id)
        )
        batch = result.scalar_one_or_none()
        if batch is None:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="Extraction batch not found",
            )
        return batch

    async def add_documents(
        self,
        db: AsyncSession,
        *,
        batch_id: UUID,
        files: List[UploadFile],
    ) -> List[ExtractionDocument]:
        if not files:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="At least one file is required.",
            )
        batch = await self.get_batch(db, batch_id)
        upload_root = resolve_under_cwd(settings.extraction_upload_dir)
        max_sz = settings.extraction_upload_max_bytes
        public_base = settings.extraction_upload_public_base_url
        out: List[ExtractionDocument] = []

        for uf in files:
            data = await uf.read()
            if len(data) > max_sz:
                raise HTTPException(
                    status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
                    detail=(
                        f"File exceeds NORMBASE_EXTRACTION_UPLOAD_MAX_BYTES ({max_sz})."
                    ),
                )
            doc_id = uuid4()
            safe = _safe_filename(uf.filename or "upload.bin")
            rel = f"batches/{batch_id}/{doc_id}_{safe}"
            abs_path = upload_root / rel
            abs_path.parent.mkdir(parents=True, exist_ok=True)
            abs_path.write_bytes(data)

            ct = uf.content_type or ""
            extracted, parse_status, _parse_err = extract_text_from_upload(
                filename=safe,
                content_type=ct or None,
                data=data,
            )

            url: Optional[str] = None
            if public_base:
                url = f"{public_base.rstrip('/')}/{rel}"

            doc = ExtractionDocument(
                id=doc_id,
                batch_id=batch_id,
                original_filename=safe,
                content_type=ct or None,
                size_bytes=len(data),
                storage_relative_path=rel,
                source_file_url=url,
                parse_status=parse_status,
                extraction_status="pending",
                extracted_text=extracted,
            )
            db.add(doc)
            out.append(doc)

        batch.status = "processing"
        await db.commit()
        for d in out:
            await db.refresh(d)
        return out

    async def extract_document(
        self,
        db: AsyncSession,
        *,
        batch_id: UUID,
        document_id: UUID,
        default_lang: str = "zh-CN",
        max_norms: int = 30,
        submitter_user_id: Optional[str] = None,
    ) -> Tuple[UUID, str, List[UUID], List[dict]]:
        batch_result = await db.execute(
            select(ExtractionBatch).where(ExtractionBatch.id == batch_id)
        )
        batch = batch_result.scalar_one_or_none()
        if batch is None:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="Extraction batch not found",
            )

        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",
            )

        text = doc.extracted_text
        if not text:
            upload_root = resolve_under_cwd(settings.extraction_upload_dir)
            raw = (upload_root / doc.storage_relative_path).read_bytes()
            extracted, _ps, perr = extract_text_from_upload(
                filename=doc.original_filename,
                content_type=doc.content_type,
                data=raw,
            )
            if not extracted:
                msg = perr or "Could not extract UTF-8 text or PDF text from this file."
                await db.execute(
                    update(ExtractionDocument)
                    .where(ExtractionDocument.id == document_id)
                    .values(
                        extraction_status="failed",
                        error_message=msg[:4000],
                        updated_at=utc_now(),
                    )
                )
                await db.commit()
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail=msg,
                )
            text = extracted

        run_values: dict = {
            "extraction_status": "running",
            "error_message": None,
            "updated_at": utc_now(),
        }
        if not doc.extracted_text and text:
            run_values["extracted_text"] = text
            run_values["parse_status"] = "done"

        await db.execute(
            update(ExtractionDocument)
            .where(ExtractionDocument.id == document_id)
            .values(**run_values)
        )
        await db.commit()

        extraction = NormExtractionService()
        try:
            session_id, rel, norm_ids, issues = await extraction.extract_from_text(
                db,
                text=text,
                default_lang=default_lang,
                field_id=batch.field_id,
                created_by=batch.created_by,
                max_norms=max_norms,
                default_authority_level="legal",
                default_formality_level="formal",
                submitter_user_id=submitter_user_id,
                llm_doc_id=str(doc.id),
                llm_doc_title=doc.original_filename,
                source_extraction_extra={
                    "mode": "batch_document",
                    "batch_id": str(batch_id),
                    "document_id": str(document_id),
                    "source_file_storage_key": doc.storage_relative_path,
                    "source_file_url": doc.source_file_url,
                },
                artifact_relative_segments=(
                    "batches",
                    str(batch_id),
                    str(document_id),
                    "process.json",
                ),
            )
        except HTTPException as e:
            await db.execute(
                update(ExtractionDocument)
                .where(ExtractionDocument.id == document_id)
                .values(
                    extraction_status="failed",
                    error_message=str(e.detail)[:4000],
                    updated_at=utc_now(),
                )
            )
            await db.commit()
            raise
        except Exception as e:
            await db.execute(
                update(ExtractionDocument)
                .where(ExtractionDocument.id == document_id)
                .values(
                    extraction_status="failed",
                    error_message=str(e)[:4000],
                    updated_at=utc_now(),
                )
            )
            await db.commit()
            raise

        await db.execute(
            update(ExtractionDocument)
            .where(ExtractionDocument.id == document_id)
            .values(
                extraction_status="succeeded",
                error_message=None,
                updated_at=utc_now(),
            )
        )
        await db.commit()

        return session_id, rel, norm_ids, issues
