import os
import sys
import tempfile
from io import BytesIO

import streamlit as st
import pandas as pd
from datetime import datetime

# Ensure project root is in python path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from src.extractor import PassportExtractor
from src.formats import format_fly_baghdad, format_fly_dubai, format_iraqi

# Page Config
st.set_page_config(page_title="Passport OCR Extractor", page_icon="🛂", layout="wide")

# Custom CSS for UI polish
st.markdown("""
    <style>
    /* Change cursor to pointer for selectbox */
    div[data-baseweb="select"] {
        cursor: pointer !important;
    }
    /* Hide the upload file list while processing if needed, 
       but Streamlit handles this naturally with spinners */
    </style>
    """, unsafe_allow_html=True)

# Initialize Extractor
@st.cache_resource
def get_extractor():
    return PassportExtractor(use_gpu=True)

EXTRACTOR = get_extractor()

MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024
ALLOWED_FILE_TYPES = ["png", "jpg", "jpeg", "pdf", "avif", "webp", "bmp", "tiff"]

def _safe_reason(problem):
    return (
        problem.get("reason")
        or problem.get("issue")
        or problem.get("message")
        or problem.get("error")
        or "<unspecified>"
    )

def _format_data(good_results, airline):
    if airline == "Fly Dubai":
        return format_fly_dubai(good_results)
    if airline == "Iraqi":
        return format_iraqi(good_results)
    if airline == "Fly Baghdad":
        return format_fly_baghdad(good_results)
    return pd.DataFrame(good_results)

def _placeholder_rows(airline, problematic_files, good_results):
    rows = []
    for problem in problematic_files:
        if airline == "Fly Dubai":
            rows.append({
                "TYPE": "Adult", "TITLE": "MR", "FIRST NAME": "•••", 
                "LAST NAME": "•••", "DOB (DDMMMYY)": "•••", "GENDER": "Male"
            })
        elif airline == "Iraqi":
            rows.append({
                "Last Name": "•••", "First Name and Middle Name": "•••", "Title": "MR",
                "PTC": "ADT", "Gender": "M", "Date of Birth": "•••",
                "Passport Last Name": "•••", "Passport First Name": "•••",
                "Passport Middle Name": "", "Passport Number": "•••",
                "Passport Nationality": "•••", "Passport Issue Country": "•••",
                "Passport Expiry Date": "•••", "Visa Number": "", "Visa Type": "",
                "Visa Issue Date": "", "Place of Birth": "", "Visa Place of Issue": "",
                "Visa Country of Application": "", "Address Type": "", "Address Country": "",
                "Address Details": "", "Address City": "", "Address State": "", "Address Zip Code": ""
            })
        elif airline == "Fly Baghdad":
            rows.append({
                "Sequence": len(good_results) + len(rows) + 1, "Traveling With": "",
                "Pax Type": "ADT", "Title": "MR", "First Name": "•••", "Last Name": "•••",
                "Gender": "MALE", "DOB (dd/mm/yyyy)": "•••", "Nationality": "•••",
                "Passport Number": "•••", "Passport Expiry (dd/mm/yyyy)": "•••",
                "Passport Issued Country": "•••"
            })
        else:
            rows.append({
                "source_file": problem["file_name"], "surname": "•••", "given_names": "•••",
                "passport_number": "•••", "nationality": "•••", "date_of_birth": "•••",
                "sex": "•••", "expiration_date": "•••", "personal_number": "•••", "mrz_found": False
            })
    return rows

def main():
    st.title("🛂 Passport OCR Extractor")
    st.markdown("Extract data from passport MRZ codes into structured tables.")

    with st.sidebar:
        st.header("Settings")
        st.info("💡 Optimized for performance - GPU enabled")
        
        airline = st.selectbox(
            "Choose Airline Format",
            options=["Default", "Fly Dubai", "Iraqi", "Fly Baghdad"],
            index=0
        )
        
        export_all_files = st.checkbox(
            "Include all files in export",
            value=False,
            help="When checked, includes files with missing/invalid data (shows ••• for missing fields)"
        )

    uploaded_files = st.file_uploader(
        "Upload Passport Files", 
        type=ALLOWED_FILE_TYPES, 
        accept_multiple_files=True
    )

    if st.button("Extract Data") and uploaded_files:
        all_results = []
        problematic_files = []
        
        # UI placeholder for results - cleared on each run
        status_container = st.empty()
        table_container = st.empty()
        download_container = st.empty()
        failed_container = st.expander("📋 Failed passport details", expanded=False)

        with st.status("Processing files...", expanded=True) as status:
            total_files = len(uploaded_files)
            
            for i, uploaded_file in enumerate(uploaded_files):
                file_name = uploaded_file.name
                ext = os.path.splitext(file_name)[1].lower()
                
                status.update(label=f"Processing {i+1}/{total_files}: {file_name}")

                if uploaded_file.size > MAX_FILE_SIZE_BYTES:
                    problematic_files.append({
                        "file_name": file_name,
                        "reason": f"File too large. Max 10 MB allowed."
                    })
                    continue

                # Save temp file for processing
                with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp_file:
                    tmp_file.write(uploaded_file.getvalue())
                    tmp_path = tmp_file.name

                file_processed = False
                results = []

                try:
                    if ext == ".pdf":
                        results = EXTRACTOR.process_pdf(tmp_path, airline=airline.lower()) or []
                        file_processed = True
                    else:
                        result = EXTRACTOR.get_data(tmp_path, airline=airline.lower())
                        if result:
                            results = [result]
                            file_processed = True
                except Exception as e:
                    problematic_files.append({"file_name": file_name, "reason": f"Processing error: {e}"})
                finally:
                    if os.path.exists(tmp_path):
                        os.remove(tmp_path)

                mrz_found_in_file = False
                for res in results:
                    res["source_file"] = file_name
                    if res.get("mrz_found", False):
                        mrz_found_in_file = True
                
                all_results.extend(results)

                if not file_processed:
                    problematic_files.append({"file_name": file_name, "reason": "Unsupported format or corrupted file"})
                elif not results:
                    problematic_files.append({"file_name": file_name, "reason": "No passport data detected"})
                elif not mrz_found_in_file:
                    problematic_files.append({"file_name": file_name, "reason": "MRZ data not found"})

            status.update(label="Extraction Complete!", state="complete", expanded=False)

        # Filtering logic
        good_results = []
        for result in all_results:
            is_problematic = any(p["file_name"] == result.get("source_file") for p in problematic_files)
            if not is_problematic and result.get("mrz_found", False):
                good_results.append(result)

        successful_extractions = len(good_results)
        failed_count = total_files - len(set(r.get("source_file", "") for r in good_results))

        if failed_count > 0:
            status_container.warning(f"⚠️ {successful_extractions} passport(s) were successfully extracted from {total_files} file(s).")
        else:
            status_container.success(f"✅ All {successful_extractions} passport(s) were successfully imported.")

        if problematic_files:
            with failed_container:
                for problem in problematic_files:
                    st.write(f"- **{problem['file_name']}**: {_safe_reason(problem)}")

        if not good_results and not (export_all_files and problematic_files):
            st.error("⚠️ No data could be extracted. Please check the files.")
            return

        # Prepare final DataFrame
        df = _format_data(good_results, airline)
        
        if export_all_files and problematic_files:
            placeholders = _placeholder_rows(airline, problematic_files, good_results)
            if placeholders:
                placeholder_df = pd.DataFrame(placeholders)
                df = placeholder_df if df.empty else pd.concat([df, placeholder_df], ignore_index=True)

        # Display results with pagination
        table_container.subheader(f"📊 Extracted Data ({len(df)} records)")
        # Streamlit's st.dataframe handles pagination automatically in recent versions
        table_container.dataframe(df, use_container_width=True, hide_index=True)

        # Download buttons
        col1, col2 = download_container.columns(2)
        
        csv = df.to_csv(index=False).encode('utf-8')
        col1.download_button(
            label="Download data as CSV",
            data=csv,
            file_name=f"passport_data_{airline.lower().replace(' ', '_')}.csv",
            mime='text/csv',
        )
        
        output = BytesIO()
        with pd.ExcelWriter(output, engine='openpyxl') as writer:
            df.to_excel(writer, index=False, sheet_name='PassportData')
        col2.download_button(
            label="Download data as Excel",
            data=output.getvalue(),
            file_name=f"passport_data_{airline.lower().replace(' ', '_')}.xlsx",
            mime='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
        )

if __name__ == "__main__":
    main()