src/document_generator_models.py

from pydantic import BaseModel
from typing import Optional, List, Dict, Any
from enum import StrEnum

# ── 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 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 column/section in the form
    nested_under: Optional[str] = None  # e.g. "personal_details"


class DocumentSchema(BaseModel):
    """Full metadata for a document type.
    Drives form rendering, session-state key management, and API calls.
    The backend uses `document_name` to ask an LLM to generate the right SQL.
    """
    document_type: DocumentType
    title: str
    icon: str = "📝"
    # section_key -> display header, order is preserved
    sections: Dict[str, str]
    fields: List[FieldSchema]
    has_trip_type: bool = False
    has_group_members: bool = False

    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 fields_in_section(self, section: str) -> List[FieldSchema]:
        return [f for f in self.fields if f.section == section]

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

    def top_level_fields(self, section: str) -> List[FieldSchema]:
        return [f for f in self.fields if f.section == section and f.nested_under is None]


def generate_person_fields(person_type: str, section: str, role_label: str) -> List[FieldSchema]:
    """Helper to keep the schema DRY by generating common personal fields for Minor LOA."""
    prefix = f"loa_minor_{person_type}"
    return [
        FieldSchema(widget_key=f"{prefix}_name",    data_key=f"{person_type}_name",    label=f"{role_label} Name",          placeholder="e.g. John Doe",           section=section),
        FieldSchema(widget_key=f"{prefix}_id",      data_key=f"{person_type}_id",      label=f"{role_label} National ID",   placeholder="e.g. 3174000000000000",   section=section),
        FieldSchema(widget_key=f"{prefix}_address", data_key=f"{person_type}_address", label=f"{role_label} Address",       placeholder="e.g. Jl. Bangka XI...",   section=section),
        FieldSchema(widget_key=f"{prefix}_city",    data_key=f"{person_type}_city",    label=f"{role_label} City/Province", placeholder="e.g. Jakarta Selatan",    section=section),
    ]


DOCUMENT_REGISTRY: Dict[str, DocumentSchema] = {
    DocumentType.COVER_LETTER_UK: DocumentSchema(
        document_type=DocumentType.COVER_LETTER_UK,
        title="UK Visa Cover Letter",
        icon="🇬🇧",
        sections={
            "trip":       "🌍 Trip & Embassy Details",
            "applicant":  "👤 Main Applicant",
            "employment": "💼 Employment & Ties",
            "financials": "💰 Financial Clarifications",
            "contact":    "📋 Contact & Documents",
        },
        fields=[
            # ── Trip & Embassy Section ────────────────────────────────
            FieldSchema(widget_key="uk_cl_city", data_key="city", label="City of Application", placeholder="e.g. Jakarta", section="trip"),
            FieldSchema(widget_key="uk_cl_emb_city", data_key="embassy_city", label="Embassy City", placeholder="e.g. Jakarta", section="trip"),
            FieldSchema(widget_key="uk_cl_emb_country", data_key="embassy_country", label="Embassy Country", placeholder="e.g. Indonesia", section="trip"),
            FieldSchema(widget_key="uk_cl_visa_type", data_key="visa_type", label="Visa Type", placeholder="e.g. Standard Visitor", section="trip"),
            FieldSchema(widget_key="uk_cl_purpose", data_key="visit_purpose", label="Visit Purpose", placeholder="e.g. tourism and cultural exploration", section="trip"),
            FieldSchema(widget_key="uk_cl_start", data_key="start_date", label="Travel Start Date", placeholder="e.g. 17 May 2026", section="trip"),
            FieldSchema(widget_key="uk_cl_end", data_key="end_date", label="Travel End Date", placeholder="e.g. 23 May 2026", section="trip"),
            FieldSchema(widget_key="uk_cl_duration", data_key="duration", label="Duration", placeholder="e.g. seven days and six nights", section="trip"),
            FieldSchema(widget_key="uk_cl_locations", data_key="locations_to_visit", label="Locations to Visit", placeholder="e.g. London and Edinburgh", section="trip"),
            FieldSchema(widget_key="uk_cl_highlight", data_key="trip_highlight", label="Trip Highlight", placeholder="e.g. attend an Orchestra Concert", field_type="text_area", section="trip"),

            # ── Applicant Section (nested under personal_details) ─────
            FieldSchema(widget_key="uk_cl_pd_name", data_key="name", label="Full Name", placeholder="e.g. John Doe", section="applicant", nested_under="personal_details"),
            FieldSchema(widget_key="uk_cl_pd_dob", data_key="dob", label="Date of Birth", placeholder="e.g. 3 October 2000", section="applicant", nested_under="personal_details"),
            FieldSchema(widget_key="uk_cl_pd_nationality", data_key="nationality", label="Nationality", placeholder="e.g. Indonesian", section="applicant", nested_under="personal_details"),
            FieldSchema(widget_key="uk_cl_pd_passport", data_key="passport_number", label="Passport Number", placeholder="e.g. X3236019", section="applicant", nested_under="personal_details"),
            FieldSchema(widget_key="uk_cl_home_country", data_key="home_country", label="Home Country", placeholder="e.g. Indonesia", section="applicant"),

            # ── Employment Section ────────────────────────────────────
            FieldSchema(widget_key="uk_cl_job_title", data_key="job_title", label="Job Title", placeholder="e.g. Front End Developer", section="employment"),
            FieldSchema(widget_key="uk_cl_company", data_key="company_name", label="Company Name", placeholder="e.g. PT Radya Anugrah", section="employment"),
            FieldSchema(widget_key="uk_cl_job_resp", data_key="job_responsibilities", label="Job Responsibilities", placeholder="e.g. designing and building web interfaces...", field_type="text_area", section="employment"),
            FieldSchema(widget_key="uk_cl_projects", data_key="project_details", label="Project Details (Ties)", placeholder="e.g. pharmaceutical tracking systems...", field_type="text_area", section="employment"),
            FieldSchema(widget_key="uk_cl_add_income", data_key="additional_income_source", label="Additional Income Source", placeholder="e.g. project-based freelance income", section="employment"),
            FieldSchema(widget_key="uk_cl_family", data_key="family_members", label="Family Members (Ties)", placeholder="e.g. parents", section="employment"),

            # ── Financials Section ────────────────────────────────────
            FieldSchema(widget_key="uk_cl_bank_names", data_key="bank_names", label="Bank Names", placeholder="e.g. BNI & Jago", section="financials"),
            FieldSchema(widget_key="uk_cl_acc_usage", data_key="account_usage_explanation", label="Account Usage Explanation", placeholder="e.g. BNI for payroll, Jago for daily operations", field_type="text_area", section="financials"),
            FieldSchema(widget_key="uk_cl_other_funds_src", data_key="other_funds_source", label="Other Funds Source", placeholder="e.g. Freelance Inflows", section="financials"),
            FieldSchema(widget_key="uk_cl_other_funds_exp", data_key="other_funds_explanation", label="Other Funds Explanation", placeholder="Explain freelance invoices...", field_type="text_area", section="financials"),
            FieldSchema(widget_key="uk_cl_turnover_rsn", data_key="turnover_reason", label="High Turnover Reason", placeholder="e.g. Investments", section="financials"),
            FieldSchema(widget_key="uk_cl_turnover_exp", data_key="turnover_explanation", label="High Turnover Explanation", placeholder="Explain high turnover...", field_type="text_area", section="financials"),
            FieldSchema(widget_key="uk_cl_outflow_rsn", data_key="outflow_reason", label="Large Outflow Reason", placeholder="e.g. One-off Asset", section="financials"),
            FieldSchema(widget_key="uk_cl_outflow_exp", data_key="outflow_explanation", label="Large Outflow Explanation", placeholder="Explain large outflow...", field_type="text_area", section="financials"),
            FieldSchema(widget_key="uk_cl_monthly_exp", data_key="monthly_expenditure_amount", label="Monthly Expenditure", placeholder="e.g. IDR 9,000,000", section="financials"),
            FieldSchema(widget_key="uk_cl_liquid", data_key="liquid_funds_amount", label="Liquid Funds Available", placeholder="e.g. IDR 75,000,000", section="financials"),

            # ── Contact & Documents Section ───────────────────────────
            FieldSchema(widget_key="uk_cl_email", data_key="email", label="Email", placeholder="e.g. m.irfan@gmail.com", section="contact"),
            FieldSchema(widget_key="uk_cl_phone", data_key="phone_number", label="Phone Number", placeholder="e.g. +6283829851734", section="contact"),
            FieldSchema(widget_key="uk_cl_add_docs", data_key="list_of_documents", label="List of Documents", placeholder="List any docs like '• Payroll slip'...", field_type="text_area", section="contact"),
        ],
        has_trip_type=True,
        has_group_members=True,
    ),
    DocumentType.COVER_LETTER_SCHENGEN: DocumentSchema(
        document_type=DocumentType.COVER_LETTER_SCHENGEN,
        title="Schengen Visa Cover Letter",
        icon="📝",
        sections={
            "trip":      "🌍 Trip Details",
            "applicant": "👤 Main Applicant Details",
            "contact":   "📋 Contact & Financials",
        },
        fields=[
            # ── Trip section ───────────────────────────────
            # data_key matches the prefill response keys returned by /prefill-document
            FieldSchema(widget_key="cl_country",    data_key="country",    label="Country of Embassy",  placeholder="e.g. Germany",          section="trip"),
            FieldSchema(widget_key="cl_city",        data_key="city",       label="City of Application", placeholder="e.g. Jakarta",          section="trip"),
            FieldSchema(widget_key="cl_purpose",     data_key="purpose",    label="Purpose of Trip",     placeholder="e.g. tourism",          section="trip"),
            FieldSchema(widget_key="cl_main_dest",   data_key="main_dest",  label="Main Destination",    placeholder="e.g. Germany",          section="trip"),
            FieldSchema(widget_key="cl_event",       data_key="event",      label="Event",               placeholder="e.g. personal vacation", section="trip"),
            FieldSchema(widget_key="cl_other_dest",  data_key="other_dest", label="Other Destinations",  placeholder="e.g. France and Italy",  section="trip"),
            FieldSchema(widget_key="cl_start",       data_key="start",      label="Travel Start Date",   placeholder="e.g. 2026-05-01",        section="trip"),
            FieldSchema(widget_key="cl_end",         data_key="end",        label="Travel End Date",     placeholder="e.g. 2026-05-14",        section="trip"),
            FieldSchema(widget_key="cl_duration",    data_key="duration",   label="Duration",            placeholder="e.g. 14 days",           section="trip"),
            # ── Applicant section (nested under personal_details) ──
            FieldSchema(widget_key="cl_pd_name",            data_key="name",            label="Full Name",       placeholder="e.g. John Doe",     section="applicant", nested_under="personal_details"),
            FieldSchema(widget_key="cl_pd_dob",             data_key="dob",             label="Date of Birth",   placeholder="e.g. 1st Jan 1990", section="applicant", nested_under="personal_details"),
            FieldSchema(widget_key="cl_pd_nationality",     data_key="nationality",     label="Nationality",     placeholder="e.g. Indonesian",   section="applicant", nested_under="personal_details"),
            FieldSchema(widget_key="cl_pd_occupation",      data_key="occupation",      label="Occupation",      placeholder="e.g. Engineer",     section="applicant", nested_under="personal_details"),
            FieldSchema(widget_key="cl_pd_passport_number", data_key="passport_number", label="Passport Number", placeholder="e.g. A1234567",     section="applicant", nested_under="personal_details"),
            # ── Contact section ────────────────────────────
            FieldSchema(widget_key="cl_trip_highlight",   data_key="trip_highlight",   label="Trip Highlight",   placeholder="Key highlights of your trip...", field_type="text_area", section="contact"),
            FieldSchema(widget_key="cl_contact_email",    data_key="contact_email",    label="Contact Email",    placeholder="e.g. john@email.com",            section="contact"),
            FieldSchema(widget_key="cl_contact_phone",    data_key="contact_phone",    label="Contact Phone",    placeholder="e.g. +62 812 345 6789",          section="contact"),
            FieldSchema(widget_key="cl_job_commitment",   data_key="job_commitment",   label="Job Commitment",   placeholder="e.g. returning to work on...",   section="contact"),
            FieldSchema(widget_key="cl_financial_status", data_key="financial_status", label="Financial Status", placeholder="e.g. sufficient funds",          section="contact"),
            FieldSchema(widget_key="cl_list_of_documents", data_key="list_of_documents", label="List of Documents", placeholder="e.g. passport, bank statement, hotel booking...", field_type="text_area", section="contact"),
        ],
        has_trip_type=True,
        has_group_members=True,
    ),
    DocumentType.LETTER_OF_AUTHORIZATION: DocumentSchema(
        document_type=DocumentType.LETTER_OF_AUTHORIZATION,
        title="Letter of Authorization for Visa",
        icon="✍️",
        sections={
            "grantor":        "👤 First Party (Grantor)",
            "representative": "🤝 Authorized Representative",
            "authorization":  "📝 Authorization Details",
        },
        fields=[
            # ── Grantor section ───────────────────────────────
            FieldSchema(widget_key="loa_grantor_name",    data_key="grantor_name",    label="Full Name based on Passport",     placeholder="e.g. Neni Diankrisna Putri", description="Full name of the applicant as it appears on their passport (maps to applicant_name or full_name from the profile). This person is granting authorization.", section="grantor"),
            FieldSchema(widget_key="loa_grantor_address", data_key="grantor_address", label="Address",                         placeholder="e.g. Jl. Mustika Jaya I...", description="Home address of the applicant (grantor).",                                                                                                                        section="grantor"),
            FieldSchema(widget_key="loa_grantor_id",      data_key="grantor_id",      label="National ID Number (Ktp)",        placeholder="e.g. 3175025510900006",      description="Indonesian national ID number (NIK/KTP) of the applicant (grantor).",                                                                                        section="grantor"),
            FieldSchema(widget_key="loa_grantor_city",    data_key="grantor_city",    label="City/Province",                   placeholder="e.g. Jakarta Timur",         description="City and/or province where the applicant (grantor) resides.",                                                                                                   section="grantor"),
            
            # ── Representative section ────────────────────────
            FieldSchema(widget_key="loa_rep_name",    data_key="rep_name",    label="Full Name",          placeholder="Select a representative",     section="representative", field_type="select", options=list(REPRESENTATIVES.keys())),
            FieldSchema(widget_key="loa_rep_id",      data_key="rep_id",      label="National ID Number", placeholder="e.g. 3174050905990008",       section="representative"),
            FieldSchema(widget_key="loa_rep_address", data_key="rep_address", label="Address",            placeholder="e.g. Jl. Kramat 1 No 15...",  section="representative"),
            FieldSchema(widget_key="loa_rep_city",    data_key="rep_city",    label="City/Province",      placeholder="e.g. Jakarta Selatan",        section="representative"),

            # ── Authorization section ─────────────────────────
            FieldSchema(widget_key="loa_passport_num",    data_key="passport_number",       label="Passport Number",  placeholder="e.g. X1373768",                       section="authorization"),
            FieldSchema(widget_key="loa_passport_name",   data_key="passport_name",         label="Passport Name",    placeholder="e.g. Neni Diankrisna Putri",          section="authorization"),
            FieldSchema(widget_key="loa_place",           data_key="city",                  label="Signing City",     placeholder="e.g. Jakarta",                        section="authorization"),
            FieldSchema(widget_key="loa_date",            data_key="date",                  label="Signing Date",     placeholder="e.g. 13 February 2026",               section="authorization"),
        ],
        has_trip_type=False,
        has_group_members=False,
    ),
    DocumentType.LETTER_OF_AUTHORIZATION_MINOR: DocumentSchema(
        document_type=DocumentType.LETTER_OF_AUTHORIZATION_MINOR,
        title="Letter of Authorization for Minor's Visa",
        icon="🧒",
        sections={
            "parents":        "👪 Parents (First Party)",
            "representative": "🤝 Authorized Representative",
            "authorization":  "📝 Authorization Details",
        },
        fields=(
            # ── Parents (Grantors) section ────────────────────────
            generate_person_fields("father", "parents", "Father's") +
            generate_person_fields("mother", "parents", "Mother's") +
            
            # ── Representative section (name is a dropdown, rest auto-filled) ──
            [FieldSchema(widget_key="loa_minor_rep_name", data_key="rep_name", label="Representative's Name",
                         section="representative", field_type="select", options=list(REPRESENTATIVES.keys()))] +
            generate_person_fields("rep", "representative", "Representative's")[1:] +

            # ── Authorization & Minor section ─────────────────────
            [
                FieldSchema(widget_key="loa_minor_passport",    data_key="passport_number", label="Child's Passport Number",  placeholder="e.g. X2588429",                section="authorization"),
                FieldSchema(widget_key="loa_minor_pass_name",   data_key="passport_name",   label="Child's Passport Name",    placeholder="e.g. Mohamad Noah...",         section="authorization"),
                FieldSchema(widget_key="loa_minor_place",       data_key="city",            label="Signing City",             placeholder="e.g. Jakarta",                 section="authorization"),
                FieldSchema(widget_key="loa_minor_date",        data_key="date",            label="Signing Date",             placeholder="e.g. 24 November 2025",        section="authorization"),
            ]
        ),
        has_trip_type=False,
        has_group_members=False,
    ),
}


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

class PrefillDocumentRequest(BaseModel):
    """
    Sent to /prefill-document (structure mode).
    `structure` describes the fields the LLM should map from DB data.
    Returns a flat dict of { field_key: value_or_null, ..., "_missing_required": [...] }.
    """
    application_id: int
    structure: List[Dict[str, Any]]


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_content: str = ""
    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]
    structure: Optional[List[Dict[str, Any]]] = None  # field definitions for core API handover
    session_uuid: Optional[str] = None


class ExportDocumentRequest(BaseModel):
    """Sent to /export-document once the user is satisfied with the draft."""
    document_content: str
    title: Optional[str] = None
    export_format: str = "gdocs"


# ── 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