"""Extraction batch and document models (norm extraction pipeline)."""
from __future__ import annotations

import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional

from sqlalchemy import CheckConstraint, ForeignKey, String, Text, Integer
from sqlalchemy.dialects.postgresql import UUID, JSONB
from sqlalchemy.orm import Mapped, mapped_column, relationship

from normbase.db.connection import Base
from normbase.models.norm import utc_now


class ExtractionBatch(Base):
    """A batch of uploaded documents for norm extraction."""

    __tablename__ = "extraction_batches"

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    title: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    field_id: Mapped[Optional[uuid.UUID]] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("fields.id", ondelete="SET NULL"),
        nullable=True,
        comment="Default information field for norms produced from this batch",
    )
    status: Mapped[str] = mapped_column(
        String(32),
        nullable=False,
        default="open",
        comment="open/processing/completed/failed",
    )
    options: Mapped[Dict[str, Any]] = mapped_column(
        JSONB, nullable=False, default=lambda: {}
    )
    created_by: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    created_at: Mapped[datetime] = mapped_column(default=utc_now, nullable=False)
    updated_at: Mapped[datetime] = mapped_column(
        default=utc_now, onupdate=utc_now, nullable=False
    )

    documents: Mapped[List["ExtractionDocument"]] = relationship(
        "ExtractionDocument",
        back_populates="batch",
        cascade="all, delete-orphan",
    )

    __table_args__ = (
        CheckConstraint(
            "status IN ('open', 'processing', 'completed', 'failed')",
            name="ck_extraction_batch_status",
        ),
        CheckConstraint(
            "created_by IS NULL "
            "OR created_by LIKE 'user:%' "
            "OR created_by LIKE 'org:%' "
            "OR created_by LIKE 'system:%' "
            "OR created_by LIKE 'external:%'",
            name="ck_extraction_batch_created_by_format",
        ),
    )


class ExtractionDocument(Base):
    """One uploaded file within an extraction batch."""

    __tablename__ = "extraction_documents"

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    batch_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("extraction_batches.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    original_filename: Mapped[str] = mapped_column(String(512), nullable=False)
    content_type: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    size_bytes: Mapped[int] = mapped_column(Integer, nullable=False)
    storage_relative_path: Mapped[str] = mapped_column(
        String(1024),
        nullable=False,
        comment="Path relative to NORMBASE_EXTRACTION_UPLOAD_DIR",
    )
    source_file_url: Mapped[Optional[str]] = mapped_column(
        Text,
        nullable=True,
        comment="Optional stable or public URL for UI (see design doc)",
    )
    parse_status: Mapped[str] = mapped_column(
        String(32),
        nullable=False,
        default="pending",
        comment="pending/skipped/done",
    )
    extraction_status: Mapped[str] = mapped_column(
        String(32),
        nullable=False,
        default="pending",
        comment="pending/queued/running/succeeded/failed",
    )
    error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    extracted_text: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    created_at: Mapped[datetime] = mapped_column(default=utc_now, nullable=False)
    updated_at: Mapped[datetime] = mapped_column(
        default=utc_now, onupdate=utc_now, nullable=False
    )

    batch: Mapped["ExtractionBatch"] = relationship(
        "ExtractionBatch", back_populates="documents"
    )

    __table_args__ = (
        CheckConstraint(
            "parse_status IN ('pending', 'skipped', 'done')",
            name="ck_extraction_document_parse_status",
        ),
        CheckConstraint(
            "extraction_status IN ('pending', 'queued', 'running', 'succeeded', 'failed')",
            name="ck_extraction_document_extraction_status",
        ),
        CheckConstraint("size_bytes >= 0", name="ck_extraction_document_size"),
    )
