import gradio as gr
import json
import time
import os
from groq import Groq
from jsonschema import validate, ValidationError

MAX_RETRIES = 3

# Initialize Groq client
client = Groq(api_key=os.getenv("GROQ_API_KEY"))


def generate_structured_output(input_text, schema_text, model_name):
    start_time = time.time()
    retries = 0

    try:
        schema = json.loads(schema_text)
    except Exception as e:
        return None, f"❌ Invalid JSON Schema: {str(e)}", 0, 0

    base_prompt = f"""
You are a strict JSON generator.

Follow this JSON schema exactly:
{schema_text}

Rules:
- Return ONLY valid JSON
- No explanation
- No markdown
- No extra text

Input:
{input_text}
"""

    while retries < MAX_RETRIES:
        try:
            response = client.chat.completions.create(
                model=model_name,
                messages=[{"role": "user", "content": base_prompt}],
                temperature=0
            )

            output = response.choices[0].message.content.strip()

            parsed = json.loads(output)
            validate(instance=parsed, schema=schema)

            latency = round(time.time() - start_time, 2)

            return (
                json.dumps(parsed, indent=2),
                "✅ Valid JSON",
                retries,
                latency,
            )

        except (json.JSONDecodeError, ValidationError) as e:
            retries += 1
            error_message = str(e)

            base_prompt = f"""
The previous output was invalid.

Error:
{error_message}

Fix the JSON to strictly match this schema:
{schema_text}

Return ONLY corrected JSON.
"""

    latency = round(time.time() - start_time, 2)

    return None, f"❌ Failed after {MAX_RETRIES} retries", retries, latency


with gr.Blocks(title="StructGuard-Groq") as demo:
    gr.Markdown("# 🚀 StructGuard-Groq")
    gr.Markdown("### Production-Grade Schema-First LLM Extractor (Powered by Groq)")

    with gr.Row():
        model_name = gr.Dropdown(
            [
                "llama3-70b-8192",
                "mixtral-8x7b-32768",
                "llama3-8b-8192"
            ],
            value="llama3-70b-8192",
            label="Groq Model"
        )

    input_text = gr.Textbox(
        lines=8,
        label="Input Text",
        placeholder="Paste text you want to structure..."
    )

    schema_text = gr.Textbox(
        lines=8,
        label="JSON Schema",
        placeholder='Paste valid JSON Schema here...'
    )

    generate_btn = gr.Button("Generate Structured JSON")

    output_json = gr.Code(label="Validated JSON Output", language="json")
    status = gr.Textbox(label="Validation Status")
    retries_box = gr.Number(label="Retries Used")
    latency_box = gr.Number(label="Latency (seconds)")

    generate_btn.click(
        generate_structured_output,
        inputs=[input_text, schema_text, model_name],
        outputs=[output_json, status, retries_box, latency_box],
    )

demo.launch()