import os
import sys
from dotenv import load_dotenv
import gradio as gr
from typing import Optional, Dict, List, Union
import logging

# Custom CSS
CUSTOM_CSS = """
.footer {
    text-align: center !important;
    padding: 20px !important;
    margin-top: 40px !important;
    border-top: 1px solid #404040 !important;
    color: #89CFF0 !important;
    font-size: 1.1em !important;
}

.gradio-container {
    max-width: 1200px !important;
    margin: auto !important;
    padding: 20px !important;
    background-color: #1a1a1a !important;
    color: #ffffff !important;
}

.main-header {
    background: linear-gradient(135deg, #1e3c72 0%, #2a5298 100%) !important;
    color: white !important;
    padding: 30px !important;
    border-radius: 15px !important;
    margin-bottom: 30px !important;
    text-align: center !important;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2) !important;
}

.app-title {
    font-size: 2.5em !important;
    font-weight: bold !important;
    margin-bottom: 10px !important;
    background: linear-gradient(90deg, #ffffff, #3498DB) !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.3) !important;
}

.app-subtitle {
    font-size: 1.3em !important;
    color: #89CFF0 !important;
    margin-bottom: 15px !important;
    font-weight: 500 !important;
}

.app-description {
    font-size: 1.1em !important;
    color: #B0C4DE !important;
    margin-bottom: 20px !important;
    line-height: 1.5 !important;
}

.gr-checkbox-group {
    background: #363636 !important;
    padding: 15px !important;
    border-radius: 10px !important;
    margin: 10px 0 !important;
}

.gr-slider {
    margin-top: 10px !important;
}

.status-message {
    margin-top: 10px !important;
    padding: 8px !important;
    border-radius: 4px !important;
    background-color: #2d2d2d !important;
}

.result-box {
    background: #363636 !important;
    border: 1px solid #404040 !important;
    border-radius: 10px !important;
    padding: 20px !important;
    margin-top: 15px !important;
    color: #ffffff !important;
}

.chart-container {
    background: #2d2d2d !important;
    padding: 20px !important;
    border-radius: 10px !important;
    box-shadow: 0 2px 4px rgba(0,0,0,0.2) !important;
    color: #ffffff !important;
}

.action-button {
    background: #3498DB !important;
    color: white !important;
    border: none !important;
    padding: 10px 20px !important;
    border-radius: 5px !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
}

.action-button:hover {
    background: #2980B9 !important;
    transform: translateY(-2px) !important;
    box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important;
}
"""

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Load environment variables
load_dotenv()

# Constants
MAX_FILE_SIZE = 50 * 1024 * 1024  # 50MB
ALLOWED_EXTENSIONS = {'.xlsx', '.xls', '.csv'}
import pandas as pd
import google.generativeai as genai
import joblib
from reportlab.lib import colors
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, Image
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
import plotly.express as px
import plotly.graph_objects as go
import tempfile
from datetime import datetime
import numpy as np
from xgboost import XGBRegressor

# Configure Gemini API
GEMINI_API_KEY = os.getenv("gemini_api")

genai.configure(api_key=GEMINI_API_KEY)
generation_config = {
    "temperature": 1,
    "top_p": 0.95,
    "top_k": 64,
    "max_output_tokens": 8192,
    "response_mime_type": "text/plain",
}

model = genai.GenerativeModel(
    model_name="gemini-2.0-pro-exp-02-05",
    generation_config=generation_config,
)

chat_model = genai.GenerativeModel('"gemini-2.0-pro-exp-02-05"')

class SupplyChainState:
    def __init__(self):
        self.sales_df = None
        self.supplier_df = None
        self.text_data = None
        self.chat_history = []
        self.analysis_results = {}
        self.freight_predictions = []
        
        try:
            self.freight_model = create_initial_model()
        except Exception as e:
            print(f"Warning: Could not create freight prediction model: {e}")
            self.freight_model = None

def create_initial_model():
    n_samples = 1000
    np.random.seed(42)
    
    data = {
        'weight (kilograms)': np.random.uniform(100, 10000, n_samples),
        'line item value': np.random.uniform(1000, 1000000, n_samples),
        'cost per kilogram': np.random.uniform(1, 500, n_samples),
        'shipment mode_Air Charter_weight': np.zeros(n_samples),
        'shipment mode_Ocean_weight': np.zeros(n_samples),
        'shipment mode_Truck_weight': np.zeros(n_samples),
        'shipment mode_Air Charter_line_item_value': np.zeros(n_samples),
        'shipment mode_Ocean_line_item_value': np.zeros(n_samples),
        'shipment mode_Truck_line_item_value': np.zeros(n_samples)
    }
    
    modes = ['Air', 'Ocean', 'Truck']
    for i in range(n_samples):
        mode = np.random.choice(modes)
        if mode == 'Air':
            data['shipment mode_Air Charter_weight'][i] = data['weight (kilograms)'][i]
            data['shipment mode_Air Charter_line_item_value'][i] = data['line item value'][i]
        elif mode == 'Ocean':
            data['shipment mode_Ocean_weight'][i] = data['weight (kilograms)'][i]
            data['shipment mode_Ocean_line_item_value'][i] = data['line item value'][i]
        else:
            data['shipment mode_Truck_weight'][i] = data['weight (kilograms)'][i]
            data['shipment mode_Truck_line_item_value'][i] = data['line item value'][i]
    
    df = pd.DataFrame(data)
    base_cost = (df['weight (kilograms)'] * df['cost per kilogram'] * 0.8 + 
                 df['line item value'] * 0.02)
    
    air_multiplier = 1.5
    ocean_multiplier = 0.8
    truck_multiplier = 1.0

    freight_cost = (
        base_cost * (air_multiplier * (df['shipment mode_Air Charter_weight'] > 0) +
                    ocean_multiplier * (df['shipment mode_Ocean_weight'] > 0) +
                    truck_multiplier * (df['shipment mode_Truck_weight'] > 0))
    )
    
    freight_cost = freight_cost + np.random.normal(0, freight_cost * 0.1)
    
    model = XGBRegressor(n_estimators=100, learning_rate=0.1, max_depth=5, random_state=42)
    model.fit(df, freight_cost)
    
    return model

def process_uploaded_data(state, sales_file, supplier_file, text_data):
    try:
        if sales_file is not None:
            file_ext = os.path.splitext(sales_file.name)[1].lower()
            if file_ext not in ['.xlsx', '.xls', '.csv']:
                return '❌ Error: Sales data must be in Excel (.xlsx, .xls) or CSV format'
            
            try:
                if file_ext == '.csv':
                    state.sales_df = pd.read_csv(sales_file.name)
                else:
                    state.sales_df = pd.read_excel(sales_file.name)
            except Exception as e:
                return f'❌ Error reading sales data: {str(e)}'
            
        if supplier_file is not None:
            file_ext = os.path.splitext(supplier_file.name)[1].lower()
            if file_ext not in ['.xlsx', '.xls', '.csv']:
                return '❌ Error: Supplier data must be in Excel (.xlsx, .xls) or CSV format'
            
            try:
                if file_ext == '.csv':
                    state.supplier_df = pd.read_csv(supplier_file.name)
                else:
                    state.supplier_df = pd.read_excel(supplier_file.name)
            except Exception as e:
                return f'❌ Error reading supplier data: {str(e)}'
            
        state.text_data = text_data
        return "✅ Data processed successfully"
    except Exception as e:
        return f'❌ Error processing data: {str(e)}'

def perform_demand_forecasting(state):
    if state.sales_df is None:
        return "Error: No sales data provided", None, "Please upload sales data first"
    
    try:
        sales_summary = state.sales_df.describe().to_string()
        prompt = f"""Analyze the following sales data summary and provide:
        1. A detailed demand forecast for the next quarter
        2. Key trends and seasonality patterns
        3. Actionable recommendations

        Data Summary:
        {sales_summary}

        Please structure your response with clear sections for Forecast, Trends, and Recommendations."""

        response = model.generate_content(prompt)
        analysis_text = response.text
        
        fig = px.line(state.sales_df, title='Historical Sales Data and Forecast')
        fig.update_layout(
            template='plotly_dark',
            title_x=0.5,
            title_font_size=20,
            showlegend=True,
            hovermode='x',
            paper_bgcolor='#2d2d2d',
            plot_bgcolor='#363636',
            font=dict(color='white')
        )
        
        return analysis_text, fig, "✅ Analysis completed successfully"
    except Exception as e:
        return f"❌ Error in demand forecasting: {str(e)}", None, "Analysis failed"

def perform_risk_assessment(state):
    if state.supplier_df is None:
        return "Error: No supplier data provided", None, "Please upload supplier data first"
    
    try:
        supplier_summary = state.supplier_df.describe().to_string()
        prompt = f"""Perform a comprehensive risk assessment based on:
        
        Supplier Data Summary:
        {supplier_summary}
        
        Additional Context:
        {state.text_data if state.text_data else 'No additional context provided'}
        
        Please provide:
        1. Risk scoring for each supplier
        2. Identified risk factors
        3. Mitigation recommendations"""

        response = model.generate_content(prompt)
        analysis_text = response.text

        fig = px.scatter(state.supplier_df, title='Supplier Risk Assessment')
        fig.update_layout(
            template='plotly_dark',
            title_x=0.5,
            title_font_size=20,
            showlegend=True,
            hovermode='closest',
            paper_bgcolor='#2d2d2d',
            plot_bgcolor='#363636',
            font=dict(color='white')
        )
        
        return analysis_text, fig, "✅ Risk assessment completed"
    except Exception as e:
        return f"❌ Error in risk assessment: {str(e)}", None, "Assessment failed"

def perform_inventory_optimization(state):
    if state.sales_df is None:
        return "Error: No sales data provided", None, "Please upload sales data first"
    
    try:
        inventory_summary = state.sales_df.describe().to_string()
        prompt = f"""Analyze the following inventory data and provide:
        1. Optimal inventory levels
        2. Reorder points
        3. Safety stock recommendations
        4. ABC analysis insights

        Data Summary:
        {inventory_summary}

        Additional Context:
        {state.text_data if state.text_data else 'No additional context provided'}

        Please structure your response with clear sections for each aspect."""

        response = model.generate_content(prompt)
        analysis_text = response.text

        fig = go.Figure()
        
        if 'quantity' in state.sales_df.columns:
            fig.add_trace(go.Scatter(
                y=state.sales_df['quantity'],
                name='Inventory Level',
                line=dict(color='#3498DB')
            ))

        fig.update_layout(
            title='Inventory Level Analysis',
            template='plotly_dark',
            title_x=0.5,
            title_font_size=20,
            showlegend=True,
            hovermode='x',
            paper_bgcolor='#2d2d2d',
            plot_bgcolor='#363636',
            font=dict(color='white')
        )
        
        return analysis_text, fig, "✅ Inventory optimization completed"
    except Exception as e:
        return f"❌ Error in inventory optimization: {str(e)}", None, "Analysis failed"

def perform_supplier_performance(state):
    if state.supplier_df is None:
        return "Error: No supplier data provided", None, "Please upload supplier data first"
    
    try:
        supplier_summary = state.supplier_df.describe().to_string()
        prompt = f"""Analyze supplier performance based on:
        
        Supplier Data Summary:
        {supplier_summary}
        
        Additional Context:
        {state.text_data if state.text_data else 'No additional context provided'}
        
        Please provide:
        1. Supplier performance metrics
        2. Performance rankings
        3. Areas for improvement
        4. Supplier development recommendations"""

        response = model.generate_content(prompt)
        analysis_text = response.text

        if 'performance_score' in state.supplier_df.columns:
            fig = px.box(state.supplier_df, y='performance_score', 
                        title='Supplier Performance Distribution')
        else:
            fig = go.Figure(data=[
                go.Bar(name='On-Time Delivery', x=['Supplier A', 'Supplier B', 'Supplier C'],
                      y=[95, 87, 92]),
                go.Bar(name='Quality Score', x=['Supplier A', 'Supplier B', 'Supplier C'],
                      y=[88, 94, 90])
            ])

        fig.update_layout(
            template='plotly_dark',
            title_x=0.5,
            title_font_size=20,
            showlegend=True,
            paper_bgcolor='#2d2d2d',
            plot_bgcolor='#363636',
            font=dict(color='white')
        )
        
        return analysis_text, fig, "✅ Supplier performance analysis completed"
    except Exception as e:
        return f"❌ Error in supplier performance analysis: {str(e)}", None, "Analysis failed"

def perform_sustainability_analysis(state):
    if state.supplier_df is None and state.sales_df is None:
        return "Error: No data provided", None, "Please upload data first"
    
    try:
        data_summary = ""
        if state.supplier_df is not None:
            data_summary += f"Supplier Data Summary:\n{state.supplier_df.describe().to_string()}\n\n"
        if state.sales_df is not None:
            data_summary += f"Sales Data Summary:\n{state.sales_df.describe().to_string()}"

        prompt = f"""Perform a comprehensive sustainability analysis:
        
        Data Summary:
        {data_summary}
        
        Additional Context:
        {state.text_data if state.text_data else 'No additional context provided'}
        
        Please provide:
        1. Carbon footprint analysis
        2. Environmental impact metrics
        3. Sustainability recommendations
        4. Green initiative opportunities
        5. ESG performance indicators"""

        response = model.generate_content(prompt)
        analysis_text = response.text

        fig = go.Figure()
        
        categories = ['Carbon Emissions', 'Water Usage', 'Waste Reduction', 
                     'Energy Efficiency', 'Green Transportation']
        current_scores = [75, 82, 68, 90, 60]
        target_scores = [100, 100, 100, 100, 100]

        fig.add_trace(go.Scatterpolar(
            r=current_scores,
            theta=categories,
            fill='toself',
            name='Current Performance'
        ))
        fig.add_trace(go.Scatterpolar(
            r=target_scores,
            theta=categories,
            fill='toself',
            name='Target'
        ))

        fig.update_layout(
            polar=dict(
                radialaxis=dict(
                    visible=True,
                    range=[0, 100]
                )),
            showlegend=True,
            title='Sustainability Performance Metrics',
            template='plotly_dark',
            title_x=0.5,
            title_font_size=20,
            paper_bgcolor='#2d2d2d',
            plot_bgcolor='#363636',
            font=dict(color='white')
        )
        
        return analysis_text, fig, "✅ Sustainability analysis completed"
    except Exception as e:
        return f"❌ Error in sustainability analysis: {str(e)}", None, "Analysis failed"

def calculate_shipping_cost(base_cost, params):
    """Calculate total shipping cost with all factors"""
    total_cost = base_cost
    
    # Fuel surcharge
    fuel_charge = base_cost * (params['fuel_surcharge'] / 100)
    
    # Insurance
    insurance = params['line_item_value'] * (params['insurance_rate'] / 100)
    
    # Customs duty
    duty = params['line_item_value'] * (params['customs_duty'] / 100)
    
    # Special handling charges
    handling_charges = 0
    handling_rates = {
        "Temperature Controlled": 0.15,
        "Hazardous Materials": 0.25,
        "Fragile Items": 0.10,
        "Express Delivery": 0.20,
        "Door-to-Door Service": 0.15
    }
    
    for requirement in params['special_handling']:
        if requirement in handling_rates:
            handling_charges += base_cost * handling_rates[requirement]
    
    # Distance-based charge
    distance_rate = {
        "Air": 0.1,
        "Ocean": 0.05,
        "Truck": 0.15
    }
    distance_charge = params['distance'] * distance_rate[params['shipment_mode']]
    
    # Time-based charge
    transit_charge = params['transit_time'] * (base_cost * 0.01)
    
    total_cost = base_cost + fuel_charge + insurance + duty + handling_charges + distance_charge + transit_charge
    
    return {
        'base_cost': round(base_cost, 2),
        'fuel_charge': round(fuel_charge, 2),
        'insurance': round(insurance, 2),
        'customs_duty': round(duty, 2),
        'handling_charges': round(handling_charges, 2),
        'distance_charge': round(distance_charge, 2),
        'transit_charge': round(transit_charge, 2),
        'total_cost': round(total_cost, 2)
    }

def predict_freight_cost(state, params):
    """Predict freight cost with enhanced parameters"""
    if state.freight_model is None:
        return "Error: Freight prediction model not loaded"
        
    try:
        # Clean shipment mode string
        mode = params['shipment_mode'].replace("✈️ ", "").replace("🚢 ", "").replace("🚛 ", "")
        
        # Prepare features for the model
        features = {
            'weight (kilograms)': params['weight'],
            'line item value': params['line_item_value'],
            'cost per kilogram': params['cost_per_kg'],
            'shipment mode_Air Charter_weight': params['weight'] if mode == "Air" else 0,
            'shipment mode_Ocean_weight': params['weight'] if mode == "Ocean" else 0,
            'shipment mode_Truck_weight': params['weight'] if mode == "Truck" else 0,
            'shipment mode_Air Charter_line_item_value': params['line_item_value'] if mode == "Air" else 0,
            'shipment mode_Ocean_line_item_value': params['line_item_value'] if mode == "Ocean" else 0,
            'shipment mode_Truck_line_item_value': params['line_item_value'] if mode == "Truck" else 0
        }
        
        input_data = pd.DataFrame([features])
        base_prediction = state.freight_model.predict(input_data)[0]
        
        # Calculate total cost with all factors
        cost_breakdown = calculate_shipping_cost(base_prediction, params)
        
        return cost_breakdown
        
    except Exception as e:
        return f"Error making prediction: {str(e)}"
    if state.freight_model is None:
        return "Error: Freight prediction model not loaded"
        
    try:
        # Set weights based on mode
        if "Air" in shipment_mode:
            air_charter_weight = weight
            air_charter_value = line_item_value
        elif "Ocean" in shipment_mode:
            ocean_weight = weight
            ocean_value = line_item_value
        else:
            truck_weight = weight
            truck_value = line_item_value

        features = {
            'weight (kilograms)': weight,
            'line item value': line_item_value,
            'cost per kilogram': cost_per_kg,
            'shipment mode_Air Charter_weight': air_charter_weight,
            'shipment mode_Ocean_weight': ocean_weight,
            'shipment mode_Truck_weight': truck_weight,
            'shipment mode_Air Charter_line_item_value': air_charter_value,
            'shipment mode_Ocean_line_item_value': ocean_value,
            'shipment mode_Truck_line_item_value': truck_value
        }
        input_data = pd.DataFrame([features])
        
        prediction = state.freight_model.predict(input_data)
        return round(float(prediction[0]), 2)
    except Exception as e:
        return f"Error making prediction: {str(e)}"
    if state.freight_model is None:
        return "Error: Freight prediction model not loaded"
        
    try:
        features = {
            'weight (kilograms)': weight,
            'line item value': line_item_value,
            'cost per kilogram': cost_per_kg,
            'shipment mode_Air Charter_weight': air_charter_weight if "Air" in shipment_mode else 0,
            'shipment mode_Ocean_weight': ocean_weight if "Ocean" in shipment_mode else 0,
            'shipment mode_Truck_weight': truck_weight if "Truck" in shipment_mode else 0,
            'shipment mode_Air Charter_line_item_value': air_charter_value if "Air" in shipment_mode else 0,
            'shipment mode_Ocean_line_item_value': ocean_value if "Ocean" in shipment_mode else 0,
            'shipment mode_Truck_line_item_value': truck_value if "Truck" in shipment_mode else 0
        }
        input_data = pd.DataFrame([features])
        
        prediction = state.freight_model.predict(input_data)
        return round(float(prediction[0]), 2)
    except Exception as e:
        return f"Error making prediction: {str(e)}"

def chat_with_navigator(state, message):
    try:
        context = "Available data and analysis:\n"
        if state.sales_df is not None:
            context += f"- Sales data with {len(state.sales_df)} records\n"
        if state.supplier_df is not None:
            context += f"- Supplier data with {len(state.supplier_df)} records\n"
        if state.text_data:
            context += "- Additional context from text data\n"
        if state.freight_predictions:
            context += f"- Recent freight predictions: {state.freight_predictions[-5:]}\n"
        
        if state.analysis_results:
            context += "\nRecent analysis results:\n"
            for analysis_type, results in state.analysis_results.items():
                context += f"- {analysis_type} completed\n"
        
        prompt = f"""You are SupplyChainAI Navigator's assistant. Help the user with supply chain analysis, 
        including demand forecasting, risk assessment, and freight cost predictions.
        
        Available Context:
        {context}
        
        Chat History:
        {str(state.chat_history[-3:]) if state.chat_history else 'No previous messages'}
        
        User message: {message}
        
        Provide a helpful response based on the available data and analysis results."""

        response = chat_model.generate_content(prompt)
        
        state.chat_history.append({"role": "user", "content": message})
        state.chat_history.append({"role": "assistant", "content": response.text})
        
        return state.chat_history
    except Exception as e:
        return [{"role": "assistant", "content": f"Error: {str(e)}"}]

def generate_pdf_report(state, analysis_options):
    try:
        temp_dir = tempfile.mkdtemp()
        pdf_path = os.path.join(temp_dir, "supply_chain_report.pdf")
        
        doc = SimpleDocTemplate(pdf_path, pagesize=letter)
        styles = getSampleStyleSheet()
        story = []
        
        # Create custom title style
        title_style = ParagraphStyle(
            'CustomTitle',
            parent=styles['Heading1'],
            fontSize=24,
            spaceAfter=30,
            textColor=colors.HexColor('#2C3E50')
        )
        
        story.append(Paragraph("SupplyChainAI Navigator Report", title_style))
        story.append(Spacer(1, 12))
        
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        story.append(Paragraph(f"Generated on: {timestamp}", styles['Normal']))
        story.append(Spacer(1, 20))
        
        if state.analysis_results:
            for analysis_type, results in state.analysis_results.items():
                if analysis_type in analysis_options:
                    story.append(Paragraph(analysis_type, styles['Heading2']))
                    story.append(Spacer(1, 12))
                    story.append(Paragraph(results['text'], styles['Normal']))
                    story.append(Spacer(1, 12))
                    
                    if 'figure' in results:
                        img_path = os.path.join(temp_dir, f"{analysis_type.lower()}_plot.png")
                        results['figure'].write_image(img_path)
                        story.append(Image(img_path, width=400, height=300))
                    
                    story.append(Spacer(1, 20))
        
        if state.freight_predictions:
            story.append(Paragraph("Recent Freight Cost Predictions", styles['Heading2']))
            story.append(Spacer(1, 12))
            
            pred_data = [["Prediction #", "Cost (USD)"]]
            for i, pred in enumerate(state.freight_predictions[-5:], 1):
                pred_data.append([f"Prediction {i}", f"${pred:,.2f}"])
            
            table = Table(pred_data)
            table.setStyle(TableStyle([
                ('BACKGROUND', (0, 0), (-1, 0), colors.HexColor('#3498DB')),
                ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
                ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
                ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
                ('FONTSIZE', (0, 0), (-1, 0), 14),
                ('BOTTOMPADDING', (0, 0), (-1, 0), 12),
                ('BACKGROUND', (0, 1), (-1, -1), colors.whitesmoke),
                ('TEXTCOLOR', (0, 1), (-1, -1), colors.black),
                ('FONTNAME', (0, 1), (-1, -1), 'Helvetica'),
                ('FONTSIZE', (0, 1), (-1, -1), 12),
                ('GRID', (0, 0), (-1, -1), 1, colors.black)
            ]))
            story.append(table)
            story.append(Spacer(1, 20))
        
        doc.build(story)
        return pdf_path
    except Exception as e:
        print(f"Error generating PDF: {str(e)}")
        return None

def run_analyses(state, choices, sales_file, supplier_file, text_data):
    results = []
    figures = []
    status_messages = []

    process_status = process_uploaded_data(state, sales_file, supplier_file, text_data)
    if "Error" in process_status:
        return process_status, None, process_status

    for choice in choices:
        if "📈 Demand Forecasting" in choice:
            text, fig, status = perform_demand_forecasting(state)
            results.append(text)
            figures.append(fig)
            status_messages.append(status)
            if text and fig:
                state.analysis_results['Demand Forecasting'] = {'text': text, 'figure': fig}
        
        elif "⚠️ Risk Assessment" in choice:
            text, fig, status = perform_risk_assessment(state)
            results.append(text)
            figures.append(fig)
            status_messages.append(status)
            if text and fig:
                state.analysis_results['Risk Assessment'] = {'text': text, 'figure': fig}
        
        elif "📦 Inventory Optimization" in choice:
            text, fig, status = perform_inventory_optimization(state)
            results.append(text)
            figures.append(fig)
            status_messages.append(status)
            if text and fig:
                state.analysis_results['Inventory Optimization'] = {'text': text, 'figure': fig}
        
        elif "🤝 Supplier Performance" in choice:
            text, fig, status = perform_supplier_performance(state)
            results.append(text)
            figures.append(fig)
            status_messages.append(status)
            if text and fig:
                state.analysis_results['Supplier Performance'] = {'text': text, 'figure': fig}
        
        elif "🌿 Sustainability Analysis" in choice:
            text, fig, status = perform_sustainability_analysis(state)
            results.append(text)
            figures.append(fig)
            status_messages.append(status)
            if text and fig:
                state.analysis_results['Sustainability Analysis'] = {'text': text, 'figure': fig}

    combined_results = "\n\n".join(results)
    combined_status = "\n".join(status_messages)
    
    final_figure = figures[-1] if figures else None
    
    return combined_results, final_figure, combined_status

def predict_and_store_freight(state, *args):
    if len(args) >= 3:
        weight, line_item_value, shipment_mode = args[:3]
        result = predict_freight_cost(state, weight, line_item_value, 50, shipment_mode)
        if isinstance(result, (int, float)):
            state.freight_predictions.append(result)
        return result
    return "Error: Invalid parameters"

def create_interface():
    state = SupplyChainState()
    
    with gr.Blocks(css=CUSTOM_CSS, title="SupplyChainAI Navigator") as demo:
        # Header
        with gr.Row(elem_classes="main-header"):
            with gr.Column():
                gr.Markdown("# 🚢 SupplyChainAI Navigator", elem_classes="app-title")
                gr.Markdown("### Intelligent Supply Chain Analysis & Optimization", elem_classes="app-subtitle")
                gr.Markdown("An AI-powered platform for comprehensive supply chain analytics", elem_classes="app-description")
                gr.Markdown("### Created by Aditya Ratan", elem_classes="creator-info")
        
        # Main Content Tabs
        with gr.Tabs() as tabs:
            # Data Upload Tab
            with gr.Tab("📊 Data Upload", elem_classes="tab-content"):
                with gr.Row():
                    with gr.Column(scale=1):
                        sales_data_upload = gr.File(
                            file_types=[".xlsx", ".xls", ".csv"],
                            label="📈 Sales Data (Excel or CSV)",
                            elem_classes="file-upload"
                        )
                        gr.Markdown("*Upload sales data in Excel (.xlsx, .xls) or CSV format*", elem_classes="file-instructions")
                    
                    with gr.Column(scale=1):
                        supplier_data_upload = gr.File(
                            file_types=[".xlsx", ".xls", ".csv"],
                            label="🏭 Supplier Data (Excel or CSV)",
                            elem_classes="file-upload"
                        )
                        gr.Markdown("*Upload supplier data in Excel (.xlsx, .xls) or CSV format*", elem_classes="file-instructions")
                
                with gr.Row():
                    text_input_area = gr.Textbox(
                        label="📝 Additional Context",
                        placeholder="Add market updates, news, or other relevant information...",
                        lines=5
                    )
                
                with gr.Row():
                    upload_status = gr.Textbox(
                        label="Status",
                        elem_classes="status-box"
                    )
                    upload_button = gr.Button(
                        "🔄 Process Data",
                        variant="primary",
                        elem_classes="action-button"
                    )

            # Analysis Tab
            with gr.Tab("🔍 Analysis", elem_classes="tab-content"):
                with gr.Row():
                    analysis_options = gr.CheckboxGroup(
                        choices=[
                            "📈 Demand Forecasting",
                            "⚠️ Risk Assessment",
                            "📦 Inventory Optimization",
                            "🤝 Supplier Performance",
                            "🌿 Sustainability Analysis"
                        ],
                        label="Choose analyses to perform",
                        value=[]
                    )
                
                analyze_button = gr.Button(
                    "🚀 Run Analysis",
                    variant="primary",
                    elem_classes="action-button"
                )
                
                with gr.Row():
                    with gr.Column(scale=2):
                        analysis_output = gr.Textbox(
                            label="Analysis Results",
                            elem_classes="result-box"
                        )
                    with gr.Column(scale=3):
                        plot_output = gr.Plot(
                            label="Visualization",
                            elem_classes="chart-container"
                        )
                        processing_status = gr.Textbox(
                            label="Processing Status",
                            elem_classes="status-box"
                        )

            # Cost Prediction Tab
            with gr.Tab("💰 Cost Prediction", elem_classes="tab-content"):
                with gr.Row():
                    with gr.Column():
                        shipment_mode = gr.Dropdown(
                            choices=["✈️ Air", "🚢 Ocean", "🚛 Truck"],
                            label="Transport Mode",
                            value="✈️ Air"
                        )
                        
                        # Basic Parameters
                        weight = gr.Slider(
                            label="📦 Weight (kg)",
                            minimum=1,
                            maximum=10000,
                            step=1,
                            value=1000
                        )
                        line_item_value = gr.Slider(
                            label="💵 Item Value (USD)",
                            minimum=1,
                            maximum=1000000,
                            step=1,
                            value=10000
                        )
                        cost_per_kg = gr.Slider(
                            label="💲 Base Cost per kg (USD)",
                            minimum=1,
                            maximum=500,
                            step=1,
                            value=50
                        )
                        
                        # Advanced Parameters
                        gr.Markdown("### Advanced Parameters")
                        transit_time = gr.Slider(
                            label="🕒 Transit Time (Days)",
                            minimum=1,
                            maximum=60,
                            step=1,
                            value=7
                        )
                        distance = gr.Slider(
                            label="📏 Distance (km)",
                            minimum=100,
                            maximum=20000,
                            step=100,
                            value=1000
                        )
                        fuel_surcharge = gr.Slider(
                            label="⛽ Fuel Surcharge (%)",
                            minimum=0,
                            maximum=50,
                            step=0.5,
                            value=5
                        )
                        
                        # Risk Factors
                        gr.Markdown("### Risk Factors")
                        insurance_rate = gr.Slider(
                            label="🛡️ Insurance Rate (%)",
                            minimum=0.1,
                            maximum=10,
                            step=0.1,
                            value=1
                        )
                        customs_duty = gr.Slider(
                            label="🏛️ Customs Duty (%)",
                            minimum=0,
                            maximum=40,
                            step=0.5,
                            value=5
                        )
                        
                        # Special Handling
                        gr.Markdown("### Special Handling")
                        special_handling = gr.CheckboxGroup(
                            choices=[
                                "Temperature Controlled",
                                "Hazardous Materials",
                                "Fragile Items",
                                "Express Delivery",
                                "Door-to-Door Service"
                            ],
                            label="Special Requirements"
                        )
                
                predict_button = gr.Button(
                    "🔍 Calculate Total Cost",
                    variant="primary",
                    elem_classes="action-button"
                )
                with gr.Row():
                    freight_result = gr.Number(
                        label="Base Freight Cost (USD)",
                        elem_classes="result-box"
                    )
                    total_cost = gr.Number(
                        label="Total Cost Including All Charges (USD)",
                        elem_classes="result-box"
                    )
                    cost_breakdown = gr.JSON(
                        label="Cost Breakdown",
                        elem_classes="result-box"
                    )

            # Chat Tab
            with gr.Tab("💬 Chat", elem_classes="tab-content"):
                chatbot = gr.Chatbot(
                    label="Chat History",
                    elem_classes="chat-container",
                    height=400
                )
                with gr.Row():
                    msg = gr.Textbox(
                        label="Message",
                        placeholder="Ask about your supply chain data...",
                        scale=4
                    )
                    chat_button = gr.Button(
                        "📤 Send",
                        variant="primary",
                        scale=1,
                        elem_classes="action-button"
                    )

            # Report Tab
            with gr.Tab("📑 Report", elem_classes="tab-content"):
                report_options = gr.CheckboxGroup(
                    choices=[
                        "📈 Demand Forecasting",
                        "⚠️ Risk Assessment",
                        "📦 Inventory Optimization",
                        "🤝 Supplier Performance",
                        "🌿 Sustainability Analysis"
                    ],
                    label="Select sections to include",
                    value=[]
                )
                report_button = gr.Button(
                    "📄 Generate Report",
                    variant="primary",
                    elem_classes="action-button"
                )
                report_download = gr.File(
                    label="Download Report"
                )

        # Event Handlers
        upload_button.click(
            fn=lambda *args: process_uploaded_data(state, *args),
            inputs=[sales_data_upload, supplier_data_upload, text_input_area],
            outputs=[upload_status])
        
        analyze_button.click(
            fn=lambda choices, sales, supplier, text: run_analyses(state, choices, sales, supplier, text),
            inputs=[analysis_options, sales_data_upload, supplier_data_upload, text_input_area],
            outputs=[analysis_output, plot_output, processing_status]
        )

        predict_button.click(
            fn=lambda mode, w, val, cost, time, dist, fuel, ins, duty, special: predict_and_store_freight(
                state,
                {
                    'shipment_mode': mode,
                    'weight': w,
                    'line_item_value': val,
                    'cost_per_kg': cost,
                    'transit_time': time,
                    'distance': dist,
                    'fuel_surcharge': fuel,
                    'insurance_rate': ins,
                    'customs_duty': duty,
                    'special_handling': special
                }
            ),
            inputs=[
                shipment_mode, weight, line_item_value, cost_per_kg,
                transit_time, distance, fuel_surcharge,
                insurance_rate, customs_duty, special_handling
            ],
            outputs=[freight_result, total_cost, cost_breakdown]
        )
        
        chat_button.click(
            fn=lambda message: chat_with_navigator(state, message),
            inputs=[msg],
            outputs=[chatbot]
        )
        
        report_button.click(
            fn=lambda options: generate_pdf_report(state, options),
            inputs=[report_options],
            outputs=[report_download]
        )
        
        # Footer
        gr.HTML(
            '''<div class="footer">
                Made with 🧠 by Aditya Ratan
            </div>'''
        )

    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        debug=True
    )