Blame for src/document_generator_models.py

import json
import pathlib
from enum import StrEnum
from typing import Any, Dict, List, Optional, Set

from pydantic import BaseModel


# ── Authorized Representatives lookup map ──────────────────────────────────────
# Keys: full name (displayed in dropdown)
# Values: dict of data_key → value for the other rep fields
REPRESENTATIVES: Dict[str, Dict[str, str]] = {
    "Ahmad Royhan": {
        "rep_id":      "3174050905990008",
        "rep_address": "Jl. Kramat 1 No 15, Kwitang",
        "rep_city":    "Jakarta Selatan",
    },
    # Add more representatives here:
    # "Full Name": {
    #     "rep_id":      "...",
    #     "rep_address": "...",
    #     "rep_city":    "...",
    # },
}


class DocumentType(StrEnum):
    COVER_LETTER_SCHENGEN = "cover_letter_schengen"
    COVER_LETTER_UK       = "cover_letter_uk"
    LETTER_OF_AUTHORIZATION = "letter_of_authorization"
    LETTER_OF_AUTHORIZATION_MINOR = "letter_of_authorization_minor"
    # INVITATION_LETTER    = "invitation_letter"
    # BANK_REFERENCE       = "bank_reference"


class SectionConfig(BaseModel):
    """Metadata for one form section, including explicit display order."""

    key: str
    title: str
    order: int = 0
    # data_keys whose widgets are rendered as read-only in this section
    disabled_keys: Set[str] = set()


class FieldSchema(BaseModel):
    """Describes one editable field for a document type."""

    widget_key: str               # Streamlit session-state key
    data_key: str                 # Key inside the data dict
    label: str
    placeholder: str = ""
    description: Optional[str] = None   # Sent to /prefill-document; falls back to placeholder
    field_type: str = "text_input"      # "text_input" | "text_area" | "select"
    options: Optional[List[str]] = None  # For field_type="select"
    section: str = "general"            # Maps to a SectionConfig.key
    nested_under: Optional[str] = None  # e.g. "personal_details"
    order: int = 0                      # Explicit render position within the section; lower = earlier
    static_value: Optional[str] = None  # Pre-seeded value; applied as fallback when DB prefill returns null


class DocumentSchema(BaseModel):
    """Full metadata for a document type.

    Drives form rendering, session-state key management, and API calls.
    Loaded from a JSON file in the src/schemas/ directory.
    """

    document_type: DocumentType
    title: str
    icon: str = "📝"
    sections: List[SectionConfig]
    fields: List[FieldSchema]
    has_trip_type: bool = False
    has_group_members: bool = False

    def ordered_sections(self) -> List[SectionConfig]:
        """Return sections sorted by SectionConfig.order (stable sort)."""
        return sorted(self.sections, key=lambda s: s.order)

    def section_title(self, key: str) -> str:
        """Return the display title for a section key."""
        for s in self.sections:
            if s.key == key:
                return s.title
        return key

    def section_disabled_keys(self, key: str) -> Set[str]:
        """Return the disabled_keys set for a section key."""
        for s in self.sections:
            if s.key == key:
                return s.disabled_keys
        return set()

    def fields_in_section(self, section: str) -> List[FieldSchema]:
        """Return fields belonging to section, sorted by FieldSchema.order."""
        return sorted(
            [f for f in self.fields if f.section == section],
            key=lambda f: f.order,
        )

    def top_level_fields(self, section: str) -> List[FieldSchema]:
        """Return non-nested fields in section, sorted by order."""
        return sorted(
            [f for f in self.fields if f.section == section and f.nested_under is None],
            key=lambda f: f.order,
        )

    def widget_keys(self) -> List[str]:
        """All Streamlit widget keys that must be cleared on undo/redo."""
        keys: List[str] = []
        if self.has_trip_type:
            keys.append(f"{self.document_type}_trip_type")
        keys.extend(f.widget_key for f in self.fields)
        if self.has_group_members:
            keys.append(f"{self.document_type}_group_members_editor")
        return keys

    def nested_fields(self, nested_under: str) -> List[FieldSchema]:
        return [f for f in self.fields if f.nested_under == nested_under]


# ── JSON-based schema loader ────────────────────────────────────────────────────

_SCHEMA_DIR: pathlib.Path = pathlib.Path(__file__).parent / "schemas"
_REPRESENTATIVES_MARKER: str = "$REPRESENTATIVES"


def _resolve_options(options: Any) -> Optional[List[str]]:
    """Replace the $REPRESENTATIVES marker with the actual representative names list."""
    if options == _REPRESENTATIVES_MARKER:
        return list(REPRESENTATIVES.keys())
    return options  # type: ignore[return-value]


def _load_document_schema(path: pathlib.Path) -> DocumentSchema:
    """Load and validate a DocumentSchema from a JSON file."""
    raw: Dict[str, Any] = json.loads(path.read_text(encoding="utf-8"))
    # Resolve dynamic markers in field options before Pydantic validation
    for field_data in raw.get("fields", []):
        if "options" in field_data:
            field_data["options"] = _resolve_options(field_data["options"])
    return DocumentSchema.model_validate(raw)


def _build_registry() -> Dict[str, DocumentSchema]:
    """Scan the schemas/ directory and build DOCUMENT_REGISTRY from JSON files."""
    registry: Dict[str, DocumentSchema] = {}
    for json_file in sorted(_SCHEMA_DIR.glob("*.json")):
        schema: DocumentSchema = _load_document_schema(json_file)
        registry[schema.document_type] = schema
    return registry


DOCUMENT_REGISTRY: Dict[str, DocumentSchema] = _build_registry()


# ── API request models ─────────────────────────────────────

class PrefillDocumentRequest(BaseModel):
    """
    Sent to /prefill-document.
    `structure` describes the fields the LLM should map from DB data.
    `sections` (optional) — if provided, the response is nested per section:
        { "sections": [{key, title, order, fields: [{key, label, order, value}]}],
          "_missing_required": [...] }
    If `sections` is omitted, falls back to the old flat dict response.
    """
    application_id: int
    structure: List[Dict[str, Any]]
    sections: Optional[List[Dict[str, Any]]] = None  # [{key, title, order}, ...]


class GenerateDocumentRequest(BaseModel):
    """
    Sent to /generate-document.
    `doc_data` carries all form fields as-is; the backend merges them on top
    of the DB data and generates the document.
    Adding a new document type requires no change to this model.
    """
    application_id: int
    document_name: str  # value from DocumentType
    doc_data: Optional[Dict[str, Any]] = None


class DocChatRequest(BaseModel):
    """Sent to /generate-document/chat (new stateless draft-revision flow)."""
    query: str
    history: List[Dict[str, str]]
    session_uuid: Optional[str] = None
    current_document: Optional[Dict[str, Any]] = None   # structured {title, city_date, main}
    current_document_content: str = ""                  # legacy plain-markdown fallback
    structure: Optional[List[Dict[str, Any]]] = None


class GenerateDraftRequest(BaseModel):
    """
    Sent to /generate-document/draft.
    Generates a Markdown draft without creating a Google Doc.
    `doc_type` and `structure` are stored on the frontend for later handover to the core API.
    """
    doc_type: str
    data: Dict[str, Any]
    application_id: Optional[int] = None  # used by the API to load persona from DB
    structure: Optional[List[Dict[str, Any]]] = None  # field definitions for core API handover
    session_uuid: Optional[str] = None


class DraftField(BaseModel):
    """A single structured field from /generate-document/draft."""
    style: str
    content: str


class ExportDocumentRequest(BaseModel):
    """Sent to /export-document once the user is satisfied with the draft."""
    document_content: Optional[str] = None  # legacy
    title: Optional[DraftField] = None
    city_date: Optional[DraftField] = None
    main: Optional[DraftField] = None
    export_format: str = "gdocs"
    document_url: Optional[str] = None  # existing Google Doc URL; if set, AI edits instead of creating


# ── Shared sub-models (used for type hints / validation internally) ──

class PersonalDetails(BaseModel):
    name: Optional[str] = None
    dob: Optional[str] = None
    nationality: Optional[str] = None
    occupation: Optional[str] = None
    passport_number: Optional[str] = None


class GroupMember(BaseModel):
    relationship: Optional[str] = None
    name: Optional[str] = None
    dob: Optional[str] = None
    occupation: Optional[str] = None
    nationality: Optional[str] = None
    passport_number: Optional[str] = None