import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# 🔁 Replace with your actual Hugging Face model name
MODEL_NAME = "subhoshripal/smolified-context-bridge-slm"

# Load model + tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# ---------- FORMAT OUTPUT ----------
def format_output(text):
    impact, business, action = "N/A", "N/A", "N/A"

    text = text.lower()

    # Try to extract sections flexibly
    if "impact" in text:
        try:
            impact = text.split("impact")[1].split("\n")[0].replace("level", "").replace(":", "").strip().title()
        except:
            pass

    if "business" in text:
        try:
            business = text.split("business")[1].split("strategic")[0].replace("translation", "").replace(":", "").strip().capitalize()
        except:
            pass

    if "strategic" in text or "action" in text:
        try:
            action = text.split("strategic")[-1].replace("action", "").replace(":", "").strip().capitalize()
        except:
            pass

    # FALLBACK: if parsing failed → use whole text smartly
    if business == "N/A":
        business = text[:120].capitalize()

    if action == "N/A":
        action = "Investigate the issue and apply appropriate fixes."

    if impact == "N/A":
        impact = "Moderate"

    return impact, business, action
    
# ---------- MODEL FUNCTION ----------
def analyze(text):
    if not text.strip():
        return "N/A", "Please enter a system log.", "N/A"

    prompt = f"""
You are an Industrial Systems Analyst.

Analyze the following technical log and respond STRICTLY in this format:

Impact Level (Low/Moderate/High/Critical):
Business Translation (1-2 lines, clear and concise):
Strategic Action (specific next step):

Log:
{text}
"""

    inputs = tokenizer(prompt, return_tensors="pt")

    with torch.no_grad():
        outputs = model.generate(
    **inputs,
    max_new_tokens=100,
    temperature=0.2,   # more stable
    top_p=0.9,
    do_sample=True,
    repetition_penalty=1.3
)

    result = tokenizer.decode(
    outputs[0][inputs["input_ids"].shape[-1]:],
    skip_special_tokens=True
    )
    impact, business, action = format_output(result)

    return impact, business, action

# ---------- UI ----------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # ContextBridge: From Logs to Decisions  
    ### Translate technical errors into business insights in seconds
    
    Paste a system log below
    """)

    input_box = gr.Textbox(
        placeholder="e.g. API timeout after 3000ms while fetching user data",
        label="System Log",
        lines=3
    )

    gr.Examples(
        examples=[
            "API timeout after 3000ms while fetching user data",
            "Unauthorized access attempt from IP 192.168.1.5",
            "Database connection refused after multiple retries"
        ],
        inputs=input_box
    )

    analyze_btn = gr.Button("Analyze ⚡")

    impact_output = gr.Textbox(label="Impact Level")
    business_output = gr.Textbox(label="Business Translation")
    action_output = gr.Textbox(label="Strategic Action")

    analyze_btn.click(
        analyze,
        inputs=input_box,
        outputs=[impact_output, business_output, action_output]
    )

demo.launch()