import gradio as gr
from transformers import pipeline

# Load lightweight open LLM (Mistral from HuggingFace hub for example)
model = pipeline(
    "text2text-generation", 
    model="google/flan-t5-small"
)

def find_flaws(argument, context):
    """Step 1: Ask LLM to find flaws in the argument, considering the context."""
    prompt = f"""
You are an argument quality analyzer. 
Here is the context in which the argument was made:
{context}

Here is the argument:
{argument}

Analyze the argument in this context and list the main flaws or weaknesses 
(e.g. relevance, missing evidence, logical gaps, clarity issues).
Flaws:
"""
    flaws_text = model(prompt, max_new_tokens=200, do_sample=True)[0]["generated_text"]
    flaws_text = flaws_text.split("Flaws:")[-1].strip()
    return flaws_text

def suggest_fixes(argument, context, flaws):
    """Step 2: Ask LLM to suggest fixes for each flaw, considering the context."""
    prompt = f"""
Context:
{context}

Argument:
{argument}

Flaws:
{flaws}

For each flaw, suggest a concrete fix or improvement that would make the argument 
stronger *within this context* (e.g., add data, clarify reasoning, connect more clearly to the question).
Fixes:
"""
    fixes_text = model(prompt, max_new_tokens=250, do_sample=True)[0]["generated_text"]
    fixes_text = fixes_text.split("Fixes:")[-1].strip()
    fixes_list = [fix.strip(" -") for fix in fixes_text.split("\n") if fix.strip()]
    return fixes_list

def integrate_fix(argument, context, flaw, solution):
    """Step 3: Rewrite the argument by integrating the chosen solution."""
    prompt = f"""
Context:
{context}

Original argument:
{argument}

Flaw:
{flaw}

Proposed solution:
{solution}

Rewrite the argument so it keeps the same overall meaning but integrates the fix 
and is more persuasive, clear, and relevant to the context.
Improved Argument:
"""
    improved = model(prompt, max_new_tokens=300, do_sample=True)[0]["generated_text"]
    improved = improved.split("Improved Argument:")[-1].strip()
    return improved

# Gradio workflow with context support
with gr.Blocks() as demo:
    gr.Markdown("# 🧭 Interactive Argument Advisor (with Context)")

    with gr.Row():
        with gr.Column():
            context_input = gr.Textbox(
                label="Provide Context",
                lines=3,
                placeholder="Describe the question, pro/con side, criteria, or other context that this argument is responding to..."
            )
            argument_input = gr.Textbox(
                label="Enter Your Argument",
                lines=4,
                placeholder="Type or paste your argument here..."
            )

    flaws_output = gr.Textbox(label="Detected Flaws", interactive=False)
    find_flaws_btn = gr.Button("🔍 Find Flaws")

    fixes_dropdown = gr.Dropdown(label="Suggested Fixes", choices=[], interactive=True)
    suggest_btn = gr.Button("💡 Suggest Fixes")

    improved_output = gr.Textbox(label="Improved Argument", lines=4)
    apply_btn = gr.Button("✅ Apply Fix")

    # Step 1: Find flaws
    find_flaws_btn.click(fn=find_flaws, inputs=[argument_input, context_input], outputs=flaws_output)

    # Step 2: Suggest fixes
    suggest_btn.click(fn=suggest_fixes, inputs=[argument_input, context_input, flaws_output], outputs=fixes_dropdown)

    # Step 3: Apply fix
    apply_btn.click(fn=integrate_fix, inputs=[argument_input, context_input, flaws_output, fixes_dropdown], outputs=improved_output)

demo.launch(share=True)