import random
import gradio as gr
from datetime import datetime
from transformers import pipeline

# ---------- Label Mapping (Unchanged) ----------
label_to_issue_type = {
    "LABEL_0": "Performance",
    "LABEL_1": "Error",
    "LABEL_2": "Security",
    "LABEL_3": "Best Practice"
}

suggestions = {
    "Performance": "Consider optimizing loops and database access. Use collections to reduce SOQL queries.",
    "Error": "Add proper error handling and null checks. Use try-catch blocks effectively.",
    "Security": "Avoid dynamic SOQL. Use binding variables to prevent SOQL injection.",
    "Best Practice": "Refactor for readability and use bulk-safe patterns, such as processing records in batches."
}

severities = {
    "Performance": "Medium",
    "Error": "High",
    "Security": "High",
    "Best Practice": "Low"
}

# ---------- Mock Salesforce Knowledge Base ----------
salesforce_knowledge_base = {
    "governor limits soql": "In Salesforce, the governor limit for SOQL queries is 100 per synchronous transaction and 200 per asynchronous transaction (e.g., Batch Apex, Queueable).",
    "governor limits dml": "The governor limit for DML statements is 150 per synchronous or asynchronous transaction.",
    "bulkify apex trigger": """
To bulkify an Apex trigger, ensure it handles multiple records to stay within governor limits:
- Use collections (e.g., Lists, Sets, Maps) to process records.
- Perform SOQL queries outside loops.
- Execute DML operations in bulk.
Example:
```apex
trigger AccountTrigger on Account (before insert) {
    Set<String> accountNames = new Set<String>();
    for (Account acc : Trigger.new) {
        accountNames.add(acc.Name);
    }
    List<Account> existingAccounts = [SELECT Name FROM Account WHERE Name IN :accountNames];
    // Process records
}
```""",
    "soql injection": """
To prevent SOQL injection:
- Use bind variables (e.g., `:variable`) instead of dynamic SOQL.
- Escape single quotes with `String.escapeSingleQuotes()`.
Example:
```apex
String userInput = 'Test';
List<Account> accounts = [SELECT Name FROM Account WHERE Name = :userInput];
```""",
    "lwc best practices": """
LWC best practices include:
- Use `@api` decorators for public properties.
- Leverage `@wire` for efficient data retrieval.
- Avoid hardcoding IDs; use dynamic queries.
- Follow SLDS for consistent UI.
Example:
```javascript
import { LightningElement, api } from 'lwc';
export default class MyComponent extends LightningElement {
    @api recordId;
}
```""",
    "batch apex": """
Batch Apex processes records in chunks to handle large data volumes:
- Implement `Database.Batchable<SObject>`.
- Use `start`, `execute`, and `finish` methods.
Example:
```apex
global class MyBatch implements Database.Batchable<SObject> {
    global Database.QueryLocator start(Database.BatchableContext bc) {
        return Database.getQueryLocator('SELECT Id FROM Account');
    }
    global void execute(Database.BatchableContext bc, List<SObject> scope) {
        // Process records
    }
    global void finish(Database.BatchableContext bc) {}
}
```""",
}

# ---------- Load QnA Model ----------
try:
    qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-large")
except Exception as e:
    print(f"Model loading error: {e}. Falling back to flan-t5-base.")
    qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")

# ---------- Local Logging (Unchanged) ----------
def log_to_console(data, log_type):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] {log_type} Log: {data}")

# ---------- Code Analyzer (Unchanged) ----------
def analyze_code(code):
    if not code.strip():
        return "No code provided.", "", ""

    label = random.choice(list(label_to_issue_type.keys()))
    issue_type = label_to_issue_type[label]
    suggestion = suggestions[issue_type]
    severity = severities[issue_type]

    log_data = {
        "Name": f"Review_{issue_type}",
        "CodeSnippet": code,
        "IssueType": issue_type,
        "Suggestion": suggestion,
        "Severity": severity
    }
    log_to_console(log_data, "Code Review")

    return issue_type, suggestion, severity

# ---------- Metadata Validator (Unchanged) ----------
def validate_metadata(metadata):
    if not metadata.strip():
        return "No metadata provided.", "", ""

    mtype = "Field"
    issue = "Unused field detected"
    recommendation = "Remove it to improve performance or document its purpose."

    log_data = {
        "Name": f"MetadataLog_{mtype}",
        "MetadataType": mtype,
        "IssueDescription": issue,
        "Recommendation": recommendation,
        "Status": "Open"
    }
    log_to_console(log_data, "Metadata Validation")

    return mtype, issue, recommendation

# ---------- Advanced Salesforce Chatbot ----------
conversation_history = []

def salesforce_chatbot(query, history=[]):
    global conversation_history
    if not query.strip():
        return "Please provide a valid Salesforce-related question."

    # Expanded Salesforce keywords
    salesforce_keywords = [
        "apex", "soql", "sosl", "trigger", "lwc", "aura", "visualforce", "salesforce",
        "governor limits", "bulkification", "dml", "query", "metadata", "flow",
        "process builder", "sobject", "schema", "lightning", "custom object",
        "validation rule", "workflow", "platform event", "batch apex", "queueable",
        "future method", "lightning web component", "api", "rest", "soap", "integration",
        "trigger", "profile", "permission set", "sharing rule", "field", "record type",
        "crm", "sfdc", "force.com"
    ]

    # Check if query is Salesforce-related
    if not any(keyword.lower() in query.lower() for keyword in salesforce_keywords):
        return "Please ask a Salesforce-related question (e.g., about Apex, SOQL, LWC, or Salesforce platform features)."

    # Check knowledge base for exact matches
    query_key = query.lower().strip()
    for kb_key, kb_answer in salesforce_knowledge_base.items():
        if kb_key in query_key:
            # Store in conversation history
            conversation_history.append((query, kb_answer))
            conversation_history = conversation_history[-6:]  # Keep last 6 exchanges
            log_to_console({"Question": query, "Answer": kb_answer}, "Chatbot Query")
            return kb_answer

    # Build conversation context
    history_summary = "\n".join([f"User: {q}\nAssistant: {a}" for q, a in conversation_history[-4:]])
    prompt = f"""
You are an expert Salesforce developer with deep knowledge of Apex, SOQL, SOSL, LWC, Aura, Visualforce, and Salesforce platform features. Your role is to provide 100% accurate answers based strictly on Salesforce official documentation and best practices (e.g., governor limits: 100 SOQL queries, 150 DML statements per transaction). Follow these guidelines:
- Provide concise, accurate answers with no speculation.
- Include code snippets in ```apex``` or ```javascript``` blocks for technical questions.
- Reference governor limits or best practices explicitly.
- If the question is ambiguous, ask for clarification within the response.
- Use bullet points for lists or steps.
- Leverage conversation history for context.
- If unsure, admit limitations and suggest checking Salesforce documentation or Trailhead.

Conversation History:
{history_summary}

Question: {query.strip()}
Answer:
"""

    try:
        # Generate response with strict parameters
        result = qa_pipeline(prompt, max_new_tokens=1024, do_sample=False, temperature=0.1, top_k=50)
        output = result[0]["generated_text"].strip()

        # Clean up response
        if output.startswith("Answer:"):
            output = output[7:].strip()

        # Validate response quality
        if len(output) < 20 or output.lower() in ["unknown", "i don't know", "not sure"]:
            output = f"I'm sorry, I couldn't find a precise answer for '{query}'. Please clarify or refer to Salesforce documentation at https://developer.salesforce.com/docs."

        # Store in conversation history
        conversation_history.append((query, output))
        conversation_history = conversation_history[-6:]  # Keep last 6 exchanges

        # Log question and answer
        log_to_console({"Question": query, "Answer": output}, "Chatbot Query")

        return output
    except Exception as e:
        error_msg = f"⚠️ Error generating response: {str(e)}. Please try rephrasing your question or check Salesforce documentation at https://developer.salesforce.com/docs."
        log_to_console({"Question": query, "Error": error_msg}, "Chatbot Error")
        return error_msg

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

    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])

    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])

    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")
        
        # State to manage chat history
        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]
        )

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