# ─────────────────────────────────────────────────────────────
# Stage 1 UX Shell – FormPilot (no backend yet)
# ─────────────────────────────────────────────────────────────
import streamlit as st
import pandas as pd
from pathlib import Path
from io import BytesIO
import os
from dotenv import load_dotenv
load_dotenv()

from qdrant_client import QdrantClient
from langchain_community.vectorstores import Qdrant
from langchain_openai import OpenAIEmbeddings
from rag.qa_chain import get_answer
from rag.ocr_azure import parse_passport_azure


if "AZURE_DOC_KEY" not in os.environ:
    st.warning("⚠️  OCR disabled – set AZURE_DOC_KEY & AZURE_DOC_ENDPOINT")

st.set_page_config(
    page_title="FormPilot – Immigration Paralegal Copilot",
    page_icon="🛡️",
    layout="wide",
)


# Add custom CSS for better display of extracted information
st.markdown("""
<style>
    .extracted-info {
        background: #f7f7f7;
        padding: 12px;
        border-radius: 5px;
        border-left: 3px solid #4CAF50;
        margin-bottom: 15px;
    }
    .field-label {
        font-weight: bold;
        color: #333;
    }
    .field-value {
        color: #1E88E5;
    }
    .success-value {
        color: #4CAF50;
        font-weight: bold;
    }
    .missing-value {
        color: #F44336;
        font-style: italic;
    }
    .verification-checklist {
        background-color: #f9f9f9;
        padding: 15px;
        border-radius: 8px;
        margin-top: 20px;
    }
</style>
""", unsafe_allow_html=True)



# ------------------------------------------------------------
# Helper: generate checklist table with extracted information
# ------------------------------------------------------------
def build_checklist(profile=None):
    fields = [
        ("I‑485 Part 1 – Full Name", "pending"),
        ("I‑485 Part 1 – A‑Number", "pending"),
        ("I‑485 Part 1 – Date of Birth", "pending"),
        ("I‑485 Part 2 – Basis of Application", "pending"),
    ]
    df = pd.DataFrame(fields, columns=["Field", "Status"])
    
    # Update with extracted information if available
    if profile:
        # Full Name
        if profile.get('FirstName') and profile.get('LastName'):
            full_name = f"{profile.get('FirstName')} {profile.get('LastName')}"
            df.loc[df["Field"].str.contains("Full Name"), "Status"] = full_name
        
        # A-Number
        if profile.get('ANumber'):
            df.loc[df["Field"].str.contains("A‑Number"), "Status"] = profile.get('ANumber')
        
        # Date of Birth
        if profile.get('DateOfBirth'):
            df.loc[df["Field"].str.contains("Date of Birth"), "Status"] = profile.get('DateOfBirth')
    
    return df

# ------------------------------------------------------------
# Helper: Format extracted information for display
# ------------------------------------------------------------
def display_extracted_info(profile):
    if not profile:
        return st.info("No information extracted from documents.")
    
    st.markdown('<div class="verification-checklist">', unsafe_allow_html=True)
    
    # Display First Name
    st.markdown(
        f'<div class="extracted-info">'
        f'<span class="field-label">First Name:</span> '
        f'<span class="{"success-value" if profile.get("FirstName") else "missing-value"}">'
        f'{profile.get("FirstName", "Not found")}</span>'
        f'</div>', 
        unsafe_allow_html=True
    )
    
    # Display Last Name
    st.markdown(
        f'<div class="extracted-info">'
        f'<span class="field-label">Last Name:</span> '
        f'<span class="{"success-value" if profile.get("LastName") else "missing-value"}">'
        f'{profile.get("LastName", "Not found")}</span>'
        f'</div>', 
        unsafe_allow_html=True
    )
    
    # Display Date of Birth
    st.markdown(
        f'<div class="extracted-info">'
        f'<span class="field-label">Date of Birth:</span> '
        f'<span class="{"success-value" if profile.get("DateOfBirth") else "missing-value"}">'
        f'{profile.get("DateOfBirth", "Not found")}</span>'
        f'</div>', 
        unsafe_allow_html=True
    )
    
    # Display A-Number (if available)
    st.markdown(
        f'<div class="extracted-info">'
        f'<span class="field-label">A-Number:</span> '
        f'<span class="{"success-value" if profile.get("ANumber") else "missing-value"}">'
        f'{profile.get("ANumber", "Not applicable")}</span>'
        f'</div>', 
        unsafe_allow_html=True
    )
    
    st.markdown('</div>', unsafe_allow_html=True)



# ------------------------------------------------------------
# Sidebar – navigation + brand
# ------------------------------------------------------------
with st.sidebar:
    st.title("⚖️+📜+🛡️+🔍 FormPilot")
    if "stage" not in st.session_state:
        st.session_state.stage = "home"
    st.markdown("---")
    if st.button("🏠 Home"):
        st.session_state.stage = "home"
    if st.button("📄 Draft Packet"):
        if "uploaded_files" in st.session_state:
            st.session_state.stage = "draft"

# ------------------------------------------------------------
# Stage A – Upload page
# ------------------------------------------------------------
if st.session_state.stage == "home":
    st.header("New Case – Build I‑485 Package")
    uploaded_files = st.file_uploader(
        "Upload client documents (passport, visa, etc.)",
        type=["pdf", "jpg", "jpeg", "png"],
        accept_multiple_files=True,
    )
    form_choice = st.selectbox(
        "Select USCIS Form to prepare",
        options=["I‑485 (Adjustment of Status)"],
    )

    # Store uploads in session so we can view later
    if st.button("🚀 Build Package"):
        if not uploaded_files:
            st.warning("Please upload at least one document.")
        else:
            st.session_state.uploaded_files = uploaded_files
            st.session_state.form_choice = form_choice
            st.session_state.stage = "draft"
            st.rerun()

    st.info(
        """
        **What happens next?**  
        In Stage 2+ the AI will ingest your uploads, retrieve instructions,
        and draft the packet. For now, we jump to a placeholder Draft view.
        """
    )

# ------------------------------------------------------------
# Stage B – Draft Packet page
# ------------------------------------------------------------
elif st.session_state.stage == "draft":
    st.header("Draft Packet (Stage 3 – Retrieval MVP))")

    # Two-column layout: preview & checklist
    col_left, col_right = st.columns([2, 1])

    with col_left:
        st.subheader("📄 PDF Preview")
        st.write(
            "A preview of the filled I‑485 will appear here once AI pre‑fill is ready."
        )
        st.image(
            "https://placehold.co/600x800?text=PDF+Preview",
            caption="Static placeholder preview",
        )

    with col_right:
        st.subheader("✅ Document Information")
        profile = {}
        
        # Try to get profile from session state first
        if "profile" in st.session_state:
            profile = st.session_state.profile
        # Otherwise, try to process the first uploaded file
        elif "uploaded_files" in st.session_state and st.session_state.uploaded_files:
            with st.spinner("Extracting information..."):
                data = st.session_state.uploaded_files[0].getvalue()
                try:
                    profile = parse_passport_azure(data)    
                    st.session_state.profile = profile
                except Exception as e:
                    st.error(f"Error extracting information: {str(e)}")
        
        # Display nicely formatted extracted information
        if profile and (profile.get("FirstName") or profile.get("LastName") or profile.get("DateOfBirth")):
            display_extracted_info(profile)
            
            # Show checklist with auto-filled fields
            st.subheader("Form Field Checklist")
            df = build_checklist(profile)
            st.dataframe(df, hide_index=True, width=350)
            
            # Show the raw JSON for debugging
            with st.expander("Raw Extracted Data"):
                st.json(profile)
        else:
            st.warning("No profile information extracted from documents.")
            df = build_checklist()
            st.dataframe(df, hide_index=True, width=350)


    st.markdown("---")
    st.subheader("Ask about I‑485 instructions")
    q = st.text_input("Your question", key="qa")
    if q:
        with st.spinner("Retrieving..."):
            ans, cites = get_answer(q)
        st.success(ans)
        st.caption("Sources: " + ", ".join(sorted(cites)))


# ------------------------------------------------------------
# (Optional) Future stage for OCR parsing
# ------------------------------------------------------------
# Example usage of doc_parse_stub
if False:
    for file in st.session_state.get("uploaded_files", []):
        data = file.read()
        parsed = parse_passport(data)
        st.write(parsed)