import streamlit as st
import pandas as pd
import os
from typing import Dict, Tuple

class TrackingLookup:
    def __init__(self):
        self.tracking_data: Dict[str, Tuple[str, str]] = {}
        
    def extract_mailbag_no(self, df: pd.DataFrame) -> str:
        for _, row in df.iterrows():
            for col in df.columns:
                val = str(row[col])
                if 'HDOD-' in val:
                    return val.strip()
        return 'unknown_mailbag'

    def clean_tracking_numbers(self, df: pd.DataFrame) -> pd.DataFrame:
        tracking_col = None
        number_col = None
        
        if '#' in df.columns and 'Tracking No' in df.columns:
            return df[['#', 'Tracking No']]
        
        for col in df.columns:
            if df[col].astype(str).str.contains('UV|UW|LP|LORD', na=False, regex=True).any():
                tracking_col = col
                break
        
        for col in df.columns:
            if df[col].astype(str).str.match('^\d+$', na=False).any():
                number_col = col
                break
        
        if tracking_col is not None and number_col is not None:
            df = df[[number_col, tracking_col]]
            df.columns = ['#', 'Tracking No']
            df = df[df['Tracking No'].notna()]
            df = df[df['Tracking No'].str.match(r'^[A-Z0-9]+$', na=False)]
            return df
        
        return pd.DataFrame(columns=['#', 'Tracking No'])

    def process_excel_file(self, uploaded_file) -> bool:
        try:
            xl = pd.ExcelFile(uploaded_file)
            processed_sheets = 0
            
            for sheet_name in xl.sheet_names:
                try:
                    df = pd.read_excel(uploaded_file, sheet_name=sheet_name)
                    if df.empty or 'Domestic Letter Bill' not in str(df.iloc[0:5].values):
                        continue
                        
                    mailbag_no = self.extract_mailbag_no(df)
                    cleaned_df = self.clean_tracking_numbers(df)
                    
                    if not cleaned_df.empty:
                        for _, row in cleaned_df.iterrows():
                            self.tracking_data[row['Tracking No']] = (mailbag_no, row['#'])
                        processed_sheets += 1
                        
                except Exception as e:
                    st.warning(f"Warning: Could not process sheet '{sheet_name}': {str(e)}")
                    continue
            
            if processed_sheets > 0:
                st.success(f"✅ Successfully processed {processed_sheets} sheet(s)")
                return True
            else:
                st.error("No valid sheets were found in the file")
                return False
                
        except Exception as e:
            st.error(f"Error processing file: {str(e)}")
            return False

    def lookup_tracking(self, tracking_no: str) -> Tuple[str, str]:
        return self.tracking_data.get(tracking_no, (None, None))

def main():
    st.set_page_config(
        page_title="Mail Tracking Lookup",
        layout="wide",
        initial_sidebar_state="collapsed"
    )
    
    st.markdown("""
        <style>
        .stApp {
            background-color: #0e1117;
        }
        
        .result-box {
            background-color: #262730;
            padding: 2rem;
            border-radius: 1rem;
            margin: 2rem auto;
            border: 1px solid #4a4a4a;
            text-align: center;
            max-width: 800px;
        }
        
        .result-box p {
            color: #ffffff !important;
            margin: 1rem 0;
            font-size: 2.5rem !important;
            line-height: 1.2;
        }
        
        .result-box strong {
            color: #4CAF50;
            font-weight: 600;
        }
        
        .uploadedFile {
            background-color: #262730 !important;
            border: 1px solid #4a4a4a !important;
        }

        /* Hide the form submit button */
        .stButton {
            display: none;
        }
        </style>
        """, unsafe_allow_html=True)

    st.title("📬 Mail Tracking Lookup System")
    
    # Initialize session state
    if 'tracker' not in st.session_state:
        st.session_state.tracker = TrackingLookup()
    if 'file_processed' not in st.session_state:
        st.session_state.file_processed = False
    if 'last_result' not in st.session_state:
        st.session_state.last_result = None
    if 'form_count' not in st.session_state:
        st.session_state.form_count = 0

    # File upload section
    if not st.session_state.file_processed:
        uploaded_file = st.file_uploader("Choose an Excel file", type=['xlsx', 'xls'])
        if uploaded_file:
            if st.session_state.tracker.process_excel_file(uploaded_file):
                st.session_state.file_processed = True
                st.rerun()

    # Tracking lookup section
    if st.session_state.file_processed:
        st.markdown("<h2 style='text-align: center; margin-top: 2rem;'>Enter Tracking Number</h2>", unsafe_allow_html=True)
        
        # Centered input with columns
        col1, col2, col3 = st.columns([1, 2, 1])
        with col2:
            # Create a unique key for each form instance
            form_key = f'tracking_form_{st.session_state.form_count}'
            with st.form(key=form_key):
                tracking_input = st.text_input(
                    "Tracking Number",
                    label_visibility="collapsed",
                    placeholder="e.g., UV743594518UZ or LORD04012"
                )
                submit_button = st.form_submit_button(label='Search', type='primary')
            
            if submit_button and tracking_input:
                mailbag_no, sequence_no = st.session_state.tracker.lookup_tracking(tracking_input)
                
                if mailbag_no:
                    st.session_state.last_result = {
                        "tracking_no": tracking_input,
                        "mailbag_no": mailbag_no,
                        "sequence_no": sequence_no,
                        "found": True
                    }
                else:
                    st.session_state.last_result = {
                        "found": False,
                        "tracking": tracking_input
                    }
                
                # Increment form count to create a new form instance
                st.session_state.form_count += 1
                st.rerun()
                
            # Display results if they exist
            if st.session_state.last_result:
                if st.session_state.last_result["found"]:
                    st.markdown(
                        f"""<div class="result-box">
                            <p><strong>Tracking Number:</strong></p>
                            <p>{st.session_state.last_result["tracking_no"]}</p>
                            <p><strong>Mail Bag No:</strong></p>
                            <p>{st.session_state.last_result["mailbag_no"]}</p>
                            <p><strong>Sequence #:</strong></p>
                            <p>{st.session_state.last_result["sequence_no"]}</p>
                        </div>""",
                        unsafe_allow_html=True
                    )
                else:
                    st.warning(f"❌ Tracking number not found in the uploaded data")

if __name__ == "__main__":
    main()