feat: cover letter--schengen

src/document_generator_models.py

@@ -8,9 +8,10 @@ from enum import StrEnum
 # Each value must match the string the backend expects.
 
 class DocumentType(StrEnum):
-    COVER_LETTER = "cover_letter"
-    # INVITATION_LETTER = "invitation_letter"
-    # BANK_REFERENCE   = "bank_reference"
+    COVER_LETTER_SCHENGEN = "cover_letter_schengen"
+    # COVER_LETTER_UK      = "cover_letter_uk"
+    # INVITATION_LETTER    = "invitation_letter"
+    # BANK_REFERENCE       = "bank_reference"
 
 
 # ── Field & schema descriptors ─────────────────────────────
@@ -64,36 +65,36 @@ class DocumentSchema(BaseModel):
 # To add a new document type:
 #   1. Add its value to DocumentType above.
 #   2. Add a DocumentSchema entry here.
-#   3. Add a tab + form in 03_Document_Generator.py.
-#   The backend receives `document_name` and generates the SQL via LLM.
+#   The frontend auto-generates a tab and form from the schema — no UI code needed.
 
 DOCUMENT_REGISTRY: Dict[str, DocumentSchema] = {
-    DocumentType.COVER_LETTER: DocumentSchema(
-        document_type=DocumentType.COVER_LETTER,
+    DocumentType.COVER_LETTER_SCHENGEN: DocumentSchema(
+        document_type=DocumentType.COVER_LETTER_SCHENGEN,
         title="Schengen Visa Cover Letter",
         icon="📝",
         sections={
             "trip":      "🌍 Trip Details",
-            "applicant": "👤 Applicant Details",
+            "applicant": "👤 Main Applicant Details",
             "contact":   "📋 Contact & Financials",
         },
         fields=[
             # ── Trip section ───────────────────────────────
-            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_of_application", 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="destination_country", 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_date",          label="Travel Start Date",   placeholder="e.g. 2026-05-01",          section="trip"),
-            FieldSchema(widget_key="cl_end",         data_key="end_date",            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"),
+            # 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"),
+            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"),
@@ -111,12 +112,12 @@ DOCUMENT_REGISTRY: Dict[str, DocumentSchema] = {
 
 class PrefillDocumentRequest(BaseModel):
     """
-    Sent to /prefill-document.
-    Backend uses `document_name` to determine which SQL to run (via LLM)
-    for the given application_id.
+    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
-    document_name: str  # value from DocumentType
+    structure: List[Dict[str, Any]]
 
 
 class GenerateDocumentRequest(BaseModel):
@@ -132,13 +133,31 @@ class GenerateDocumentRequest(BaseModel):
 
 
 class DocChatRequest(BaseModel):
-    """Sent to /generate-document/chat."""
+    """Sent to /generate-document/chat (new stateless draft-revision flow)."""
     query: str
     history: List[Dict[str, str]]
-    language: str = "ID"
-    session_uuid: str
-    application_id: Optional[int] = None
-    doc_data: Optional[Dict[str, Any]] = None
+    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) ──

src/pages/03_Document_Generator.py

@@ -11,10 +11,12 @@ from dotenv import load_dotenv
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
 from src.document_generator_models import (
     DocumentType,
+    DocumentSchema,
     DOCUMENT_REGISTRY,
     PrefillDocumentRequest,
-    GenerateDocumentRequest,
     DocChatRequest,
+    GenerateDraftRequest,
+    ExportDocumentRequest,
 )
 
 load_dotenv()
@@ -22,97 +24,111 @@ load_dotenv()
 API_BASE_URL = os.getenv("URL_CE_BOT", "http://localhost:5000")
 TODAY = datetime.date.today()
 
+
 # ─────────────────────────────────────────────────────────
-#  COVER LETTER (CL) SESSION STATE & HELPERS
+#  GENERIC SESSION STATE
+#  Keys: doc_{doc_type}_{suffix}  (one set per registered doc type)
 # ─────────────────────────────────────────────────────────
 
-# Widget keys derived from the registry — no need to update manually when fields change
-CL_FORM_KEYS = DOCUMENT_REGISTRY[DocumentType.COVER_LETTER].widget_keys()
+def _sk(doc_type: str, suffix: str) -> str:
+    """Compose a session-state key from doc type + suffix."""
+    return f"doc_{doc_type}_{suffix}"
 
 
-def _init_cl_state():
-    defaults = {
-        "cl_data": None,             # Current cover letter form data
-        "cl_chat_messages": [],      # Chat history (display)
-        "cl_api_history": [],        # Simplified history (API)
-        "cl_history_stack": [],      # Undo stack
-        "cl_redo_stack": [],         # Redo stack
-        "cl_session_uuid": str(uuid.uuid4()),  # Server-side session ID
-        "cl_app_id": "",             # Last successfully loaded application ID
-        "cl_doc_url": None,          # Generated document URL
-    }
-    for key, val in defaults.items():
-        if key not in st.session_state:
-            st.session_state[key] = val
+def _init_all_doc_states():
+    """Initialize session-state slots for every registered document type."""
+    for doc_type in DOCUMENT_REGISTRY:
+        defaults = {
+            _sk(doc_type, "data"):             None,
+            _sk(doc_type, "chat_messages"):    [],
+            _sk(doc_type, "api_history"):      [],
+            _sk(doc_type, "history_stack"):    [],
+            _sk(doc_type, "redo_stack"):       [],
+            _sk(doc_type, "session_uuid"):     str(uuid.uuid4()),
+            _sk(doc_type, "app_id"):           "",
+            _sk(doc_type, "draft_content"):    "",
+            _sk(doc_type, "draft_generated"):  False,
+            _sk(doc_type, "doc_structure"):    None,
+            _sk(doc_type, "doc_url"):          None,
+        }
+        for key, val in defaults.items():
+            if key not in st.session_state:
+                st.session_state[key] = val
 
 
-_init_cl_state()
+_init_all_doc_states()
 
 
+# ─────────────────────────────────────────────────────────
+#  UNDO / REDO
+# ─────────────────────────────────────────────────────────
 
-def _cl_clear_form_keys():
-    """Delete widget session-state keys so they re-initialize from `value=` on next render."""
-    for key in CL_FORM_KEYS:
+def _clear_widget_keys(doc_type: str):
+    """Pop all widget keys for a doc type so widgets reset from value= on next render."""
+    for key in DOCUMENT_REGISTRY[doc_type].widget_keys():
         st.session_state.pop(key, None)
 
 
-def cl_push_history(data):
-    if data is not None:
-        st.session_state.cl_history_stack.append({
-            "data": copy.deepcopy(data),
-            "chat": copy.deepcopy(st.session_state.cl_chat_messages),
-            "api": copy.deepcopy(st.session_state.cl_api_history),
-        })
-        st.session_state.cl_redo_stack.clear()
-
-
-def cl_undo():
-    if st.session_state.cl_history_stack:
-        st.session_state.cl_redo_stack.append({
-            "data": copy.deepcopy(st.session_state.cl_data),
-            "chat": copy.deepcopy(st.session_state.cl_chat_messages),
-            "api": copy.deepcopy(st.session_state.cl_api_history),
-        })
-        prev = st.session_state.cl_history_stack.pop()
-        st.session_state.cl_data = prev["data"]
-        st.session_state.cl_chat_messages = prev["chat"]
-        st.session_state.cl_api_history = prev["api"]
-        _cl_clear_form_keys()
-
-
-def cl_redo():
-    if st.session_state.cl_redo_stack:
-        st.session_state.cl_history_stack.append({
-            "data": copy.deepcopy(st.session_state.cl_data),
-            "chat": copy.deepcopy(st.session_state.cl_chat_messages),
-            "api": copy.deepcopy(st.session_state.cl_api_history),
-        })
-        nxt = st.session_state.cl_redo_stack.pop()
-        st.session_state.cl_data = nxt["data"]
-        st.session_state.cl_chat_messages = nxt["chat"]
-        st.session_state.cl_api_history = nxt["api"]
-        _cl_clear_form_keys()
-
-
-def call_doc_chat_api(query, history, application_id, session_uuid, doc_data=None):
-    """Call /generate-document/chat endpoint."""
+def push_history(doc_type: str, data):
+    if data is None:
+        return
+    st.session_state[_sk(doc_type, "history_stack")].append({
+        "data": copy.deepcopy(data),
+        "chat": copy.deepcopy(st.session_state[_sk(doc_type, "chat_messages")]),
+        "api":  copy.deepcopy(st.session_state[_sk(doc_type, "api_history")]),
+    })
+    st.session_state[_sk(doc_type, "redo_stack")].clear()
+
+
+def do_undo(doc_type: str):
+    history = st.session_state[_sk(doc_type, "history_stack")]
+    if not history:
+        return
+    st.session_state[_sk(doc_type, "redo_stack")].append({
+        "data": copy.deepcopy(st.session_state[_sk(doc_type, "data")]),
+        "chat": copy.deepcopy(st.session_state[_sk(doc_type, "chat_messages")]),
+        "api":  copy.deepcopy(st.session_state[_sk(doc_type, "api_history")]),
+    })
+    prev = history.pop()
+    st.session_state[_sk(doc_type, "data")]          = prev["data"]
+    st.session_state[_sk(doc_type, "chat_messages")] = prev["chat"]
+    st.session_state[_sk(doc_type, "api_history")]   = prev["api"]
+    _clear_widget_keys(doc_type)
+
+
+def do_redo(doc_type: str):
+    redo = st.session_state[_sk(doc_type, "redo_stack")]
+    if not redo:
+        return
+    st.session_state[_sk(doc_type, "history_stack")].append({
+        "data": copy.deepcopy(st.session_state[_sk(doc_type, "data")]),
+        "chat": copy.deepcopy(st.session_state[_sk(doc_type, "chat_messages")]),
+        "api":  copy.deepcopy(st.session_state[_sk(doc_type, "api_history")]),
+    })
+    nxt = redo.pop()
+    st.session_state[_sk(doc_type, "data")]          = nxt["data"]
+    st.session_state[_sk(doc_type, "chat_messages")] = nxt["chat"]
+    st.session_state[_sk(doc_type, "api_history")]   = nxt["api"]
+    _clear_widget_keys(doc_type)
+
+
+# ─────────────────────────────────────────────────────────
+#  API HELPERS  (doc-type-agnostic)
+# ─────────────────────────────────────────────────────────
+
+def call_doc_chat_api(query, history, session_uuid, current_document_content="", structure=None):
+    """POST /generate-document/chat — stateless revision."""
     try:
-        app_id_int = None
-        if application_id:
-            try:
-                app_id_int = int(application_id)
-            except (ValueError, TypeError):
-                pass
-        request = DocChatRequest(
+        req = DocChatRequest(
             query=query,
             history=history,
             session_uuid=session_uuid,
-            application_id=app_id_int,
-            doc_data=doc_data,
+            current_document_content=current_document_content,
+            structure=structure,
         )
         resp = requests.post(
             f"{API_BASE_URL}/generate-document/chat",
-            json=request.model_dump(exclude_none=True),
+            json=req.model_dump(exclude_none=True),
             timeout=120,
         )
         resp.raise_for_status()
@@ -121,87 +137,238 @@ def call_doc_chat_api(query, history, application_id, session_uuid, doc_data=Non
         return None, str(e)
 
 
-def call_generate_document(app_id, document_name, form_data):
-    """Call /generate-document endpoint.
-    `document_name` tells the backend which SQL to generate (via LLM).
-    `form_data` is forwarded as `doc_data` for the backend to merge on top of DB data.
-    """
+def call_generate_draft(schema: DocumentSchema, form_data, session_uuid=None):
+    """POST /generate-document/draft — returns Markdown, no Google Doc created."""
     try:
-        request = GenerateDocumentRequest(
-            application_id=int(app_id),
-            document_name=document_name,
-            doc_data=form_data or None,
+        structure = [
+            {"key": f.data_key, "label": f.label, "required": False}
+            for f in schema.fields
+        ]
+        req = GenerateDraftRequest(
+            doc_type=str(schema.document_type),
+            data=form_data or {},
+            structure=structure,
+            session_uuid=session_uuid,
         )
         resp = requests.post(
-            f"{API_BASE_URL}/generate-document",
-            json=request.model_dump(exclude_none=True),
+            f"{API_BASE_URL}/generate-document/draft",
+            json=req.model_dump(exclude_none=True),
             timeout=120,
         )
         resp.raise_for_status()
         return resp.json(), None
-    except (ValueError, TypeError) as e:
-        return None, f"Invalid application_id: {e}"
     except Exception as e:
         return None, str(e)
 
 
+def call_export_document(document_content, title=None):
+    """POST /export-document — converts Markdown to Google Doc."""
+    try:
+        req = ExportDocumentRequest(document_content=document_content, title=title)
+        resp = requests.post(
+            f"{API_BASE_URL}/export-document",
+            json=req.model_dump(exclude_none=True),
+            timeout=120,
+        )
+        resp.raise_for_status()
+        return resp.json(), None
+    except Exception as e:
+        return None, str(e)
 
 
 # ─────────────────────────────────────────────────────────
-#  PAGE CONFIG & TABS
+#  FORM RENDERING  (driven by DocumentSchema metadata)
 # ─────────────────────────────────────────────────────────
 
-st.set_page_config(page_title="Document Generator", layout="wide")
-st.title("🌍 Document Generator")
+def _get_field_value(data: dict, field) -> str:
+    """Read a field's current value, respecting nested_under."""
+    if field.nested_under:
+        return (data.get(field.nested_under) or {}).get(field.data_key, "") or ""
+    return data.get(field.data_key, "") or ""
+
+
+def _apply_field_change(updated_data: dict, original_data: dict, field, new_val: str):
+    """Write a changed value into updated_data, respecting nested_under."""
+    if field.nested_under:
+        if updated_data.get(field.nested_under) is None:
+            updated_data[field.nested_under] = copy.deepcopy(
+                original_data.get(field.nested_under) or {}
+            )
+        updated_data[field.nested_under][field.data_key] = new_val
+    else:
+        updated_data[field.data_key] = new_val
+
+
+def _render_section(schema: DocumentSchema, section_key: str, data: dict, updated_data: dict) -> bool:
+    """
+    Render all fields for one section.
+    Returns True if any field value changed.
+    """
+    st.subheader(schema.sections[section_key])
+    changed = False
+    for field in schema.fields_in_section(section_key):
+        cur = _get_field_value(data, field)
+        widget = st.text_area if field.field_type == "text_area" else st.text_input
+        new = widget(field.label, value=cur, placeholder=field.placeholder, key=field.widget_key)
+        if new != cur:
+            _apply_field_change(updated_data, data, field, new)
+            changed = True
+    return changed
+
+
+def render_doc_form(schema: DocumentSchema):
+    """
+    Render the editable form driven entirely by DocumentSchema.
 
-tab1, = st.tabs([
-    "📝 Cover Letter",
-])
+    Layout rules:
+      - 1 section  → full width
+      - 2+ sections → first N//2 sections in left column, remainder in right column
 
+    Special elements:
+      - has_trip_type=True  → radio above columns
+      - has_group_members=True → data editor below columns (shown when trip_type=="Group"
+                                 or when the schema has no trip_type toggle)
+    """
+    doc_type = str(schema.document_type)
+    data = st.session_state[_sk(doc_type, "data")]
+    updated_data = copy.deepcopy(data)
+    form_changed = False
+
+    st.subheader("✏️ Edit Details")
+
+    # ── Trip type radio ──────────────────────────────────
+    trip_type = None
+    if schema.has_trip_type:
+        opts = ["Individual", "Group"]
+        cur = data.get("trip_type", "Individual")
+        if cur not in opts:
+            cur = "Individual"
+        trip_type = st.radio(
+            "Trip Type", opts,
+            index=opts.index(cur),
+            horizontal=True,
+            key=f"{doc_type}_trip_type",
+        )
+        if trip_type != cur:
+            updated_data["trip_type"] = trip_type
+            form_changed = True
+
+    # ── Section columns ──────────────────────────────────
+    section_keys = list(schema.sections.keys())
+    if len(section_keys) == 1:
+        changed = _render_section(schema, section_keys[0], data, updated_data)
+        form_changed = form_changed or changed
+    else:
+        mid = len(section_keys) // 2          # e.g. 3 sections → mid=1 (left:1, right:2)
+        left_sections  = section_keys[:mid]
+        right_sections = section_keys[mid:]
+        col1, col2 = st.columns(2)
+        with col1:
+            for sk in left_sections:
+                changed = _render_section(schema, sk, data, updated_data)
+                form_changed = form_changed or changed
+        with col2:
+            for sk in right_sections:
+                changed = _render_section(schema, sk, data, updated_data)
+                form_changed = form_changed or changed
+
+    # ── Group members data editor ────────────────────────
+    if schema.has_group_members:
+        show_group = (trip_type == "Group") if schema.has_trip_type else True
+        if show_group:
+            st.subheader("👥 Group Members")
+            gm_list = data.get("group_members", []) or []
+            gm_cols = ["relationship", "name", "dob", "occupation", "nationality", "passport_number"]
+            gm_df = pd.DataFrame(
+                gm_list or [dict.fromkeys(gm_cols, "")],
+                columns=gm_cols,
+            )
+            edited_gm = st.data_editor(
+                gm_df,
+                num_rows="dynamic",
+                use_container_width=True,
+                key=f"{doc_type}_group_members_editor",
+            )
+            new_gm = edited_gm.to_dict("records")
+            if new_gm != gm_list:
+                updated_data["group_members"] = new_gm
+                form_changed = True
+
+    # ── Persist changes ──────────────────────────────────
+    if form_changed:
+        push_history(doc_type, data)
+        st.session_state[_sk(doc_type, "data")] = updated_data
 
-# ═══════════════════════════════════════════════════════════
-#  TAB 1: COVER LETTER  (undo/redo + chat, data from DB)
-# ═══════════════════════════════════════════════════════════
 
-with tab1:
-    st.header("📝 Schengen Visa Cover Letter")
-    st.caption("Load applicant data by Application ID, edit the fields directly, chat to revise, then generate the Google Doc.")
+# ─────────────────────────────────────────────────────────
+#  FULL TAB RENDERER
+# ─────────────────────────────────────────────────────────
+
+def render_document_tab(schema: DocumentSchema):
+    """Render the complete UI for one document type inside its tab."""
+    doc_type = str(schema.document_type)
+
+    def sk(suffix: str) -> str:
+        return _sk(doc_type, suffix)
+
+    st.header(f"{schema.icon} {schema.title}")
+    st.caption(
+        "Load applicant data by Application ID, edit the fields directly, "
+        "chat to revise, then export the Google Doc."
+    )
 
-    # ── Application ID row ──
+    # ── Application ID row ──────────────────────────────
     load_c1, load_c2, load_c3 = st.columns([2, 1, 3])
     with load_c1:
         app_id_input = st.text_input(
             "Application ID",
-            value=st.session_state.cl_app_id,
+            value=st.session_state[sk("app_id")],
             placeholder="Enter Application ID (e.g. 5786)",
-            key="cl_app_id_input",
+            key=f"{doc_type}_app_id_input",
             label_visibility="collapsed",
         )
     with load_c2:
-        load_clicked = st.button("🔍 Load from DB", use_container_width=True, key="cl_load_btn")
+        load_clicked = st.button(
+            "🔍 Load from DB",
+            use_container_width=True,
+            key=f"{doc_type}_load_btn",
+        )
     with load_c3:
-        if st.session_state.cl_data is not None and st.session_state.cl_app_id:
+        if st.session_state[sk("data")] is not None and st.session_state[sk("app_id")]:
             info_c, clear_c = st.columns([4, 1])
             with info_c:
-                st.caption(f"✅ Loaded: Application ID **{st.session_state.cl_app_id}**")
+                st.caption(f"✅ Loaded: Application ID **{st.session_state[sk('app_id')]}**")
             with clear_c:
-                if st.button("✕", key="cl_clear_btn", help="Clear and start over"):
-                    cl_push_history(st.session_state.cl_data)
-                    _cl_clear_form_keys()
-                    st.session_state.cl_data = None
-                    st.session_state.cl_app_id = ""
-                    st.session_state.cl_doc_url = None
-                    st.session_state.cl_chat_messages = []
-                    st.session_state.cl_api_history = []
-                    st.session_state.cl_session_uuid = str(uuid.uuid4())
+                if st.button("✕", key=f"{doc_type}_clear_btn", help="Clear and start over"):
+                    push_history(doc_type, st.session_state[sk("data")])
+                    _clear_widget_keys(doc_type)
+                    st.session_state[sk("data")]            = None
+                    st.session_state[sk("app_id")]          = ""
+                    st.session_state[sk("doc_url")]         = None
+                    st.session_state[sk("draft_content")]   = ""
+                    st.session_state[sk("draft_generated")] = False
+                    st.session_state[sk("doc_structure")]   = None
+                    st.session_state[sk("chat_messages")]   = []
+                    st.session_state[sk("api_history")]     = []
+                    st.session_state[sk("session_uuid")]    = str(uuid.uuid4())
                     st.rerun()
 
+    # ── Load from DB ────────────────────────────────────
     if load_clicked:
         if app_id_input:
             try:
+                structure = [
+                    {
+                        "key": f.data_key,
+                        "label": f.label,
+                        "description": f.placeholder,
+                        "required": False,
+                    }
+                    for f in schema.fields
+                ]
                 prefill_req = PrefillDocumentRequest(
                     application_id=int(app_id_input),
-                    document_name=DocumentType.COVER_LETTER,
+                    structure=structure,
                 )
                 resp = requests.post(
                     f"{API_BASE_URL}/prefill-document",
@@ -209,15 +376,25 @@ with tab1:
                     timeout=60,
                 )
                 if resp.status_code == 200:
-                    fetched = resp.json()
-                    cl_push_history(st.session_state.cl_data)
-                    _cl_clear_form_keys()
-                    st.session_state.cl_data = fetched
-                    st.session_state.cl_app_id = app_id_input
-                    st.session_state.cl_doc_url = None
-                    st.session_state.cl_chat_messages = []
-                    st.session_state.cl_api_history = []
-                    st.session_state.cl_session_uuid = str(uuid.uuid4())
+                    flat = resp.json()
+                    # Drop metadata keys (_missing_required, etc.) and reconstruct
+                    # nested fields (e.g. personal_details.name) from the flat response.
+                    fetched = {k: v for k, v in flat.items() if not k.startswith("_")}
+                    for field in schema.fields:
+                        if field.nested_under and field.data_key in fetched:
+                            fetched.setdefault(field.nested_under, {})
+                            fetched[field.nested_under][field.data_key] = fetched.pop(field.data_key)
+                    push_history(doc_type, st.session_state[sk("data")])
+                    _clear_widget_keys(doc_type)
+                    st.session_state[sk("data")]            = fetched
+                    st.session_state[sk("app_id")]          = app_id_input
+                    st.session_state[sk("doc_url")]         = None
+                    st.session_state[sk("draft_content")]   = ""
+                    st.session_state[sk("draft_generated")] = False
+                    st.session_state[sk("doc_structure")]   = None
+                    st.session_state[sk("chat_messages")]   = []
+                    st.session_state[sk("api_history")]     = []
+                    st.session_state[sk("session_uuid")]    = str(uuid.uuid4())
                     st.success(f"✅ Loaded data for Application ID: {app_id_input}")
                     st.rerun()
                 else:
@@ -235,205 +412,203 @@ with tab1:
 
     st.divider()
 
-    # ── Chat history display ──
-    for msg in st.session_state.cl_chat_messages:
-        with st.chat_message(msg["role"]):
-            st.markdown(msg["content"])
-
-    if not st.session_state.cl_data and not st.session_state.cl_chat_messages:
-        st.info("📌 Enter an Application ID above to load data from the database, or use the chat below to describe your cover letter needs.")
-
-    # ── Generated document URL ──
-    if st.session_state.cl_doc_url:
-        st.success("🎉 Document Generated!")
-        if "docs.google.com" in st.session_state.cl_doc_url:
-            st.link_button("🌐 Open in Google Docs", st.session_state.cl_doc_url, use_container_width=True)
-        else:
-            st.markdown(f"[📄 View Document]({st.session_state.cl_doc_url})")
+    if not st.session_state[sk("data")]:
+        st.info("📌 Enter an Application ID above to load data from the database.")
+        return
 
-    # ── Editable form (shown when data is loaded) ──
-    if st.session_state.cl_data is not None:
-        data = st.session_state.cl_data
-        updated_data = copy.deepcopy(data)
-        form_changed = False
+    # ── Editable form ────────────────────────────────────
+    render_doc_form(schema)
 
-        st.subheader("✏️ Edit Details")
-
-        # Trip type
-        trip_type_opts = ["Individual", "Group"]
-        cur_trip = data.get("trip_type", "Individual")
-        if cur_trip not in trip_type_opts:
-            cur_trip = "Individual"
-        trip_type = st.radio(
-            "Trip Type", trip_type_opts,
-            index=trip_type_opts.index(cur_trip),
-            horizontal=True, key="cl_trip_type",
+    # ── Toolbar: Undo / Redo / Generate Draft ───────────
+    st.divider()
+    tb = st.columns([1, 1, 2, 4, 2])
+    with tb[0]:
+        st.button(
+            "↩️ Undo",
+            on_click=do_undo,
+            args=(doc_type,),
+            disabled=not st.session_state[sk("history_stack")],
+            use_container_width=True,
+            key=f"{doc_type}_undo_btn",
+        )
+    with tb[1]:
+        st.button(
+            "↪️ Redo",
+            on_click=do_redo,
+            args=(doc_type,),
+            disabled=not st.session_state[sk("redo_stack")],
+            use_container_width=True,
+            key=f"{doc_type}_redo_btn",
+        )
+    with tb[2]:
+        h = len(st.session_state[sk("history_stack")])
+        st.caption(f"📜 {h} version{'s' if h != 1 else ''}")
+    with tb[4]:
+        draft_clicked = st.button(
+            "📝 Generate Draft",
+            type="primary",
+            use_container_width=True,
+            key=f"{doc_type}_gen_btn",
         )
-        if trip_type != cur_trip:
-            updated_data["trip_type"] = trip_type
-            form_changed = True
-
-        col1, col2 = st.columns(2)
-
-        with col1:
-            st.subheader("🌍 Trip Details")
-            for wk, dk, label, ph in [
-                ("cl_country",    "country",              "Country of Embassy",    "e.g. Germany"),
-                ("cl_city",       "city_of_application",  "City of Application",   "e.g. Jakarta"),
-                ("cl_purpose",    "purpose",              "Purpose of Trip",       "e.g. tourism"),
-                ("cl_main_dest",  "destination_country",  "Main Destination",      "e.g. Germany"),
-                ("cl_event",      "event",                "Event",                 "e.g. personal vacation"),
-                ("cl_other_dest", "other_dest",           "Other Destinations",    "e.g. France and Italy"),
-                ("cl_start",      "start_date",           "Travel Start Date",     "e.g. 2026-05-01"),
-                ("cl_end",        "end_date",             "Travel End Date",       "e.g. 2026-05-14"),
-                ("cl_duration",   "duration",             "Duration",              "e.g. 14 days"),
-            ]:
-                cur = data.get(dk, "") or ""
-                new = st.text_input(label, value=cur, placeholder=ph, key=wk)
-                if new != cur:
-                    updated_data[dk] = new
-                    form_changed = True
 
-        with col2:
-            st.subheader("👤 Applicant Details")
-            pd_data = data.get("personal_details", {}) or {}
-            for wk, fk, label, ph in [
-                ("cl_pd_name",            "name",            "Full Name",       "e.g. John Doe"),
-                ("cl_pd_dob",             "dob",             "Date of Birth",   "e.g. 1st Jan 1990"),
-                ("cl_pd_nationality",     "nationality",     "Nationality",     "e.g. Indonesian"),
-                ("cl_pd_occupation",      "occupation",      "Occupation",      "e.g. Engineer"),
-                ("cl_pd_passport_number", "passport_number", "Passport Number", "e.g. A1234567"),
-            ]:
-                cur = pd_data.get(fk, "") or ""
-                new = st.text_input(label, value=cur, placeholder=ph, key=wk)
-                if new != cur:
-                    if updated_data.get("personal_details") is None:
-                        updated_data["personal_details"] = copy.deepcopy(pd_data)
-                    updated_data["personal_details"][fk] = new
-                    form_changed = True
-
-            st.subheader("📋 Contact & Financials")
-            for wk, dk, label, ph, is_area in [
-                ("cl_trip_highlight",   "trip_highlight",   "Trip Highlight",   "Key highlights of your trip...",  True),
-                ("cl_contact_email",    "contact_email",    "Contact Email",    "e.g. john@email.com",             False),
-                ("cl_contact_phone",    "contact_phone",    "Contact Phone",    "e.g. +62 812 345 6789",           False),
-                ("cl_job_commitment",   "job_commitment",   "Job Commitment",   "e.g. returning to work on...",    False),
-                ("cl_financial_status", "financial_status", "Financial Status", "e.g. sufficient funds",           False),
-            ]:
-                cur = data.get(dk, "") or ""
-                new = (st.text_area if is_area else st.text_input)(label, value=cur, placeholder=ph, key=wk)
-                if new != cur:
-                    updated_data[dk] = new
-                    form_changed = True
-
-        # Group members table
-        if trip_type == "Group":
-            st.subheader("👥 Group Members")
-            gm_list = data.get("group_members", []) or []
-            gm_df = pd.DataFrame(
-                gm_list or [{"relationship": "", "name": "", "dob": "", "occupation": "", "nationality": "", "passport_number": ""}],
-                columns=["relationship", "name", "dob", "occupation", "nationality", "passport_number"],
+    if draft_clicked:
+        with st.spinner("Generating document draft..."):
+            result, err = call_generate_draft(
+                schema,
+                st.session_state[sk("data")],
+                session_uuid=st.session_state[sk("session_uuid")],
             )
-            edited_gm = st.data_editor(gm_df, num_rows="dynamic", use_container_width=True, key="cl_group_members_editor")
-            new_gm = edited_gm.to_dict("records")
-            if new_gm != gm_list:
-                updated_data["group_members"] = new_gm
-                form_changed = True
-
-        # Apply inline form changes
-        if form_changed:
-            cl_push_history(st.session_state.cl_data)
-            st.session_state.cl_data = updated_data
-            data = updated_data
+        if err:
+            st.error(f"❌ Error generating draft: {err}")
+        elif result and result.get("document_content"):
+            st.session_state[sk("draft_content")]   = result["document_content"]
+            st.session_state[sk("draft_generated")] = True
+            # Store field structure for future core API handover
+            st.session_state[sk("doc_structure")] = [
+                {"key": f.data_key, "label": f.label, "required": False}
+                for f in schema.fields
+            ]
+            st.session_state[sk("chat_messages")] = []
+            st.session_state[sk("api_history")]   = []
+            st.rerun()
+        else:
+            st.error("❌ Failed to generate draft.")
 
-        # ── Undo / Redo / Generate Toolbar ──
+    # ── Draft preview + export ───────────────────────────
+    if st.session_state[sk("draft_generated")]:
         st.divider()
-        tb = st.columns([1, 1, 2, 4, 2])
-        with tb[0]:
-            st.button(
-                "↩️ Undo", on_click=cl_undo,
-                disabled=not st.session_state.cl_history_stack,
-                use_container_width=True, key="cl_undo_btn",
-            )
-        with tb[1]:
-            st.button(
-                "↪️ Redo", on_click=cl_redo,
-                disabled=not st.session_state.cl_redo_stack,
-                use_container_width=True, key="cl_redo_btn",
+        st.subheader("📄 Document Draft")
+        st.caption(
+            "Data values are **highlighted in bold**. "
+            "Use the chat below to revise, or expand the raw editor to edit directly."
+        )
+
+        # Rendered markdown preview (bold highlights are visible here)
+        with st.container(border=True):
+            st.markdown(st.session_state[sk("draft_content")])
+
+        # Collapsible raw Markdown editor for direct edits
+        with st.expander("✏️ Edit Raw Markdown"):
+            edited_raw = st.text_area(
+                "Raw editor",
+                value=st.session_state[sk("draft_content")],
+                height=420,
+                key=f"{doc_type}_draft_editor",
+                label_visibility="collapsed",
             )
-        with tb[2]:
-            h = len(st.session_state.cl_history_stack)
-            st.caption(f"📜 {h} version{'s' if h != 1 else ''}")
-        with tb[4]:
-            generate_clicked = st.button(
-                "🌐 Generate Document", type="primary",
-                use_container_width=True, key="cl_gen_btn",
+            if edited_raw != st.session_state[sk("draft_content")]:
+                st.session_state[sk("draft_content")] = edited_raw
+                st.rerun()
+
+        export_col, link_col = st.columns([2, 3])
+        with export_col:
+            export_clicked = st.button(
+                "🌐 Export to Google Docs",
+                type="primary",
+                use_container_width=True,
+                key=f"{doc_type}_export_btn",
             )
+        with link_col:
+            if st.session_state[sk("doc_url")]:
+                st.link_button(
+                    "📂 Open Exported Doc",
+                    st.session_state[sk("doc_url")],
+                    use_container_width=True,
+                )
 
-        if generate_clicked:
-            if not st.session_state.cl_app_id:
-                st.error("⚠️ Please load an Application ID first to generate the document.")
+        if export_clicked:
+            data = st.session_state[sk("data")] or {}
+            applicant_name = (data.get("personal_details") or {}).get("name") or ""
+            doc_label = schema.title
+            export_title = (
+                f"{doc_label} – {applicant_name}".strip(" –")
+                if applicant_name
+                else doc_label
+            )
+            with st.spinner("Exporting to Google Docs..."):
+                result, err = call_export_document(
+                    st.session_state[sk("draft_content")],
+                    title=export_title,
+                )
+            if err:
+                st.error(f"❌ Export failed: {err}")
+            elif result and result.get("url"):
+                st.session_state[sk("doc_url")] = result["url"]
+                st.rerun()
             else:
-                with st.spinner("Generating Google Doc... This may take a moment."):
-                    result, err = call_generate_document(
-                        st.session_state.cl_app_id,
-                        DocumentType.COVER_LETTER,
-                        st.session_state.cl_data,
-                    )
-                if err:
-                    st.error(f"❌ Error: {err}")
-                elif result and result.get("url"):
-                    st.session_state.cl_doc_url = result["url"]
-                    st.rerun()
-                else:
-                    st.error("❌ Failed to generate document (no URL returned).")
+                st.error("❌ Export failed (no URL returned).")
 
-    # ── Chat input (always visible) ──
-    st.divider()
-    placeholder_text = (
-        "Chat to revise: e.g. 'change destination to France' or 'update passport number to A9876543'"
-        if st.session_state.cl_data is not None
-        else "Describe your cover letter needs, or mention an Application ID to load data..."
-    )
-    chat_c1, chat_c2 = st.columns([5, 1])
-    with chat_c1:
-        chat_text = st.text_input(
-            "Chat", placeholder=placeholder_text,
-            key="cl_chat_input", label_visibility="collapsed",
+        # ── Chat ────────────────────────────────────────
+        st.divider()
+        st.caption(
+            "💬 Chat to revise the draft — "
+            "e.g. *'make the second paragraph more formal'* or *'update the passport number to A9876543'*"
         )
-    with chat_c2:
-        send_btn = st.button("Send ➤", type="primary", use_container_width=True, key="cl_send_btn")
-
-    if send_btn and chat_text:
-        st.session_state.cl_chat_messages.append({"role": "user", "content": chat_text})
-        st.session_state.cl_api_history.append({"role": "user", "content": chat_text})
-
-        active_app_id = st.session_state.cl_app_id or app_id_input or None
-
-        with st.spinner("🤖 Processing your request..."):
-            result, err = call_doc_chat_api(
-                chat_text,
-                st.session_state.cl_api_history,
-                active_app_id,
-                st.session_state.cl_session_uuid,
-                doc_data=st.session_state.cl_data,
+        chat_c1, chat_c2 = st.columns([5, 1])
+        with chat_c1:
+            chat_text = st.text_input(
+                "Chat",
+                placeholder="Describe your revision...",
+                key=f"{doc_type}_chat_input",
+                label_visibility="collapsed",
+            )
+        with chat_c2:
+            send_btn = st.button(
+                "Send ➤",
+                type="primary",
+                use_container_width=True,
+                key=f"{doc_type}_send_btn",
             )
 
-        if err:
-            st.session_state.cl_chat_messages.append({"role": "assistant", "content": f"❌ Error: {err}"})
-        elif result:
-            answer = result.get("answer", "No response received.")
-            st.session_state.cl_chat_messages.append({"role": "assistant", "content": answer})
-            st.session_state.cl_api_history.append({"role": "assistant", "content": answer})
-            if result.get("doc_url"):
-                st.session_state.cl_doc_url = result["doc_url"]
-            if result.get("doc_data"):
-                cl_push_history(st.session_state.cl_data)
-                _cl_clear_form_keys()
-                st.session_state.cl_data = result["doc_data"]
-
-        st.session_state.pop("cl_chat_input", None)
-        st.rerun()
-
-    elif send_btn:
-        st.warning("Please type a message first.")
+        if send_btn and chat_text:
+            st.session_state[sk("chat_messages")].append({"role": "user", "content": chat_text})
+            st.session_state[sk("api_history")].append({"role": "user", "content": chat_text})
+
+            with st.spinner("🤖 Revising document..."):
+                result, err = call_doc_chat_api(
+                    chat_text,
+                    st.session_state[sk("api_history")],
+                    st.session_state[sk("session_uuid")],
+                    current_document_content=st.session_state[sk("draft_content")],
+                    structure=st.session_state[sk("doc_structure")],
+                )
+
+            if err:
+                st.session_state[sk("chat_messages")].append(
+                    {"role": "assistant", "content": f"❌ Error: {err}"}
+                )
+            elif result:
+                answer = result.get("answer", "Document updated.")
+                st.session_state[sk("chat_messages")].append({"role": "assistant", "content": answer})
+                st.session_state[sk("api_history")].append({"role": "assistant", "content": answer})
+                if result.get("updated_document_content"):
+                    st.session_state[sk("draft_content")] = result["updated_document_content"]
+                if result.get("structure") is not None:
+                    st.session_state[sk("doc_structure")] = result["structure"]
+
+            st.session_state.pop(f"{doc_type}_chat_input", None)
+            st.rerun()
+
+        elif send_btn:
+            st.warning("Please type a message first.")
+
+        # ── Chat history ─────────────────────────────────
+        for msg in st.session_state[sk("chat_messages")]:
+            with st.chat_message(msg["role"]):
+                st.markdown(msg["content"])
+
+
+# ─────────────────────────────────────────────────────────
+#  PAGE ENTRY POINT
+# ─────────────────────────────────────────────────────────
+
+st.set_page_config(page_title="Document Generator", layout="wide")
+st.title("🌍 Document Generator")
+
+# Tabs are generated from DOCUMENT_REGISTRY — add a new DocumentSchema to get a new tab
+tab_labels = [f"{schema.icon} {schema.title}" for schema in DOCUMENT_REGISTRY.values()]
+tabs = st.tabs(tab_labels)
+
+for tab, (doc_type, schema) in zip(tabs, DOCUMENT_REGISTRY.items()):
+    with tab:
+        render_document_tab(schema)