import gradio as gr
import os
import re
import json
from mistralai import Mistral

# -----------------------------
# Knowledge Base
# -----------------------------

ABB_ERROR_CODES = {
    "E-458": {
        "description": "Servo Thermal Overload",
        "category": "Thermal Protection",
        "severity_base": 7
    },
    "E-302": {
        "description": "Encoder Signal Loss",
        "category": "Sensor Failure",
        "severity_base": 8
    }
}

with open("documentation_store.json", "r") as f:
    DOC_STORE = json.load(f)

# -----------------------------
# Utility Functions
# -----------------------------

def extract_error(log):
    match = re.search(r"E-\d+", log)
    return match.group(0) if match else None

def retrieve_documentation(error_code):
    if error_code and error_code in DOC_STORE:
        doc = DOC_STORE[error_code]
        return f"""
Manufacturer Documentation Reference:
Title: {doc['title']}
Description: {doc['description']}
Source: {doc['source']}
"""
    return ""

def compute_risk(log_input, error_code):
    score = 0
    text = log_input.lower()

    if "stopped" in text:
        score += 25

    if "temperature reached" in text or "temperature spike" in text:
        score += 35

    if error_code and error_code in ABB_ERROR_CODES:
        score += ABB_ERROR_CODES[error_code]["severity_base"] * 4

    return min(score, 100)

# -----------------------------
# Model Initialization
# -----------------------------

api_key = os.getenv("MISTRAL_API_KEY")
client = Mistral(api_key=api_key)

# -----------------------------
# Core Analysis Function
# -----------------------------

def analyze_fault(log_input):

    if not log_input.strip():
        return "Please provide a factory log."

    error_code = extract_error(log_input)
    documentation_context = retrieve_documentation(error_code)

    try:
        prompt = f"""
You are an industrial automation expert working in an ABB Smart Factory.

Use Kepner-Tregoe structured problem analysis.

If manufacturer documentation is provided below,
you MUST explicitly reference it when determining root cause.
Do not ignore it.

{documentation_context}

Factory Log:
---
{log_input}
---

Return structured JSON only.
"""
        response = client.chat.complete(
            model="mistral-large-latest",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.2,
            max_tokens=600
        )

        ai_output = response.choices[0].message.content
        risk_score = compute_risk(log_input, error_code)

        return f"{ai_output}\n\n---\nRisk Score: {risk_score}/100"

    except Exception as e:
        return f"Error occurred: {str(e)}"

# -----------------------------
# UI Layer
# -----------------------------

with gr.Blocks(title="ABB Smart Factory – KT Fault Intelligence") as demo:
    gr.Markdown("## 🏭 ABB Smart Factory – AI KT Fault Intelligence System")

    log_input = gr.Textbox(label="Factory Log Input", lines=12)
    output = gr.Textbox(label="AI KT Analysis Output", lines=20)

    analyze_button = gr.Button("Analyze Fault")
    analyze_button.click(analyze_fault, inputs=log_input, outputs=output)

# -----------------------------
# Launch
# -----------------------------

demo.launch(server_name="0.0.0.0", server_port=7860)