import streamlit as st
import os
from smolagents import CodeAgent, LiteLLMModel, tool
import json
from datetime import datetime
import random

# Hugging Face Spaces configuration
st.set_page_config(
    page_title="AI Procurement Agent", 
    page_icon="🤖", 
    layout="wide"
)

# Initialize OpenAI model for Spaces
@st.cache_resource
def get_model():
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        st.error("⚠️ OpenAI API Key not found in Spaces secrets!")
        st.info("Add OPENAI_API_KEY to your Spaces secrets in Settings.")
        st.stop()
    
    return LiteLLMModel(
        model_id="gpt-3.5-turbo",
        api_key=api_key,
        temperature=0.2
    )

# [All your tool definitions remain exactly the same]
@tool
def query_sap_pr(pr_number: str) -> str:
    """
    Query SAP purchase requisition details
    
    Args:
        pr_number: The purchase requisition number to query
    """
    materials = ["Steel Pipes", "Cement", "Electrical Cables", "Safety Equipment"]
    mock_data = {
        "pr_number": pr_number,
        "materials": f"{random.choice(materials)} - {random.randint(50, 200)} units",
        "total_value": random.randint(30000, 80000),
        "requestor": random.choice(["John Doe", "Jane Smith", "Mike Johnson"]),
        "department": random.choice(["Construction", "Manufacturing", "Maintenance"]),
        "urgency": random.choice(["High", "Medium", "Low"]),
        "delivery_date": "2025-10-15",
        "status": "Approved"
    }
    return json.dumps(mock_data, indent=2)

@tool  
def check_vendor_solvency(vendor_name: str) -> str:
    """
    Check vendor financial solvency from external sources
    
    Args:
        vendor_name: Name of the vendor to check financial solvency for
    """
    scores = ["A+", "A", "A-", "B+", "B"]
    risk_levels = ["Low", "Medium", "High"]
    mock_data = {
        "vendor_name": vendor_name,
        "credit_rating": random.choice(scores),
        "financial_stability": "Stable",
        "solvency_score": random.randint(65, 95),
        "debt_ratio": round(random.uniform(0.2, 0.6), 2),
        "risk_level": random.choice(risk_levels[:2]),
        "recommendation": "Approved for business"
    }
    return json.dumps(mock_data, indent=2)

@tool
def get_vendor_performance(vendor_name: str) -> str:
    """
    Get historical vendor performance metrics
    
    Args:
        vendor_name: Name of the vendor to get performance history for
    """
    mock_data = {
        "vendor_name": vendor_name,
        "evaluation_period": "12 months",
        "total_orders": random.randint(15, 80),
        "on_time_delivery": f"{random.randint(85, 98)}%",
        "quality_score": round(random.uniform(3.5, 5.0), 1),
        "order_accuracy": f"{random.randint(92, 99)}%",
        "avg_lead_time": f"{random.randint(7, 21)} days",
        "overall_rating": round(random.uniform(3.5, 4.8), 1),
        "issues_resolved": random.randint(0, 3)
    }
    return json.dumps(mock_data, indent=2)

@tool
def compare_material_rates(material: str) -> str:
    """
    Compare current market rates for materials
    
    Args:
        material: Type of material to compare rates for
    """
    vendors = ["Alpha Corp", "Beta Supplies", "Gamma Materials", "Delta Industries"]
    base_price = random.randint(80, 250)
    
    quotes = []
    for vendor in random.sample(vendors, 3):
        quotes.append({
            "vendor": vendor,
            "price_per_unit": base_price + random.randint(-25, 30),
            "lead_time": f"{random.randint(5, 20)} days",
            "min_order": random.randint(10, 50)
        })
    
    mock_data = {
        "material": material,
        "market_average": base_price,
        "vendor_quotes": quotes,
        "best_value_vendor": min(quotes, key=lambda x: x["price_per_unit"])["vendor"],
        "price_trend": random.choice(["Stable", "Increasing", "Decreasing"])
    }
    return json.dumps(mock_data, indent=2)

@tool
def optimize_vendor_selection(requirements: str) -> str:
    """
    Run optimization to find best vendor combination
    
    Args:
        requirements: Procurement requirements and constraints to optimize for
    """
    vendors = ["ABC Supplies", "XYZ Materials", "DEF Corp", "GHI Industries"]
    selected_vendor = random.choice(vendors)
    
    mock_data = {
        "requirements_analyzed": requirements,
        "recommended_vendor": selected_vendor,
        "optimization_score": round(random.uniform(7.5, 9.5), 1),
        "total_estimated_cost": random.randint(40000, 70000),
        "potential_savings": f"{random.randint(8, 18)}%",
        "risk_assessment": "Low to Medium",
        "confidence_level": f"{random.randint(85, 96)}%",
        "key_factors": ["Cost efficiency", "Quality track record", "Delivery reliability"]
    }
    return json.dumps(mock_data, indent=2)

@tool
def create_purchase_order(vendor: str, amount: str) -> str:
    """
    Create purchase order in SAP system
    
    Args:
        vendor: Name of the vendor to create purchase order for
        amount: Total amount for the purchase order
    """
    po_number = f"PO{datetime.now().strftime('%Y%m%d%H%M%S')[-8:]}"
    
    mock_data = {
        "po_number": po_number,
        "vendor": vendor,
        "total_amount": amount,
        "currency": "USD",
        "status": "Created Successfully",
        "payment_terms": "NET 30",
        "delivery_terms": "FOB Destination",
        "created_by": "AI_Procurement_Agent",
        "created_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "approval_required": random.choice([True, False])
    }
    return json.dumps(mock_data, indent=2)

# Initialize Agent with OpenAI - FIXED: Removed max_iterations
@st.cache_resource
def get_agent():
    tools = [
        query_sap_pr,
        check_vendor_solvency, 
        get_vendor_performance,
        compare_material_rates,
        optimize_vendor_selection,
        create_purchase_order
    ]
    
    return CodeAgent(
        tools=tools,
        model=get_model()  # Using OpenAI via LiteLLM
    )

def main():
    st.title("🤖 AI Procurement Agent")
    st.markdown("*Demo powered by Hugging Face SmolAgents + OpenAI GPT-4*")
    
    with st.expander("ℹ️ About this Demo"):
        st.markdown("""
        This AI agent can help automate procurement workflows by:
        - Querying purchase requisitions
        - Evaluating vendor financial health
        - Analyzing vendor performance history
        - Comparing material market rates
        - Optimizing vendor selection
        - Creating purchase orders
        
        **Note:** This uses mock data for demonstration purposes.
        """)
    
    # Sidebar
    with st.sidebar:
        st.header("🎯 Demo Scenarios")
        scenario = st.selectbox(
            "Choose a workflow:",
            [
                "🔄 Full Procurement Workflow",
                "📊 Vendor Evaluation", 
                "💰 Price Analysis",
                "🤖 Custom AI Query"
            ]
        )
        
        st.markdown("---")
        st.markdown("**🚀 Using OpenAI GPT-4**")
        
    # Main content
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.header("📝 Input")
        
        if scenario == "🔄 Full Procurement Workflow":
            pr_number = st.text_input("Enter PR Number:", value="PR-2025-001")
            
            if st.button("🚀 Start Full Workflow", type="primary", use_container_width=True):
                process_full_workflow(pr_number)
                
        elif scenario == "📊 Vendor Evaluation":
            vendor_name = st.text_input("Enter Vendor Name:", value="ABC Supplies")
            
            if st.button("📊 Evaluate This Vendor", use_container_width=True):
                evaluate_vendor(vendor_name)
                
        elif scenario == "💰 Price Analysis":
            material = st.text_input("Enter Material Type:", value="Steel Pipes")
            
            if st.button("💰 Analyze Market Prices", use_container_width=True):
                analyze_prices(material)
                
        else:  # Custom Query
            custom_query = st.text_area(
                "Enter your procurement query:",
                value="Evaluate vendor XYZ Corp and check if they're suitable for a 50K USD steel pipes order",
                height=80
            )
            
            if st.button("🤖 Ask AI Agent", use_container_width=True):
                run_custom_query(custom_query)
    
    with col2:
        st.header("🤖 AI Agent Response")
        
        if 'agent_response' not in st.session_state:
            st.info("👈 Select a scenario and click a button to see the AI agent in action!")
        else:
            with st.container():
                st.success("✅ Agent completed successfully!")
                
                with st.expander("📋 Full Agent Response", expanded=True):
                    st.markdown(st.session_state['agent_response'])

# [Include all the workflow functions exactly as before]
def process_full_workflow(pr_number):
    with st.spinner("🤖 AI Agent is working on your procurement request..."):
        prompt = f"""
        Execute a complete procurement workflow for PR {pr_number}:
        
        1. Query the SAP purchase requisition details
        2. For any vendors found, check their financial solvency
        3. Review their historical performance 
        4. Compare current market rates for the materials
        5. Run optimization to select the best vendor option
        6. Create a purchase order for the recommended vendor
        
        Provide a clear summary with your final recommendation and reasoning.
        """
        
        try:
            agent = get_agent()
            result = agent.run(prompt)
            st.session_state['agent_response'] = result
            st.rerun()
        except Exception as e:
            st.error(f"Error: {str(e)}")

def evaluate_vendor(vendor_name):
    with st.spinner(f"🤖 Evaluating {vendor_name}..."):
        prompt = f"""
        Conduct a comprehensive evaluation of vendor '{vendor_name}':
        
        1. Check their financial solvency and credit rating
        2. Analyze their historical performance metrics
        3. Provide a clear recommendation on working with them
        4. Highlight any risks or benefits
        """
        
        try:
            agent = get_agent()
            result = agent.run(prompt)
            st.session_state['agent_response'] = result
            st.rerun()
        except Exception as e:
            st.error(f"Error: {str(e)}")

def analyze_prices(material):
    with st.spinner(f"🤖 Analyzing prices for {material}..."):
        prompt = f"""
        Perform price analysis for '{material}':
        
        1. Get current market rates from multiple vendors
        2. Compare pricing and identify the best value option
        3. Consider factors like lead time and minimum orders
        4. Provide pricing recommendations
        """
        
        try:
            agent = get_agent()
            result = agent.run(prompt)
            st.session_state['agent_response'] = result
            st.rerun()
        except Exception as e:
            st.error(f"Error: {str(e)}")

def run_custom_query(query):
    with st.spinner("🤖 Processing your custom query..."):
        try:
            agent = get_agent()
            result = agent.run(query)
            st.session_state['agent_response'] = result
            st.rerun()
        except Exception as e:
            st.error(f"Error: {str(e)}")

if __name__ == "__main__":
    main()