import gradio as gr
import pandas as pd
from simple_salesforce import Salesforce
from datetime import datetime
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global connection
sf_connection = None

def salesforce_data_loader(username, password, security_token, sandbox, operation, csv_file):
    """Main function that handles all Salesforce operations"""
    global sf_connection
    
    # Step 1: Connect to Salesforce
    if not username or not password or not security_token:
        return "❌ Please provide username, password, and security token"
    
    try:
        domain = 'test' if sandbox else None
        sf_connection = Salesforce(
            username=username,
            password=password,
            security_token=security_token,
            domain=domain
        )
        
        connection_msg = f"✅ Connected to Salesforce as {username}\n"
        
        # Step 2: Handle file upload if provided
        if csv_file is not None and operation != "connect_only":
            try:
                # Read the file
                if csv_file.name.endswith('.csv'):
                    df = pd.read_csv(csv_file.name)
                elif csv_file.name.endswith(('.xlsx', '.xls')):
                    df = pd.read_excel(csv_file.name)
                else:
                    return connection_msg + "❌ Please upload a CSV or Excel file"
                
                if df.empty:
                    return connection_msg + "❌ The uploaded file is empty"
                
                # Clean data
                records = df.to_dict('records')
                cleaned_records = []
                for record in records:
                    cleaned_record = {k: v for k, v in record.items() if pd.notna(v)}
                    cleaned_records.append(cleaned_record)
                
                # Determine object based on columns (simple heuristic)
                columns = df.columns.str.lower()
                if any(col in columns for col in ['firstname', 'lastname', 'email']):
                    object_name = "Contact"
                elif any(col in columns for col in ['company', 'name']):
                    object_name = "Account"
                else:
                    object_name = "Lead"
                
                # Perform operation using bulk API correctly
                if operation == "insert":
                    result = sf_connection.bulk.__getattr__(object_name).insert(cleaned_records)
                elif operation == "update":
                    result = sf_connection.bulk.__getattr__(object_name).update(cleaned_records)
                else:
                    return connection_msg + "❌ Invalid operation. Use 'insert' or 'update'"
                
                # Process results
                success_count = sum(1 for r in result if r.get('success'))
                error_count = len(result) - success_count
                
                upload_msg = f"\n📤 Upload Results:\n"
                upload_msg += f"Object: {object_name}\n"
                upload_msg += f"Total records: {len(records)}\n"
                upload_msg += f"✅ Successful: {success_count}\n"
                upload_msg += f"❌ Failed: {error_count}\n"
                
                return connection_msg + upload_msg
                
            except Exception as e:
                return connection_msg + f"❌ Upload error: {str(e)}"
        
        # Step 3: Handle export operation
        elif operation == "export":
            try:
                # Export some Account records as example
                query = "SELECT Id, Name, Type, Phone, Website FROM Account LIMIT 100"
                result = sf_connection.query_all(query)
                records = result['records']
                
                if records:
                    df = pd.DataFrame(records)
                    if 'attributes' in df.columns:
                        df = df.drop('attributes', axis=1)
                    
                    export_msg = f"\n📥 Export Results:\n"
                    export_msg += f"Records exported: {len(records)}\n"
                    export_msg += f"Fields: {', '.join(df.columns)}\n"
                    export_msg += f"Sample data:\n{df.head().to_string()}"
                    
                    return connection_msg + export_msg
                else:
                    return connection_msg + "\n❌ No records found to export"
                    
            except Exception as e:
                return connection_msg + f"\n❌ Export error: {str(e)}"
        
        else:
            return connection_msg + "\n💡 Connection successful! Upload a file to insert/update data, or select 'export' to download data."
            
    except Exception as e:
        error_msg = str(e)
        if "INVALID_LOGIN" in error_msg:
            return "❌ Invalid credentials. Please check your username, password, and security token."
        elif "API_DISABLED_FOR_ORG" in error_msg:
            return "❌ API access is disabled. Contact your Salesforce admin."
        elif "LOGIN_MUST_USE_SECURITY_TOKEN" in error_msg:
            return "❌ Security token required. Append it to your password."
        else:
            return f"❌ Connection failed: {error_msg}"

# Create the interface
demo = gr.Interface(
    fn=salesforce_data_loader,
    inputs=[
        gr.Textbox(label="Username", placeholder="your.email@company.com"),
        gr.Textbox(label="Password", type="password"),
        gr.Textbox(label="Security Token", type="password"),
        gr.Checkbox(label="Sandbox Environment"),
        gr.Dropdown(
            label="Operation", 
            choices=["connect_only", "insert", "update", "export"],
            value="connect_only"
        ),
        gr.File(label="CSV/Excel File (optional)", file_types=[".csv", ".xlsx", ".xls"])
    ],
    outputs=gr.Textbox(label="Results", lines=10),
    title="🚀 Salesforce Data Loader",
    description="""
    **Simple Salesforce Data Management Tool**
    
    1. **Connect**: Enter your credentials and select 'connect_only'
    2. **Upload**: Select 'insert' or 'update' and upload a CSV/Excel file  
    3. **Export**: Select 'export' to download Account data
    
    **Note**: For uploads, the tool auto-detects object type based on column names.
    """,
    examples=[
        ["user@company.com", "password123", "token123", False, "connect_only", None],
    ]
)

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