import os
import json
import pandas as pd
import gradio as gr
from typing import Annotated, Dict, List, Any
from tavily import TavilyClient
from autogen import AssistantAgent, UserProxyAgent, register_function, Cache
from autogen.agentchat import GroupChat, GroupChatManager
import plotly.graph_objects as go
import plotly.express as px
from datetime import datetime
import asyncio
import threading
import time
import io
import docx
from pptx import Presentation
import PyPDF2
import boto3
from botocore.exceptions import ClientError

# Disable Docker globally
os.environ["AUTOGEN_USE_DOCKER"] = "0"

class SupplyChainOptimizer:
    def __init__(self):
        # Initialize API keys from environment variables
        self.aws_access_key = os.environ.get("AWS_ACCESS_KEY_ID")
        self.aws_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
        self.tavily_api_key = os.environ.get("TAVILY_API_KEY") or os.environ.get("TAVILY_KEY")
        
        # For development/testing, allow demo mode
        self.demo_mode = False
        if not self.aws_access_key or not self.aws_secret_key or not self.tavily_api_key:
            print("API keys not found. Running in demo mode.")
            self.demo_mode = True
        else:
            # Initialize Bedrock client
            try:
                self.bedrock_client = boto3.client(
                    'bedrock-runtime',
                    aws_access_key_id=self.aws_access_key,
                    aws_secret_access_key=self.aws_secret_key,
                    region_name='us-east-1'  # or your preferred region
                )
            except Exception as e:
                print(f"Error initializing Bedrock: {e}")
                self.demo_mode = True

        # Initialize Tavily client (separate from demo mode since it's optional)
        if self.tavily_api_key:
            try:
                self.tavily = TavilyClient(api_key=self.tavily_api_key)
                print("Tavily client initialized successfully")
            except Exception as e:
                print(f"Error initializing Tavily: {e}")
                self.tavily = None
        else:
            print("Tavily API key not found. Search functionality will use demo mode.")
            self.tavily = None

        # Initialize agents
        self._setup_agents()
        
        # Store results
        self.latest_analysis = ""
        self.latest_optimization = ""
        self.search_results = ""

    def call_claude_api(self, prompt, system_message=""):
        """Call Claude via AWS Bedrock"""
        if self.demo_mode:
            return "Demo mode response - AI analysis would appear here with real API keys"
        
        try:
            body = {
                "anthropic_version": "bedrock-2023-05-31",
                "max_tokens": 4000,
                "system": system_message,
                "messages": [
                    {
                        "role": "user",
                        "content": prompt
                    }
                ]
            }
            response = self.bedrock_client.invoke_model(
                modelId="anthropic.claude-3-haiku-20240307-v1:0",
                body=json.dumps(body)
            )
            response_body = json.loads(response['body'].read())
            return response_body['content'][0]['text']
        except Exception as e:
            return f"Error calling Claude API: {str(e)}"   

    def _setup_agents(self):
        """Setup agents with AWS Bedrock Claude"""
        # Keep user proxy for compatibility but won't be used with direct API calls
        self.user_proxy = None
        
        # Only setup if not in demo mode
        if not self.demo_mode:
            # Initialize Bedrock client if not already done
            if not hasattr(self, 'bedrock_client'):
                try:
                    self.bedrock_client = boto3.client(
                        'bedrock-runtime',
                        aws_access_key_id=self.aws_access_key,
                        aws_secret_access_key=self.aws_secret_key,
                        region_name='us-east-1'  # or your preferred region
                    )
                except Exception as e:
                    print(f"Error initializing Bedrock: {e}")
                    self.demo_mode = True
                    self.reasoning_agent = None
                    self.optimization_agent = None
                    return
            
            # Store system messages for direct API calls
            self.reasoning_system_message = """You are an expert supply chain analyst. Analyze forecast data using real-time search results.
            Evaluate if forecasts are reasonable based on current market conditions, events, and trends.
            Provide detailed reasoning and recommendations. Format your response clearly with proper structure and avoid using asterisks for emphasis."""
            
            self.optimization_system_message = """You are a supply chain optimization expert. Create detailed redistribution plans.
            Consider costs, travel time, inventory levels, and demand forecasts.
            Provide step-by-step optimization plans with clear recommendations and cost analysis. Use professional formatting without asterisks."""
            
            # Set agents as enabled (we'll use direct API calls)
            self.reasoning_agent = "claude-enabled"
            self.optimization_agent = "claude-enabled"
        else:
            self.reasoning_agent = None
            self.optimization_agent = None

    def tavily_search_tool(self, query: Annotated[str, "Market Research Query"]) -> Annotated[str, "Search results"]:
        """Search tool using Tavily API"""
        if self.demo_mode or not self.tavily:
            return f"Demo Mode: Search query '{query}' - Market conditions show stable tourism activity with moderate demand fluctuations in the specified regions."
        
        try:
            results = self.tavily.get_search_context(query=query, search_depth="advanced")
            self.search_results = results
            return results
        except Exception as e:
            return f"Search error: {str(e)}"

    def parse_file_content(self, file_path, file_type):
        """Parse various file formats and extract data"""
        try:
            if file_type == "excel":
                # Try to read Excel file
                df = pd.read_excel(file_path, sheet_name=None)
                return self._process_excel_data(df)
            elif file_type == "csv":
                df = pd.read_csv(file_path)
                return self._process_csv_data(df)
            elif file_type == "pdf":
                return self._process_pdf_data(file_path)
            elif file_type == "word":
                return self._process_word_data(file_path)
            elif file_type == "ppt":
                return self._process_ppt_data(file_path)
            else:
                return "Unsupported file format"
        except Exception as e:
            return f"Error processing file: {str(e)}"

    def _process_excel_data(self, excel_data):
        """Process Excel data and extract forecast, inventory, and route information"""
        processed_data = {
            'forecast': [],
            'inventory': [],
            'routes': []
        }
        
        for sheet_name, df in excel_data.items():
            if 'forecast' in sheet_name.lower():
                processed_data['forecast'] = df.to_dict('records')
            elif 'inventory' in sheet_name.lower():
                processed_data['inventory'] = df.to_dict('records')
            elif 'route' in sheet_name.lower():
                processed_data['routes'] = df.to_dict('records')
        
        return processed_data

    def _process_csv_data(self, df):
        # Clean column names first
        df.columns = df.columns.str.strip().str.lower()
        
        # Let AI analyze the structure and map columns
        column_analysis = self._analyze_columns_with_ai(df)
        
        return {
            'data': df.to_dict('records'),
            'column_mapping': column_analysis,
            'original_columns': df.columns.tolist()
        }

    def _process_pdf_data(self, file_path):
        """Extract text from PDF"""
        try:
            with open(file_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                text = ""
                for page in pdf_reader.pages:
                    text += page.extract_text()
                return {'text': text}
        except Exception as e:
            return f"Error reading PDF: {str(e)}"

    def _process_word_data(self, file_path):
        """Extract text from Word document"""
        try:
            doc = docx.Document(file_path)
            text = ""
            for paragraph in doc.paragraphs:
                text += paragraph.text + "\n"
            return {'text': text}
        except Exception as e:
            return f"Error reading Word document: {str(e)}"

    def _process_ppt_data(self, file_path):
        """Extract text from PowerPoint"""
        try:
            prs = Presentation(file_path)
            text = ""
            for slide in prs.slides:
                for shape in slide.shapes:
                    if hasattr(shape, "text"):
                        text += shape.text + "\n"
            return {'text': text}
        except Exception as e:
            return f"Error reading PowerPoint: {str(e)}"

    def _analyze_columns_with_ai(self, df):
        """Use AI to understand column structure and map to standard format"""
        sample_data = df.head(3).to_string()
        columns = df.columns.tolist()
    
        prompt = f"""
        Analyze this data structure and map columns to standard supply chain format:
    
        Columns: {columns}
        Sample Data:
        {sample_data}
    
        Map these columns to:
        - city/location: (identify city/location column)
        - product: (identify product column) 
        - demand/forecast: (identify demand/forecast column)
        - stock/inventory: (identify stock/inventory column)
        - cost: (identify cost column)
        - distance: (identify distance column)
    
        Return JSON mapping like: {{"city": "actual_column_name", "product": "actual_column_name", ...}}
        """
    
        if self.demo_mode:
            # Return best guess mapping
            mapping = {}
            for col in columns:
                col_lower = col.lower()
                if any(word in col_lower for word in ['city', 'location', 'destination', 'source']):
                    mapping['city'] = col
                elif any(word in col_lower for word in ['product', 'item', 'sku']):
                    mapping['product'] = col
                elif any(word in col_lower for word in ['demand', 'forecast', 'required']):
                    mapping['demand'] = col
                elif any(word in col_lower for word in ['stock', 'inventory', 'level']):
                    mapping['stock'] = col
            return mapping   
        else:
            response = self.call_claude_api(prompt, "You are a data analyst expert at understanding file structures.")
            # Parse JSON response
            try:
                return json.loads(response)
            except:
                return self._fallback_column_mapping(columns)

    def _fallback_column_mapping(self, columns):
        """Fallback column mapping if AI parsing fails"""
        mapping = {}
        for col in columns:
            col_lower = col.lower()
            if any(word in col_lower for word in ['city', 'location', 'destination', 'source']):
                mapping['city'] = col
            elif any(word in col_lower for word in ['product', 'item', 'sku']):
                mapping['product'] = col
            elif any(word in col_lower for word in ['demand', 'forecast', 'required']):
                mapping['demand'] = col
            elif any(word in col_lower for word in ['stock', 'inventory', 'level']):
                mapping['stock'] = col
            elif any(word in col_lower for word in ['cost', 'price']):
                mapping['cost'] = col
            elif any(word in col_lower for word in ['distance', 'km', 'miles']):
                mapping['distance'] = col
        return mapping

    def analyze_file_with_ai(self, file_obj, data_type):
        """Analyze uploaded file and standardize data format"""
        try:
            # Get file extension
            file_name = file_obj.name
            if file_name.endswith('.csv'):
                df = pd.read_csv(file_obj.name)
            elif file_name.endswith(('.xlsx', '.xls')):
                df = pd.read_excel(file_obj.name)
            else:
                return {'standardized_data': [], 'detected_columns': [], 'error': 'Unsupported file format'}
            
            # Clean column names
            df.columns = df.columns.str.strip()
            detected_columns = df.columns.tolist()
            
            # Map columns based on data type
            column_mapping = self._analyze_columns_with_ai(df)
            
            # Standardize data based on type
            standardized_data = self._standardize_data(df, column_mapping, data_type)
            
            return {
                'standardized_data': standardized_data,
                'detected_columns': detected_columns,
                'column_mapping': column_mapping
            }
        except Exception as e:
            return {'standardized_data': [], 'detected_columns': [], 'error': str(e)}

    def _standardize_data(self, df, column_mapping, data_type):
        """Standardize data format based on type"""
        standardized = []
        
        try:
            if data_type == 'forecast':
                for _, row in df.iterrows():
                    item = {
                        'City': row.get(column_mapping.get('city', ''), 'Unknown'),
                        'Product': row.get(column_mapping.get('product', ''), 'Unknown'),
                        'Forecasted_Demand': int(row.get(column_mapping.get('demand', ''), 0)),
                        'Month': 'December'  # Default month
                    }
                    standardized.append(item)
            
            elif data_type == 'inventory':
                for _, row in df.iterrows():
                    item = {
                        'City': row.get(column_mapping.get('city', ''), 'Unknown'),
                        'Product': row.get(column_mapping.get('product', ''), 'Unknown'),
                        'Stock_Level': int(row.get(column_mapping.get('stock', ''), 0))
                    }
                    standardized.append(item)
            
            elif data_type == 'routes':
                for _, row in df.iterrows():
                    item = {
                        'Source': row.get(column_mapping.get('source', ''), 'Unknown'),
                        'Destination': row.get(column_mapping.get('destination', ''), 'Unknown'),
                        'Distance_km': float(row.get(column_mapping.get('distance', ''), 0)),
                        'Cost_per_km': float(row.get(column_mapping.get('cost', ''), 0)),
                        'Average_Travel_Time_hrs': float(row.get(column_mapping.get('time', ''), 0))
                    }
                    standardized.append(item)
        
        except Exception as e:
            print(f"Error standardizing data: {e}")
            return []
        
        return standardized

    def generate_data_from_text(self, text_input):
        """Generate sample data based on text description"""
        prompt = f"""
        Based on this business description, generate sample supply chain data:
        
        Text: {text_input}
        
        Generate realistic data for:
        1. Forecast data (cities, products, demand)
        2. Inventory data (cities, products, stock levels)
        3. Route data (source, destination, distance, cost, travel time)
        
        Return as JSON with keys: forecast, inventory, routes
        Each should be a list of dictionaries with appropriate fields.
        """
        
        if self.demo_mode:
            # Return default data
            return {
                'forecast': DEFAULT_FORECAST,
                'inventory': DEFAULT_INVENTORY,
                'routes': DEFAULT_ROUTES
            }
        else:
            try:
                response = self.call_claude_api(prompt, "You are a supply chain data expert.")
                return json.loads(response)
            except:
                return {
                    'forecast': DEFAULT_FORECAST,
                    'inventory': DEFAULT_INVENTORY,
                    'routes': DEFAULT_ROUTES
                }

    def create_forecast_visualization(self, forecast_data):
        """Create interactive forecast visualization with vibrant colors"""
        if not forecast_data:
            forecast_data = DEFAULT_FORECAST
            
        df = pd.DataFrame(forecast_data)
        
        fig = go.Figure()
        
        # Vibrant color palette
        colors = ['#FF6B35', '#F7931E', '#FFD23F', '#06FFA5', '#4ECDC4', '#45B7D1', '#96CEB4', '#FFEAA7', '#DDA0DD', '#FA8072']
        
        for i, (city, group) in enumerate(df.groupby('City')):
            fig.add_trace(go.Bar(
                name=city,
                x=group['Product'],
                y=group['Forecasted_Demand'],
                marker_color=colors[i % len(colors)],
                text=group['Forecasted_Demand'],
                textposition='auto',
                textfont=dict(size=14, color='white', family='Arial Black'),
            ))
        
        fig.update_layout(
            title={
                'text': "Demand Forecast by City & Product",
                'x': 0.5,
                'xanchor': 'center',
                'font': {'size': 20, 'color': '#B8860B', 'family': 'Arial Black'}
            },
            xaxis_title="Products",
            yaxis_title="Forecasted Demand",
            xaxis=dict(title=dict(font=dict(size=16, color='#8B4513', family='Arial Black'))),
            yaxis=dict(title=dict(font=dict(size=16, color='#8B4513', family='Arial Black'))),
            barmode='group',
            plot_bgcolor='rgba(255,248,220,0.9)',
            paper_bgcolor='rgba(255,248,220,0.9)',
            font=dict(color='#8B4513', size=12, family='Arial'),
            height=500,
            legend=dict(font=dict(size=12, color='#8B4513', family='Arial'))
        )
        
        return fig

    def create_inventory_chart(self, inventory_data, forecast_data):
        """Create inventory vs demand comparison with vibrant styling"""
        if not inventory_data:
            inventory_data = DEFAULT_INVENTORY
        if not forecast_data:
            forecast_data = DEFAULT_FORECAST
            
        inv_df = pd.DataFrame(inventory_data)
        fore_df = pd.DataFrame(forecast_data)
        
        # Merge data
        merged = pd.merge(inv_df, fore_df, on=['City', 'Product'], how='outer')
        merged = merged.fillna(0)
        
        fig = go.Figure()
        
        fig.add_trace(go.Bar(
            name='Current Stock',
            x=[f"{row['City']} - {row['Product']}" for _, row in merged.iterrows()],
            y=merged['Stock_Level'],
            marker_color='#FF4757',
            opacity=0.9,
            text=merged['Stock_Level'],
            textposition='auto',
            textfont=dict(size=12, color='white', family='Arial Black')
        ))
        
        fig.add_trace(go.Bar(
            name='Forecasted Demand',
            x=[f"{row['City']} - {row['Product']}" for _, row in merged.iterrows()],
            y=merged['Forecasted_Demand'],
            marker_color='#FFA502',
            opacity=0.9,
            text=merged['Forecasted_Demand'],
            textposition='auto',
            textfont=dict(size=12, color='white', family='Arial Black')
        ))
        
        fig.update_layout(
            title={
                'text': "Inventory vs Demand Analysis",
                'x': 0.5,
                'xanchor': 'center',
                'font': {'size': 20, 'color': '#B8860B', 'family': 'Arial Black'}
            },
            xaxis_title="City & Product",
            yaxis_title="Units",
            xaxis=dict(title=dict(font=dict(size=16, color='#8B4513', family='Arial Black'))),
            yaxis=dict(title=dict(font=dict(size=16, color='#8B4513', family='Arial Black'))),
            barmode='group',
            plot_bgcolor='rgba(255,248,220,0.9)',
            paper_bgcolor='rgba(255,248,220,0.9)',
            font=dict(color='#8B4513', size=12, family='Arial'),
            height=500,
            legend=dict(font=dict(size=12, color='#8B4513', family='Arial'))
        )
        
        return fig

    def create_route_network(self, route_data):
        """Create route network visualization with vibrant colors"""
        if not route_data:
            route_data = DEFAULT_ROUTES
            
        df = pd.DataFrame(route_data)
        
        fig = go.Figure()
        
        fig.add_trace(go.Scatter(
            x=df['Distance_km'],
            y=df['Cost_per_km'],
            mode='markers+text',
            marker=dict(
                size=[cost*2 for cost in df['Cost_per_km']],
                color=df['Average_Travel_Time_hrs'],
                colorscale=[[0, '#FFD700'], [0.5, '#FF6347'], [1, '#DC143C']],
                showscale=True,
                colorbar=dict(title=dict(text="Travel Time (hrs)", font=dict(size=14, color='#8B4513', family='Arial Black')))
            ),
            text=[f"{row['Source']} → {row['Destination']}" for _, row in df.iterrows()],
            textposition="top center",
            textfont=dict(size=12, color='#8B0000', family='Arial Black'),
            name="Routes"
        ))
        
        fig.update_layout(
            title={
                'text': "Route Analysis: Distance vs Cost",
                'x': 0.5,
                'xanchor': 'center',
                'font': {'size': 20, 'color': '#B8860B', 'family': 'Arial Black'}
            },
            xaxis_title="Distance (km)",
            yaxis_title="Cost per km (₹)",
            xaxis=dict(title=dict(font=dict(size=16, color='#8B4513', family='Arial Black'))),
            yaxis=dict(title=dict(font=dict(size=16, color='#8B4513', family='Arial Black'))),
            plot_bgcolor='rgba(255,248,220,0.9)',
            paper_bgcolor='rgba(255,248,220,0.9)',
            font=dict(color='#8B4513', size=12, family='Arial'),
            height=500
        )
        
        return fig

    def optimize_supply_chain(self, forecast_data, inventory_data, route_data, search_query, progress=gr.Progress()):
        """Main optimization function"""
        progress(0.1, desc="Conducting market research...")
        
        # Conduct search
        search_results = self.tavily_search_tool(search_query)
        
        progress(0.3, desc="Analyzing forecast data...")
        
        if self.demo_mode:
            time.sleep(2)
            progress(0.6, desc="Optimizing redistribution plan...")
            
            self.latest_analysis = """
DEMO MODE - Market Analysis & Forecast Reasoning

Market Conditions Assessment:
- Current tourism trends show moderate activity in hill stations
- Seasonal demand patterns indicate December is peak season for tourist destinations
- Economic indicators suggest stable consumer spending on FMCG products

Forecast Accuracy Evaluation:
- Goa: Forecasted demand of 1,500 units appears reasonable given tourist influx
- Coorg: 700 units forecast aligns with typical seasonal patterns
- Mahabaleshwar: 1,000 units seems appropriate for weekend destination
- Lonavala: 850 units matches proximity to major cities
- Ooty: 400 units may be conservative given popularity

Risk Factors:
- Weather conditions could impact transportation
- Festival seasons may create demand spikes
- Competition from local suppliers

Recommendations:
- Monitor real-time booking data
- Prepare for demand fluctuations
- Consider safety stock adjustments

Note: Set your API keys to get real-time market intelligence
            """
            
            time.sleep(2)
            progress(0.8, desc="Generating insights...")
            
            self.latest_optimization = """
DEMO MODE - Optimized Redistribution Plan

Priority Actions:

1. Immediate Redistribution (Week 1)
   - Move 300 Biscuit units from Lonavala to Coorg (Cost: ₹6,480)
   - Route: Lonavala → Goa → Coorg (Total: 1,000 km)
   - Expected delivery: 17 hours

2. Strategic Rebalancing (Week 2)
   - Reduce Goa soap inventory by 500 units
   - Distribute to Mahabaleshwar and Ooty based on demand
   - Utilize cost-effective Goa → Ooty route (₹16,500)

3. Cost Optimization Strategies
   - Consolidate shipments to reduce per-km costs
   - Use off-peak travel times for better rates
   - Implement just-in-time delivery schedules

4. Expected Outcomes
   - Total redistribution cost: ₹45,000
   - Inventory optimization: 15% reduction in carrying costs
   - Service level improvement: 98% demand fulfillment
   - Risk mitigation: 20% safety stock maintained

5. Timeline & Priorities
   - High Priority: Coorg biscuits shortage (2 days)
   - Medium Priority: Ooty soap rebalancing (1 week)
   - Low Priority: General inventory optimization (2 weeks)

Cost-Benefit Analysis:
- Investment: ₹45,000 in transportation
- Savings: ₹65,000 in carrying costs + lost sales prevention
- Net Benefit: ₹20,000 + improved customer satisfaction

Note: Set your API keys for AI-powered optimization with real market data
            """
            
        else:
            # Real AI agent processing (similar structure but without asterisks in system messages)
            reasoning_prompt = f"""
            Analyze the forecast data against the real-time search results and evaluate if the forecast is reasonable.
            
            Forecast Data: {json.dumps(forecast_data, indent=2)}
            Search Results: {search_results}
            
            Please provide detailed reasoning on:
            1. Market conditions affecting demand
            2. External factors from search results
            3. Forecast accuracy assessment
            4. Risk factors to consider
            
            Please format your response professionally without using asterisks for emphasis.
            """
            
            try:
                system_msg = "You are an expert supply chain analyst. Analyze forecast data using real-time search results. Evaluate if forecasts are reasonable based on current market conditions, events, and trends. Provide detailed reasoning and recommendations. Format your response clearly with proper structure."
                
                self.latest_analysis = self.call_claude_api(reasoning_prompt, system_msg)
            
            except Exception as e:
                self.latest_analysis = f"Analysis error: {str(e)}"
            
            progress(0.6, desc="Optimizing redistribution plan...")
            
            # Optimization phase
            optimization_prompt = f"""
            Create an optimized redistribution plan using the analyzed data.
            
            Forecast Data: {json.dumps(forecast_data, indent=2)}
            Inventory Data: {json.dumps(inventory_data, indent=2)}
            Route Data: {json.dumps(route_data, indent=2)}
            
            Analysis Context: {self.latest_analysis}
            
            Please provide:
            1. Redistribution recommendations
            2. Cost optimization strategies
            3. Risk mitigation plans
            4. Timeline and priorities
            5. Expected cost savings
            
            Please format your response professionally without using asterisks for emphasis.
            """
            
            try:
                system_msg = "You are a supply chain optimization expert. Create detailed redistribution plans. Consider costs, travel time, inventory levels, and demand forecasts. Provide step-by-step optimization plans with clear recommendations and cost analysis."
                
                self.latest_optimization = self.call_claude_api(optimization_prompt, system_msg)
            except Exception as e:
                self.latest_optimization = f"Optimization error: {str(e)}"
        
        progress(1.0, desc="Optimization complete!")
        
        return self.latest_analysis, self.latest_optimization, search_results

# Initialize the optimizer
try:
    optimizer = SupplyChainOptimizer()
    startup_message = "Supply Chain Optimizer initialized successfully!"
    if optimizer.demo_mode:
        startup_message = "Running in DEMO mode. Set API keys for full functionality."
except Exception as e:
    print(f"Error initializing optimizer: {e}")
    class DemoOptimizer:
        def __init__(self):
            self.demo_mode = True
        
        def create_forecast_visualization(self, data):
            return go.Figure().add_annotation(text="Set API keys to enable charts")
        
        def create_inventory_chart(self, inv, fore):
            return go.Figure().add_annotation(text="Set API keys to enable charts")
        
        def create_route_network(self, routes):
            return go.Figure().add_annotation(text="Set API keys to enable charts")
        
        def optimize_supply_chain(self, f, i, r, q, progress=None):
            return "Demo mode", "Demo mode", "Demo mode"
        
        def parse_file_content(self, path, file_type):
            return "Demo mode"
        
        def analyze_file_with_ai(self, file_obj, data_type):
            return {'standardized_data': DEFAULT_FORECAST if data_type == 'forecast' else DEFAULT_INVENTORY if data_type == 'inventory' else DEFAULT_ROUTES, 'detected_columns': [], 'error': None}
        
        def generate_data_from_text(self, text):
            return {'forecast': DEFAULT_FORECAST, 'inventory': DEFAULT_INVENTORY, 'routes': DEFAULT_ROUTES}
    
    optimizer = DemoOptimizer()
    startup_message = "Running in minimal demo mode due to initialization error."

# Default data
DEFAULT_FORECAST = [
    {"City": "Goa", "Product": "Soap", "Forecasted_Demand": 1500, "Month": "December"},
    {"City": "Coorg", "Product": "Biscuits", "Forecasted_Demand": 700, "Month": "December"},
    {"City": "Mahabaleshwar", "Product": "Soap", "Forecasted_Demand": 1000, "Month": "December"},
    {"City": "Lonavala", "Product": "Biscuits", "Forecasted_Demand": 850, "Month": "December"},
    {"City": "Ooty", "Product": "Soap", "Forecasted_Demand": 400, "Month": "December"}
]

DEFAULT_INVENTORY = [
    {"City": "Goa", "Product": "Soap", "Stock_Level": 2000},
    {"City": "Coorg", "Product": "Biscuits", "Stock_Level": 400},
    {"City": "Mahabaleshwar", "Product": "Soap", "Stock_Level": 1100},
    {"City": "Lonavala", "Product": "Biscuits", "Stock_Level": 800},
    {"City": "Ooty", "Product": "Soap", "Stock_Level": 500}
]

DEFAULT_ROUTES = [
    {"Source": "Goa", "Destination": "Coorg", "Distance_km": 550, "Cost_per_km": 20, "Average_Travel_Time_hrs": 10},
    {"Source": "Goa", "Destination": "Ooty", "Distance_km": 750, "Cost_per_km": 22, "Average_Travel_Time_hrs": 15},
    {"Source": "Coorg", "Destination": "Goa", "Distance_km": 550, "Cost_per_km": 20, "Average_Travel_Time_hrs": 10},
    {"Source": "Coorg", "Destination": "Mahabaleshwar", "Distance_km": 400, "Cost_per_km": 18, "Average_Travel_Time_hrs": 8},
    {"Source": "Lonavala", "Destination": "Goa", "Distance_km": 450, "Cost_per_km": 19, "Average_Travel_Time_hrs": 7},
    {"Source": "Ooty", "Destination": "Coorg", "Distance_km": 800, "Cost_per_km": 21, "Average_Travel_Time_hrs": 16}
]

def process_files_and_optimize(forecast_file, inventory_file, routes_file, text_input, search_query):
    """Process uploaded files and text input for optimization"""
    try:
        # Let AI analyze files instead of using defaults
        forecast_data = []
        inventory_data = []
        route_data = []
        file_contents = []
        
        # AI-powered file processing
        if forecast_file:
            analyzed_data = optimizer.analyze_file_with_ai(forecast_file, 'forecast')
            forecast_data = analyzed_data['standardized_data']
            file_contents.append(f"Forecast file analyzed: {forecast_file.name} - Found columns: {analyzed_data.get('detected_columns', 'N/A')}")
        
        if inventory_file:
            analyzed_data = optimizer.analyze_file_with_ai(inventory_file, 'inventory')
            inventory_data = analyzed_data['standardized_data']
            file_contents.append(f"Inventory file analyzed: {inventory_file.name} - Found columns: {analyzed_data.get('detected_columns', 'N/A')}")
        
        if routes_file:
            analyzed_data = optimizer.analyze_file_with_ai(routes_file, 'routes')
            route_data = analyzed_data['standardized_data']
            file_contents.append(f"Routes file analyzed: {routes_file.name} - Found columns: {analyzed_data.get('detected_columns', 'N/A')}")
        
        # If no files uploaded, use defaults or generate from text
        if not any([forecast_file, inventory_file, routes_file]):
            if text_input and text_input.strip():
                ai_generated_data = optimizer.generate_data_from_text(text_input)
                forecast_data = ai_generated_data.get('forecast', DEFAULT_FORECAST)
                inventory_data = ai_generated_data.get('inventory', DEFAULT_INVENTORY)
                route_data = ai_generated_data.get('routes', DEFAULT_ROUTES)
                file_contents.append("AI generated data from text description")
            else:
                # Use default data
                forecast_data = DEFAULT_FORECAST
                inventory_data = DEFAULT_INVENTORY
                route_data = DEFAULT_ROUTES
                file_contents.append("Using default sample data")
        
        # Process text input if provided
        if text_input and text_input.strip():
            file_contents.append(f"Text context processed: {len(text_input)} characters")
        
        # Create visualizations
        forecast_chart = optimizer.create_forecast_visualization(forecast_data)
        inventory_chart = optimizer.create_inventory_chart(inventory_data, forecast_data)
        route_chart = optimizer.create_route_network(route_data)
        
        # Run optimization
        analysis, optimization, search_results = optimizer.optimize_supply_chain(
            forecast_data, inventory_data, route_data, search_query
        )
        
        processing_summary = "Files processed:\n" + "\n".join(file_contents) if file_contents else "No data provided"
        
        return (
            forecast_chart,
            inventory_chart, 
            route_chart,
            analysis,
            optimization,
            search_results[:2000] + "..." if len(search_results) > 2000 else search_results,
            processing_summary
        )
        
    except Exception as e:
        error_msg = f"Processing error: {str(e)}"
        empty_fig = go.Figure().add_annotation(text=f"Error: {str(e)}", x=0.5, y=0.5, showarrow=False)
        return empty_fig, empty_fig, empty_fig, error_msg, error_msg, error_msg, error_msg

# Create Gradio interface with updated warm color scheme
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');

.gradio-container {
    font-family: 'Inter', sans-serif;
    background: linear-gradient(135deg, #FFF8DC 0%, #FFEBCD 100%);
    min-height: 100vh;
}

.main-header {
    text-align: center;
    background: linear-gradient(135deg, #DC143C 0%, #FF6347 50%, #FFD700 100%);
    color: white;
    padding: 2.5rem;
    border-radius: 15px;
    margin-bottom: 2rem;
    box-shadow: 0 8px 25px rgba(220, 20, 60, 0.3);
    border: 2px solid #B8860B;
}

.main-header h1 {
    font-size: 2.5rem;
    font-weight: 700;
    margin-bottom: 1rem;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
}

.main-header p {
    font-size: 1.1rem;
    font-weight: 500;
    text-shadow: 1px 1px 2px rgba(0,0,0,0.2);
}

.section-header {
    background: linear-gradient(90deg, #FF4757 0%, #FFA502 100%);
    color: white;
    padding: 1.2rem;
    border-radius: 10px;
    margin: 1rem 0;
    text-align: center;
    font-weight: 600;
    font-size: 1.1rem;
    text-shadow: 1px 1px 2px rgba(0,0,0,0.2);
    border: 2px solid #B8860B;
    box-shadow: 0 4px 15px rgba(255, 71, 87, 0.2);
}

.results-container {
    background: linear-gradient(135deg, #FFFACD 0%, #F5DEB3 100%);
    padding: 2rem;
    border-radius: 12px;
    margin: 1rem 0;
    box-shadow: 0 4px 20px rgba(184, 134, 11, 0.15);
    border: 2px solid #DAA520;
}

.footer {
    text-align: center;
    color: #8B4513;
    padding: 2rem;
    border-top: 2px solid #DAA520;
    margin-top: 2rem;
    background: linear-gradient(135deg, #FFF8DC 0%, #FFEBCD 100%);
    border-radius: 10px;
    font-weight: 500;
}

.upload-area {
    border: 3px dashed #DAA520;
    border-radius: 10px;
    padding: 1.5rem;
    margin: 0.5rem 0;
    background: linear-gradient(135deg, #FFFACD 0%, #F0E68C 100%);
    transition: all 0.3s ease;
}

.upload-area:hover {
    border-color: #B8860B;
    background: linear-gradient(135deg, #F0E68C 0%, #DAA520 100%);
    transform: translateY(-2px);
}

.gradio-button {
    background: linear-gradient(135deg, #FF4757 0%, #FFA502 100%) !important;
    color: white !important;
    font-weight: 600 !important;
    border: 2px solid #B8860B !important;
    border-radius: 8px !important;
    padding: 12px 24px !important;
    font-size: 1rem !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(255, 71, 87, 0.2) !important;
}

.gradio-button:hover {
    background: linear-gradient(135deg, #FFA502 0%, #FF4757 100%) !important;
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(255, 71, 87, 0.3) !important;
}

.gradio-textbox, .gradio-dropdown {
    border: 2px solid #DAA520 !important;
    border-radius: 8px !important;
    background: linear-gradient(135deg, #FFFACD 0%, #F5DEB3 100%) !important;
}

.gradio-textbox:focus, .gradio-dropdown:focus {
    border-color: #B8860B !important;
    box-shadow: 0 0 10px rgba(184, 134, 11, 0.3) !important;
}
"""

# Create the Gradio interface
with gr.Blocks(css=custom_css, title="AI-Powered Supply Chain Optimizer") as interface:
    # Header
    gr.HTML("""
    <div class="main-header">
        <h1>🚚 AI-Powered Supply Chain Optimizer</h1>
        <p>Optimize your supply chain with real-time market intelligence and advanced analytics</p>
    </div>
    """)
    
    # Status message
    gr.HTML(f"""
    <div style="text-align: center; padding: 1rem; background: #E6FFE6; border-radius: 8px; margin-bottom: 1rem; border: 2px solid #90EE90;">
        <strong>System Status:</strong> {startup_message}
    </div>
    """)
    
    with gr.Tabs():
        # Tab 1: File Upload and Input
        with gr.TabItem("📁 Data Input", elem_id="input-tab"):
            gr.HTML('<div class="section-header">Upload Your Data Files</div>')
            
            with gr.Row():
                with gr.Column():
                    forecast_file = gr.File(
                        label="📊 Forecast Data (Excel/CSV)",
                        file_types=[".xlsx", ".xls", ".csv"],
                        elem_classes=["upload-area"]
                    )
                    inventory_file = gr.File(
                        label="📦 Inventory Data (Excel/CSV)",
                        file_types=[".xlsx", ".xls", ".csv"],
                        elem_classes=["upload-area"]
                    )
                with gr.Column():
                    routes_file = gr.File(
                        label="🗺️ Routes Data (Excel/CSV)",
                        file_types=[".xlsx", ".xls", ".csv"],
                        elem_classes=["upload-area"]
                    )
                    text_input = gr.Textbox(
                        label="📝 Additional Context (Optional)",
                        placeholder="Enter any additional business context, constraints, or special requirements...",
                        lines=4,
                        elem_classes=["upload-area"]
                    )
            
            search_query = gr.Textbox(
                label="🔍 Market Research Query",
                placeholder="Enter search terms for real-time market analysis (e.g., 'tourism trends December 2024 hill stations')",
                value="tourism trends December 2024 hill stations demand forecast",
                lines=2
            )
            
            optimize_btn = gr.Button(
                "🚀 Optimize Supply Chain",
                variant="primary",
                size="lg",
                elem_classes=["gradio-button"]
            )
            
            processing_status = gr.Textbox(
                label="📋 Processing Status",
                interactive=False,
                lines=3
            )
        
        # Tab 2: Visualizations
        with gr.TabItem("📈 Analytics Dashboard", elem_id="viz-tab"):
            gr.HTML('<div class="section-header">Interactive Data Visualizations</div>')
            
            with gr.Row():
                forecast_plot = gr.Plot(label="📊 Demand Forecast Analysis")
                inventory_plot = gr.Plot(label="📦 Inventory vs Demand")
            
            route_plot = gr.Plot(label="🗺️ Route Network Analysis")
        
        # Tab 3: AI Analysis Results
        with gr.TabItem("🤖 AI Analysis", elem_id="analysis-tab"):
            gr.HTML('<div class="section-header">AI-Powered Market Analysis & Recommendations</div>')
            
            with gr.Row():
                with gr.Column():
                    gr.HTML('<h3 style="color: #B8860B; text-align: center;">🔍 Market Intelligence & Forecast Analysis</h3>')
                    analysis_output = gr.Textbox(
                        label="",
                        lines=15,
                        interactive=False,
                        elem_classes=["results-container"]
                    )
                
                with gr.Column():
                    gr.HTML('<h3 style="color: #B8860B; text-align: center;">⚡ Optimization Recommendations</h3>')
                    optimization_output = gr.Textbox(
                        label="",
                        lines=15,
                        interactive=False,
                        elem_classes=["results-container"]
                    )
        
        # Tab 4: Search Results
        with gr.TabItem("🌐 Market Research", elem_id="search-tab"):
            gr.HTML('<div class="section-header">Real-Time Market Intelligence</div>')
            
            search_output = gr.Textbox(
                label="🔍 Market Research Results",
                lines=20,
                interactive=False,
                elem_classes=["results-container"]
            )
    
    # Footer
    gr.HTML("""
    <div class="footer">
        <p><strong>AI-Powered Supply Chain Optimizer</strong> | Advanced Analytics & Real-Time Intelligence</p>
        <p>🔧 Built with AutoGen, Tavily API, and Claude | 🚀 Powered by AWS Bedrock</p>
    </div>
    """)
    
    # Connect the optimization function
    optimize_btn.click(
        fn=process_files_and_optimize,
        inputs=[forecast_file, inventory_file, routes_file, text_input, search_query],
        outputs=[
            forecast_plot,
            inventory_plot,
            route_plot,
            analysis_output,
            optimization_output,
            search_output,
            processing_status
        ],
        show_progress=True
    )

# Launch the interface
if __name__ == "__main__":
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        show_error=True,
        debug=True,
        inbrowser=True
    )