modules/models.py

"""
models.py — Pydantic models for type safety and validation
Ensures data integrity across all modules.
"""

from __future__ import annotations

import re
from enum import Enum
from typing import Any, Optional

from pydantic import BaseModel, Field, field_validator


# ── Enums ──────────────────────────────────────────────────────


class Category(str, Enum):
    TADEEL = "tadeel"
    JARH = "jarh"


class Severity(str, Enum):
    PASS = "PASS"
    WARN = "WARN"
    FAIL = "FAIL"


class JobStatus(str, Enum):
    IDLE = "idle"
    SEARCHING = "searching"
    SEARCH_DONE = "search_done"
    WRITING = "writing"
    WRITE_DONE = "write_done"
    REVIEWING = "reviewing"
    REVIEW_DONE = "review_done"
    FAILED = "failed"


# ── Search Result ──────────────────────────────────────────────


class SearchResultModel(BaseModel):
    source: str = ""
    term: str = ""
    narrator: str = ""
    text: str = ""
    book: str = ""
    author: str = ""
    volume: str = ""
    page: str = ""
    link: str = ""
    section: int = Field(ge=0, le=5, default=0)
    context: str = ""
    qattan_evidence: str = ""
    needs_review: bool = False
    is_negated: bool = False
    relevance_score: float = Field(ge=0.0, le=8.0, default=0.0)

    @field_validator("link")
    @classmethod
    def validate_link(cls, v: str) -> str:
        if v and not v.startswith("http"):
            return ""
        return v

    @field_validator("text")
    @classmethod
    def validate_text(cls, v: str) -> str:
        return v[:5000] if v else ""


# ── Section Data ───────────────────────────────────────────────


class SectionResult(BaseModel):
    name: str
    description: str = ""
    count: int = 0
    results: list[dict[str, Any]] = Field(default_factory=list)


# ── Draft (Collector Output) ───────────────────────────────────


class NarratorInfo(BaseModel):
    name: str
    book: str = ""
    author: str = ""
    link: str = ""
    context: str = ""


class Draft(BaseModel):
    term: str
    category: Category
    total_results: int = 0
    by_source: dict[str, int] = Field(default_factory=dict)
    by_section: dict[int, SectionResult] = Field(default_factory=dict)
    sections_used: list[int] = Field(default_factory=list)
    narrators: list[NarratorInfo] = Field(default_factory=list)
    books_count: int = 0
    authors_count: int = 0
    estimated_citations: int = 0
    qattan_evidence_count: int = 0
    needs_review_count: int = 0


# ── Section Audit Stats ────────────────────────────────────────


class SectionAuditStats(BaseModel):
    name: str
    exists: bool = False
    words: int = 0
    citations: int = 0
    paragraphs: int = 0
    paragraphs_without_citation: int = 0
    density: float = Field(ge=0.0, le=1.0, default=0.0)


# ── Audit Issue ────────────────────────────────────────────────


class AuditIssue(BaseModel):
    type: str
    severity: str = "warning"  # "error" | "warning"
    message: str
    section: str | None = None
    details: Any = None


# ── Audit Result ───────────────────────────────────────────────


class AuditStats(BaseModel):
    inline_citations: int = 0
    numbered_citations: int = 0
    total_citations: int = 0
    total_paragraphs: int = 0
    paragraphs_without_citation: int = 0
    citation_density: float = 0.0
    words: int = 0
    chars: int = 0
    lines: int = 0
    english_words: int = 0
    required_sections: int = 7
    found_sections: int = 0
    total_links: int = 0
    hallucinated_links: int = 0
    raw_links_total: int = 0
    raw_links_used: int = 0
    raw_utilization: float = 0.0
    sections: dict[str, SectionAuditStats] = Field(default_factory=dict)


class AuditResult(BaseModel):
    term: str
    category: str
    severity: Severity = Severity.PASS
    stats: AuditStats = Field(default_factory=AuditStats)
    issues: list[AuditIssue] = Field(default_factory=list)
    passed: bool = True


# ── Job State ──────────────────────────────────────────────────


class JobState(BaseModel):
    job_id: str
    term: str
    category: Category
    status: str = "جاري التجهيز..."
    search_progress: int = 0
    write_progress: int = 0
    review_progress: int = 0
    done: bool = False
    search_done: bool = False
    write_done: bool = False
    review_done: bool = False
    review_in_progress: bool = False
    study_reviewed: bool = False
    error: str | None = None
    write_error: str | None = None
    review_error: str | None = None
    log: list[str] = Field(default_factory=list)
    section: str = "all"
    by_source: dict[str, int] = Field(default_factory=dict)
    by_section: dict[str, Any] = Field(default_factory=dict)
    section_details: dict[str, Any] = Field(default_factory=dict)
    qattan_evidence_count: int = 0
    needs_review_count: int = 0
    total_results: int = 0
    draft: dict[str, Any] | None = None
    prompt_path: str = ""
    super_prompt_path: str = ""
    study_text: str = ""
    study_file: str = ""
    created_at: float = 0.0
    api_key_override: str | None = None
    audit: dict[str, Any] | None = None