import gradio as gr
import torch
import json
from pathlib import Path
from model import (
    GreesyGPT, 
    generate_moderation, 
    ReasoningMode, 
    OutputFormat, 
    DEVICE,
    describe_reasoning_modes
)

# 1. Initialize Model
model = GreesyGPT()
weights_path = Path("greesy_gpt.pt")

if weights_path.exists():
    model.load_state_dict(torch.load(weights_path, map_location=DEVICE))
    print(f"Loaded weights from {weights_path}")
else:
    print("Warning: No trained weights found. Using fresh initialization.")

model.to(DEVICE)
model.eval()

def moderate(text, mode_str, format_str):
    if not text.strip():
        return "Please enter some text to analyze.", ""

    # Map strings from UI to Enums
    mode = ReasoningMode(mode_str.lower())
    fmt = OutputFormat(format_str.lower())

    # Generate
    result = generate_moderation(
        model, 
        prompt=text, 
        mode=mode, 
        output_format=fmt
    )

    verdict_output = result["verdict_fmt"]
    # If JSON format is selected, prettify it for the textbox
    if fmt == OutputFormat.JSON:
        verdict_output = json.dumps(verdict_output, indent=2)
    
    thinking_process = result.get("thinking", "No reasoning generated.")
    
    return verdict_output, thinking_process

# 2. Build Gradio UI
theme = gr.themes.Soft(
    primary_hue="orange",
    secondary_hue="gray",
)

with gr.Blocks(theme=theme, title="GreesyGPT Content Moderation") as demo:
    gr.Markdown("# 🛡️ GreesyGPT Content Moderation")
    gr.Markdown(
        "A reasoning-based safety model that analyzes content for violations "
        "using chain-of-thought deliberation."
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            input_text = gr.Textbox(
                label="Message to Review",
                placeholder="Type the message you want to moderate here...",
                lines=5
            )
            
            with gr.Row():
                mode_dropdown = gr.Dropdown(
                    choices=[m.value for m in ReasoningMode],
                    value="low",
                    label="Reasoning Mode",
                    info="Higher modes are more thorough but slower."
                )
                format_dropdown = gr.Dropdown(
                    choices=[f.value for f in OutputFormat],
                    value="markdown",
                    label="Output Format"
                )
            
            submit_btn = gr.Button("Analyze Content", variant="primary")

        with gr.Column(scale=3):
            output_verdict = gr.Markdown(label="Verdict")
            
            with gr.Accordion("View Internal Reasoning (Thinking Process)", open=False):
                output_thinking = gr.Textbox(
                    label="Chain of Thought",
                    interactive=False,
                    lines=10
                )

    gr.Examples(
        examples=[
            ["You're so stupid, nobody likes you.", "medium", "markdown"],
            ["How do I fix a bug in my Python code?", "none", "markdown"],
            ["CONGRATULATIONS! You won a $1000 gift card! Click here!", "low", "json"],
            ["I feel really hopeless and don't want to continue.", "high", "markdown"],
        ],
        inputs=[input_text, mode_dropdown, format_dropdown]
    )

    with gr.Expander("System Information"):
        gr.Code(describe_reasoning_modes(), language="text")

    submit_btn.click(
        fn=moderate,
        inputs=[input_text, mode_dropdown, format_dropdown],
        outputs=[output_verdict, output_thinking]
    )

if __name__ == "__main__":
    demo.launch()