import json
import textwrap
from typing import Dict, Any, List, Tuple

import gradio as gr
import requests
import matplotlib.pyplot as plt
from matplotlib.figure import Figure


# ============================================================
#  LLM CALLER (GPT-4.1 BY DEFAULT)
# ============================================================

def call_chat_completion(
    api_key: str,
    base_url: str,
    model: str,
    system_prompt: str,
    user_prompt: str,
    max_completion_tokens: int = 2000,
) -> str:
    """
    OpenAI-compatible chat completion call.

    - Uses new-style `max_completion_tokens` (for GPT-4.1, GPT-4o, etc.)
    - Falls back to `max_tokens` if the provider doesn't support it.
    - No temperature / top_p to avoid incompatibility with some models.
    """
    if not api_key:
        raise ValueError("API key is required.")

    if not base_url:
        base_url = "https://api.openai.com"

    url = base_url.rstrip("/") + "/v1/chat/completions"

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        "max_completion_tokens": max_completion_tokens,
    }

    resp = requests.post(url, headers=headers, json=payload, timeout=60)

    # Fallback for providers still expecting `max_tokens`
    if resp.status_code == 400 and "max_completion_tokens" in resp.text:
        payload.pop("max_completion_tokens", None)
        payload["max_tokens"] = max_completion_tokens
        resp = requests.post(url, headers=headers, json=payload, timeout=60)

    if resp.status_code != 200:
        raise RuntimeError(
            f"LLM API error: {resp.status_code} - {resp.text[:400]}"
        )

    data = resp.json()
    try:
        return data["choices"][0]["message"]["content"]
    except Exception as e:
        raise RuntimeError(f"Unexpected LLM response format: {e}\n\n{json.dumps(data, indent=2)}") from e


# ============================================================
#  SOP PROMPT + PARSING
# ============================================================

SOP_SYSTEM_PROMPT = """
You are an expert process engineer. Produce SOPs strictly as JSON with this schema:

{
  "title": "string",
  "purpose": "string",
  "scope": "string",
  "definitions": ["string", ...],
  "roles": [
    {
      "name": "string",
      "responsibilities": ["string", ...]
    }
  ],
  "prerequisites": ["string", ...],
  "steps": [
    {
      "step_number": 1,
      "title": "string",
      "description": "string",
      "owner_role": "string",
      "inputs": ["string", ...],
      "outputs": ["string", ...]
    }
  ],
  "escalation": ["string", ...],
  "metrics": ["string", ...],
  "risks": ["string", ...],
  "versioning": {
    "version": "1.0",
    "owner": "string",
    "last_updated": "string"
  }
}

Return ONLY JSON. No explanation or commentary.
"""

def build_user_prompt(
    sop_title: str,
    description: str,
    industry: str,
    tone: str,
    detail_level: str,
) -> str:
    return f"""
SOP Title: {sop_title or "Untitled SOP"}
Context: {description or "N/A"}
Industry: {industry or "General"}
Tone: {tone or "Professional"}
Detail Level: {detail_level or "Standard"}
Audience: mid-career professionals who need clarity and accountability.
""".strip()


def parse_sop_json(raw_text: str) -> Dict[str, Any]:
    """Extract JSON from LLM output, stripping code fences if present."""
    txt = raw_text.strip()

    if txt.startswith("```"):
        parts = txt.split("```")
        # choose the first part that looks like JSON
        txt = next((p for p in parts if "{" in p and "}" in p), parts[-1])

    first = txt.find("{")
    last = txt.rfind("}")
    if first == -1 or last == -1:
        raise ValueError("No JSON object detected in model output.")
    txt = txt[first:last + 1]

    return json.loads(txt)


def sop_to_markdown(sop: Dict[str, Any]) -> str:
    """Render SOP JSON → readable Markdown document."""

    def bullet(items):
        if not items:
            return "_None provided._"
        return "\n".join(f"- {i}" for i in items)

    md = []

    md.append(f"# {sop.get('title', 'Standard Operating Procedure')}\n")

    md.append("## 1. Purpose")
    md.append(sop.get("purpose", "N/A"))

    md.append("\n## 2. Scope")
    md.append(sop.get("scope", "N/A"))

    md.append("\n## 3. Definitions")
    md.append(bullet(sop.get("definitions", [])))

    md.append("\n## 4. Roles & Responsibilities")
    for role in sop.get("roles", []):
        md.append(f"### {role.get('name', 'Role')}")
        md.append(bullet(role.get("responsibilities", [])))

    md.append("\n## 5. Prerequisites")
    md.append(bullet(sop.get("prerequisites", [])))

    md.append("\n## 6. Procedure (Step-by-Step)")
    for step in sop.get("steps", []):
        md.append(f"### Step {step.get('step_number', '?')}: {step.get('title', 'Step')}")
        md.append(f"**Owner:** {step.get('owner_role', 'N/A')}")
        md.append(step.get("description", ""))
        md.append("**Inputs:**")
        md.append(bullet(step.get("inputs", [])))
        md.append("**Outputs:**")
        md.append(bullet(step.get("outputs", [])))

    md.append("\n## 7. Escalation")
    md.append(bullet(sop.get("escalation", [])))

    md.append("\n## 8. Metrics")
    md.append(bullet(sop.get("metrics", [])))

    md.append("\n## 9. Risks")
    md.append(bullet(sop.get("risks", [])))

    v = sop.get("versioning", {})
    md.append("\n## 10. Version Control")
    md.append(f"- Version: {v.get('version', '1.0')}")
    md.append(f"- Owner: {v.get('owner', 'N/A')}")
    md.append(f"- Last Updated: {v.get('last_updated', 'N/A')}")

    return "\n\n".join(md)


# ============================================================
#  IMPROVED DIAGRAM — AUTO-SIZED CARDS, NO OVERFLOW
# ============================================================

def create_sop_steps_figure(sop: Dict[str, Any]) -> Figure:
    """
    Draw each step as a stacked card with:
    - dynamic height based on description length
    - number block on the left
    - title + owner + wrapped description inside card
    """

    steps = sop.get("steps", [])
    if not steps:
        fig, ax = plt.subplots(figsize=(7, 2))
        ax.text(0.5, 0.5, "No steps available to visualize.", ha="center", va="center")
        ax.axis("off")
        fig.tight_layout()
        return fig

    # First pass: determine required height for each card
    card_heights = []
    total_height = 0.0

    for step in steps:
        desc_lines = textwrap.wrap(step.get("description", ""), width=70)
        # base height (title + owner) + 0.3 per line of description
        base = 1.0  # title + owner + padding
        per_line = 0.32
        h = base + per_line * max(len(desc_lines), 1)
        h += 0.3  # bottom padding
        card_heights.append(h)
        total_height += h

    # Add spacing between cards
    spacing = 0.4
    total_height += spacing * (len(steps) + 1)

    fig_height = min(20, max(5, total_height))
    fig, ax = plt.subplots(figsize=(10, fig_height))
    ax.set_xlim(0, 1)
    ax.set_ylim(0, total_height)

    y = total_height - spacing  # start from top

    for step, h in zip(steps, card_heights):
        y_bottom = y - h
        y_top = y

        # Card boundaries
        x0 = 0.05
        x1 = 0.95

        # Draw outer card
        ax.add_patch(
            plt.Rectangle(
                (x0, y_bottom),
                x1 - x0,
                h,
                fill=False,
                linewidth=1.8,
            )
        )

        # Number block
        num_block_w = 0.08
        ax.add_patch(
            plt.Rectangle(
                (x0, y_bottom),
                num_block_w,
                h,
                fill=False,
                linewidth=1.6,
            )
        )

        # Step number text in the center of the number block
        ax.text(
            x0 + num_block_w / 2,
            y_bottom + h / 2,
            str(step.get("step_number", "?")),
            ha="center",
            va="center",
            fontsize=13,
            fontweight="bold",
        )

        # Text area start
        text_x = x0 + num_block_w + 0.02

        # Title
        ax.text(
            text_x,
            y_top - 0.25,
            step.get("title", ""),
            ha="left",
            va="top",
            fontsize=12,
            fontweight="bold",
        )

        # Owner
        owner = step.get("owner_role", "")
        if owner:
            owner_y = y_top - 0.55
            ax.text(
                text_x,
                owner_y,
                f"Owner: {owner}",
                ha="left",
                va="top",
                fontsize=10,
                style="italic",
            )
        else:
            owner_y = y_top - 0.5

        # Description (wrapped)
        desc_lines = textwrap.wrap(step.get("description", ""), width=70)
        desc_y = owner_y - 0.4
        for line in desc_lines:
            ax.text(
                text_x,
                desc_y,
                line,
                ha="left",
                va="top",
                fontsize=9,
            )
            desc_y -= 0.3  # vertical spacing per line

        y = y_bottom - spacing  # move down for next card

    ax.axis("off")
    fig.tight_layout()
    return fig


# ============================================================
#  SAMPLE SCENARIOS
# ============================================================

SAMPLE_SOPS: Dict[str, Dict[str, str]] = {
    "Volunteer Onboarding": {
        "title": "Volunteer Onboarding",
        "description": "Onboard new volunteers including application review, background checks, orientation, training, and site placement.",
        "industry": "Nonprofit / Youth Development",
    },
    "Remote Employee Onboarding": {
        "title": "Remote Employee Onboarding",
        "description": "Design a remote onboarding SOP for hybrid employees including IT setup, HR paperwork, and culture onboarding.",
        "industry": "HR / General",
    },
    "IT Outage Response": {
        "title": "IT Outage Incident Response",
        "description": "Major outage response SOP including detection, triage, escalation, communication, restoration, and post-mortem.",
        "industry": "IT / Operations",
    },
}

def load_sample(sample_name: str) -> Tuple[str, str, str]:
    if not sample_name or sample_name not in SAMPLE_SOPS:
        return "", "", "General"
    s = SAMPLE_SOPS[sample_name]
    return s["title"], s["description"], s["industry"]


# ============================================================
#  MAIN HANDLER FOR GRADIO
# ============================================================

def generate_sop_ui(
    api_key_state: str,
    api_key_input: str,
    base_url: str,
    model_name: str,
    sop_title: str,
    description: str,
    industry: str,
    tone: str,
    detail_level: str,
) -> Tuple[str, str, Figure, str]:

    api_key = api_key_input or api_key_state
    if not api_key:
        return (
            "⚠️ Please enter your API key in the left panel.",
            "",
            create_sop_steps_figure({"steps": []}),
            api_key_state,
        )

    model = model_name or "gpt-4.1"

    user_prompt = build_user_prompt(sop_title, description, industry, tone, detail_level)

    try:
        raw = call_chat_completion(
            api_key=api_key,
            base_url=base_url,
            model=model,
            system_prompt=SOP_SYSTEM_PROMPT,
            user_prompt=user_prompt,
            max_completion_tokens=2000,
        )

        sop = parse_sop_json(raw)
        md = sop_to_markdown(sop)
        fig = create_sop_steps_figure(sop)
        json_out = json.dumps(sop, indent=2, ensure_ascii=False)

        return md, json_out, fig, api_key  # persist key in session state

    except Exception as e:
        return (
            f"❌ Error generating SOP:\n\n{e}",
            "",
            create_sop_steps_figure({"steps": []}),
            api_key_state,
        )


# ============================================================
#  GRADIO UI
# ============================================================

with gr.Blocks(title="ZEN Simple SOP Builder") as demo:
    gr.Markdown(
        """
# 🧭 ZEN Simple SOP Builder

Generate clean, professional Standard Operating Procedures (SOPs) from a short description,  
plus an auto-generated visual diagram of the steps.

Powered by your own API key (GPT-4.1 by default).
"""
    )

    api_key_state = gr.State("")

    with gr.Row():
        # LEFT COLUMN — API + Samples
        with gr.Column(scale=1):
            gr.Markdown("### Step 1 — API & Model Settings")

            api_key_input = gr.Textbox(
                label="LLM API Key",
                placeholder="Enter your OpenAI (or compatible) API key",
                type="password",
            )

            base_url = gr.Textbox(
                label="Base URL",
                value="https://api.openai.com",
                placeholder="e.g. https://api.openai.com or custom OpenAI-compatible endpoint",
            )

            model_name = gr.Textbox(
                label="Model Name",
                value="gpt-4.1",
                placeholder="e.g. gpt-4.1, gpt-4o, etc.",
            )

            gr.Markdown("### Load a Sample SOP")

            sample_dropdown = gr.Dropdown(
                label="Sample scenarios",
                choices=list(SAMPLE_SOPS.keys()),
                value=None,
                info="Optional: load a ready-made example to test the tool.",
            )

            load_button = gr.Button("Load Sample into Form")

        # RIGHT COLUMN — SOP Description
        with gr.Column(scale=2):
            gr.Markdown("### Step 2 — Describe the SOP")

            sop_title = gr.Textbox(
                label="SOP Title",
                placeholder="e.g. Volunteer Onboarding Workflow",
            )

            description = gr.Textbox(
                label="Describe the process / context",
                placeholder="What should this SOP cover? Who is it for? Any constraints?",
                lines=6,
            )

            industry = gr.Textbox(
                label="Industry / Domain",
                value="General",
                placeholder="e.g. Nonprofit, HR, Education, Healthcare, IT",
            )

            tone = gr.Dropdown(
                label="Tone",
                choices=["Professional", "Executive", "Supportive", "Direct", "Compliance-focused"],
                value="Professional",
            )

            detail_level = gr.Dropdown(
                label="Detail Level",
                choices=["Standard", "High detail", "Checklist-style", "Overview only"],
                value="Standard",
            )

            generate_button = gr.Button("🚀 Generate SOP", variant="primary")

    gr.Markdown("### Step 3 — Generated SOP")

    with gr.Row():
        with gr.Column(scale=3):
            sop_output = gr.Markdown(
                label="SOP (Markdown)",
                value="Your SOP will appear here after generation.",
            )
        with gr.Column(scale=2):
            sop_json_output = gr.Code(
                label="Raw SOP JSON (for automation / export)",
                language="json",
            )

    gr.Markdown("### Step 4 — Visual Workflow Diagram")
    sop_figure = gr.Plot(label="SOP Steps Diagram")

    # Wire up actions
    load_button.click(
        fn=load_sample,
        inputs=[sample_dropdown],
        outputs=[sop_title, description, industry],
    )

    generate_button.click(
        fn=generate_sop_ui,
        inputs=[
            api_key_state,
            api_key_input,
            base_url,
            model_name,
            sop_title,
            description,
            industry,
            tone,
            detail_level,
        ],
        outputs=[sop_output, sop_json_output, sop_figure, api_key_state],
    )


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