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]] = {}  # {tracking_no: (mailbag_no, #)}
        
    def extract_mailbag_no(self, df: pd.DataFrame) -> str:
        """Extract Mail Bag No from the DataFrame"""
        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:
        """Extract only # and Tracking No columns, removing any empty rows"""
        tracking_col = None
        number_col = None
        
        # First try to find by exact column names
        if '#' in df.columns and 'Tracking No' in df.columns:
            return df[['#', 'Tracking No']]
        
        # Try to find tracking number column
        for col in df.columns:
            # Check if column contains any tracking number patterns
            if df[col].astype(str).str.contains('UV|UW|LP|LORD', na=False, regex=True).any():
                tracking_col = col
                break
        
        # Find the number column (usually the first column with 1,2,3...)
        for col in df.columns:
            if df[col].astype(str).str.match('^\d+$', na=False).any():
                number_col = col
                break
        
        # If we found both columns
        if tracking_col is not None and number_col is not None:
            df = df[[number_col, tracking_col]]
            df.columns = ['#', 'Tracking No']
            
            # Clean data
            df = df[df['Tracking No'].notna()]
            # Updated pattern to include LORD format
            df = df[df['Tracking No'].str.match(r'^[A-Z0-9]+$', na=False)]
            return df
        
        # If we couldn't find the required columns, return empty DataFrame
        return pd.DataFrame(columns=['#', 'Tracking No'])

    def process_excel_file(self, uploaded_file) -> bool:
        """Process uploaded Excel file and store tracking data"""
        try:
            # Read all sheets
            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)
                    
                    # Skip empty sheets or sheets without the expected structure
                    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:
                        # Store tracking numbers with their mailbag and sequence numbers
                        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.info(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]:
        """Look up mailbag number and sequence number for a tracking number"""
        return self.tracking_data.get(tracking_no, (None, None))

def main():
    # Configure the page
    st.set_page_config(
        page_title="Mail Tracking Lookup",
        layout="wide",
        initial_sidebar_state="collapsed"
    )
    
    # Custom CSS for dark theme
    st.markdown("""
        <style>
        /* Main app */
        .stApp {
            background-color: #0e1117;
        }
        
        /* Headers */
        h1, h2, h3, h4, h5, h6 {
            color: #ffffff !important;
        }
        
        /* Result box */
        .result-box {
            background-color: #262730;
            padding: 1.5rem;
            border-radius: 0.5rem;
            margin: 1rem 0;
            border: 1px solid #4a4a4a;
        }
        
        .result-box p {
            color: #ffffff !important;
            margin: 0.5rem 0;
        }
        
        .result-box .big-font {
            font-size: 1.25rem;
            font-weight: 600;
            margin-bottom: 1rem;
            color: #ffffff;
        }
        
        .result-box strong {
            color: #4CAF50;
            font-weight: 600;
        }
        
        /* File uploader */
        .uploadedFile {
            background-color: #262730 !important;
            border: 1px solid #4a4a4a !important;
        }
        
        /* Button */
        .stButton button {
            background-color: #4CAF50 !important;
            color: white !important;
            border: none !important;
        }
        
        .stButton button:hover {
            background-color: #45a049 !important;
        }
        </style>
        """, unsafe_allow_html=True)

    st.title("📬 Mail Tracking Lookup System")
    
    # Initialize tracker in session state if it doesn't exist
    if 'tracker' not in st.session_state:
        st.session_state['tracker'] = TrackingLookup()
    
    # File upload section
    st.header("1. Upload Excel File")
    uploaded_file = st.file_uploader("Choose an Excel file", type=['xlsx', 'xls'])
    
    if uploaded_file:
        col1, col2, col3 = st.columns([1, 2, 1])
        with col2:
            if st.button("Process File", use_container_width=True):
                success = st.session_state.tracker.process_excel_file(uploaded_file)
                if success:
                    st.success("✅ File processed successfully!")
                    st.session_state.file_processed = True
    
    # Tracking number lookup section
    st.header("2. Lookup Tracking Number")
    
    tracking_input = st.text_input(
        "Enter tracking number:",
        placeholder="e.g., UV743594518UZ or LORD04012",
        help="Enter the tracking number and press Enter"
    )
    
    if tracking_input:
        mailbag_no, sequence_no = st.session_state.tracker.lookup_tracking(tracking_input)
        
        if mailbag_no:
            st.markdown(
                f"""<div class="result-box">
                    <p class="big-font">📦 Tracking Details</p>
                    <p><strong>Mail Bag No:</strong> {mailbag_no}</p>
                    <p><strong>Sequence #:</strong> {sequence_no}</p>
                </div>""",
                unsafe_allow_html=True
            )
        else:
            st.warning("❌ Tracking number not found in the uploaded data")

if __name__ == "__main__":
    main()