import os
import re
import gradio as gr
from groq import Groq

GROQ_API_KEY = os.environ.get("GROQ_API_KEY")

def strip_code_fences(text):
    text = (text or "").strip()
    m = re.match(r"^```(?:python)?\s*\n?(.*?)\n?```\s*$", text, re.DOTALL | re.IGNORECASE)
    return m.group(1).strip() if m else text

def generate_network_script(user_goal, library, device_type):
    if not GROQ_API_KEY:
        return "# Error\nGROQ_API_KEY is missing. Add it in Space Settings → Variables and secrets."
    goal = (user_goal or "").strip()
    if not goal:
        return "# Error\nPlease describe the task."

    system_prompt = """You are a Senior Network Automation Architect.
Output ONLY raw Python code. No markdown fences. No extra text."""
    user_prompt = f"Library: {library}\nDevice: {device_type}\nTask: {goal}"

    try:
        client = Groq(api_key=GROQ_API_KEY)
        r = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0.1,
        )
        return strip_code_fences(r.choices[0].message.content)
    except Exception as e:
        return f"# Error\n{str(e)}"

custom_css = """
.gradio-container { font-family: 'Inter', sans-serif; }
.header-box { text-align: center; padding: 20px; background: linear-gradient(90deg, #1e3a8a, #3b82f6);
  color: white; border-radius: 10px; margin-bottom: 20px; }
.header-box h1 { margin: 0; font-size: 28px; color: white !important; }
"""

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo"), css=custom_css) as app:
    gr.HTML('<div class="header-box"><h1>⚙️ AutoNet Architect</h1><p>AI script generator</p></div>')
    with gr.Row():
        with gr.Column(scale=1):
            library = gr.Dropdown(["Netmiko", "NAPALM", "Nornir"], value="Netmiko", label="Library")
            device_type = gr.Dropdown(
                ["cisco_ios", "cisco_nxos", "arista_eos", "juniper_junos"],
                value="cisco_ios",
                label="Device OS",
            )
            user_goal = gr.Textbox(lines=5, label="Describe the task")
            generate_btn = gr.Button("✨ Generate Code", variant="primary")
        with gr.Column(scale=2):
            output_code = gr.Code(label="Python Script", language="python", interactive=True)

    generate_btn.click(
        fn=generate_network_script,
        inputs=[user_goal, library, device_type],
        outputs=output_code,
    )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))