# [All your existing imports]
import random
import gradio as gr
from datetime import datetime
from transformers import pipeline
from simple_salesforce import Salesforce, SalesforceLogin, SFType

# [Your existing mappings, knowledge base, and model loading code remain unchanged]

# ---------- Salesforce Connection Setup ----------
# Replace these with your actual Salesforce credentials or use environment variables for security
SF_USERNAME = "your_username@example.com"
SF_PASSWORD = "your_password"
SF_SECURITY_TOKEN = "your_security_token"

try:
    session_id, instance = SalesforceLogin(
        username=SF_USERNAME,
        password=SF_PASSWORD,
        security_token=SF_SECURITY_TOKEN
    )
    sf = Salesforce(instance=instance, session_id=session_id)
    print("✅ Connected to Salesforce successfully")
except Exception as e:
    sf = None
    print(f"❌ Failed to connect to Salesforce: {e}")

# ---------- Function to Create Salesforce Record ----------
def create_salesforce_record(object_type, field_name, field_value):
    if not sf:
        return "Salesforce connection not established."

    try:
        sobject = SFType(object_type, sf.session_id, sf.sf_instance)
        result = sobject.create({field_name: field_value})
        if result.get("success"):
            return f"✅ Record created successfully in {object_type} with ID: {result['id']}"
        else:
            return f"❌ Failed to create record: {result}"
    except Exception as e:
        return f"⚠️ Error: {str(e)}"

# [Your existing functions: analyze_code, validate_metadata, salesforce_chatbot, etc. remain unchanged]

# ---------- Gradio UI ----------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🤖 Advanced Salesforce AI Code Review & Chatbot")

    # Tab 1: Code Review
    with gr.Tab("Code Review"):
        code_input = gr.Textbox(label="Apex / LWC Code", lines=8, placeholder="Enter your Apex or LWC code here")
        issue_type = gr.Textbox(label="Issue Type")
        suggestion = gr.Textbox(label="AI Suggestion")
        severity = gr.Textbox(label="Severity")
        code_button = gr.Button("Analyze Code")
        code_button.click(analyze_code, inputs=code_input, outputs=[issue_type, suggestion, severity])

    # Tab 2: Metadata Validation
    with gr.Tab("Metadata Validation"):
        metadata_input = gr.Textbox(label="Metadata XML", lines=8, placeholder="Enter your metadata XML here")
        mtype = gr.Textbox(label="Type")
        issue = gr.Textbox(label="Issue")
        recommendation = gr.Textbox(label="Recommendation")
        metadata_button = gr.Button("Validate Metadata")
        metadata_button.click(validate_metadata, inputs=metadata_input, outputs=[mtype, issue, recommendation])

    # Tab 3: Salesforce Chatbot
    with gr.Tab("Salesforce Chatbot"):
        gr.Markdown("### Ask a Salesforce Question\nGet expert answers on Apex, SOQL, LWC, and more!")
        chatbot_output = gr.Chatbot(label="Conversation History", height=400)
        query_input = gr.Textbox(label="Your Question", placeholder="e.g., How do I bulkify an Apex trigger?")
        with gr.Row():
            chatbot_button = gr.Button("Ask")
            clear_button = gr.Button("Clear Chat")

        chat_state = gr.State(value=[])

        def update_chatbot(query, chat_history):
            if not query.strip():
                return chat_history, "Please enter a valid question."
            response = salesforce_chatbot(query, chat_history)
            chat_history.append((query, response))
            return chat_history, ""

        def clear_chat():
            global conversation_history
            conversation_history = []
            return [], ""

        chatbot_button.click(fn=update_chatbot, inputs=[query_input, chat_state], outputs=[chatbot_output, query_input])
        clear_button.click(fn=clear_chat, inputs=None, outputs=[chatbot_output, query_input])

    # Tab 4: Create Salesforce Record
    with gr.Tab("Create Salesforce Record"):
        gr.Markdown("### 📇 Create a New Salesforce Record")
        object_type = gr.Textbox(label="Salesforce Object API Name", placeholder="e.g., Account")
        field_name = gr.Textbox(label="Field API Name", placeholder="e.g., Name")
        field_value = gr.Textbox(label="Field Value", placeholder="e.g., Test Company")
        create_button = gr.Button("Create Record")
        record_result = gr.Textbox(label="Result")

        create_button.click(create_salesforce_record, inputs=[object_type, field_name, field_value], outputs=record_result)

# ---------- Start UI ----------
if __name__ == "__main__":
    demo.launch()