import gradio as gr
from transformers import pipeline

# Load Hugging Face models
code_analyzer = pipeline("text-classification", model="microsoft/codebert-base")
nlp_model = pipeline("text2text-generation", model="google/flan-t5-large")

# --------------------------
# Code Review Function
# --------------------------
def analyze_code(code):
    if not code.strip():
        return "No code provided.", "", ""
    result = code_analyzer(code)
    return result[0]["label"], "Consider refactoring for better performance", "Medium"

# --------------------------
# Metadata Validator (Mock)
# --------------------------
def validate_metadata(metadata):
    if not metadata.strip():
        return "No metadata provided.", "", ""
    return "Field", "Unused field detected", "Remove it to improve performance"

# --------------------------
# AI Q&A Generator (No fallback)
# --------------------------
def process_nlp_query(query):
    if not query.strip():
        return "No query provided."

    prompt = f"""You are a certified Salesforce Apex expert. Answer this question clearly and accurately:\n\nQuestion: {query}\n\nAnswer:"""

    result = nlp_model(
        prompt,
        max_length=256,
        temperature=0.7,
        top_k=50,
        top_p=0.9,
        repetition_penalty=1.3,
        do_sample=True
    )

    output = result[0]["generated_text"]

    if "Answer:" in output:
        output = output.split("Answer:")[-1]

    lines = output.strip().splitlines()
    seen = set()
    unique_lines = [line.strip() for line in lines if line.strip() not in seen and not seen.add(line.strip())]

    return "\n".join(unique_lines).strip()

# --------------------------
# Gradio UI
# --------------------------
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 Salesforce AI Code Review & Metadata Assistant")

    with gr.Tab("Code Review"):
        code_input = gr.Textbox(label="Apex / LWC Code", lines=8)
        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)
        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("Ask AI (Natural Language)"):
        query_input = gr.Textbox(label="Your question", lines=2, placeholder="e.g. What is a governor limit in Apex?")
        response_output = gr.Textbox(label="AI Response", lines=8)
        nlp_button = gr.Button("Ask")
        nlp_button.click(process_nlp_query, inputs=query_input, outputs=response_output)

demo.launch()