app.py

"""
app.py — CircuitSense Multimodal Inspection System
MLS-1 | Multimodal Agentic AI (v4)

Changes from v1:
  - Fix 1: Ground truth removed from Vision Agent prompt
  - Fix 2: AgentView isolation enforced via _carry_pipeline_fields()
  - Point 1: Category-specific visual cues in Vision Agent prompt
  - Supervisor: confidence check fires BEFORE defect_observed check

Deployment:
  Local:         streamlit run app.py
  Hugging Face:  Push to HF Space with requirements.txt
                 Set OPENAI_API_KEY as a Repository Secret in HF Spaces Settings

Run order:
  1. Run circuitsense_inspection.ipynb  →  creates df_enriched.csv
  2. streamlit run app.py
"""

import os
import json
import re
import base64
import datetime
from io import BytesIO
from pathlib import Path
from dataclasses import dataclass
from typing import TypedDict, List, Dict, Any

import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
from PIL import Image
from openai import OpenAI
from langgraph.graph import StateGraph, END

# ─── PAGE CONFIG ──────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="CircuitSense — AI Inspection",
    page_icon="🔬",
    layout="wide",
    initial_sidebar_state="expanded"
)

st.markdown("""
<style>
.main-header {
    background: linear-gradient(135deg, #0d1b2a 0%, #1b4332 100%);
    color: white; padding: 20px 30px; border-radius: 10px; margin-bottom: 20px;
}
.badge-pass     { background:#4CAF50; color:white; padding:6px 18px; border-radius:20px; font-weight:bold; font-size:1.1em; }
.badge-rework   { background:#FF9800; color:white; padding:6px 18px; border-radius:20px; font-weight:bold; font-size:1.1em; }
.badge-scrap    { background:#F44336; color:white; padding:6px 18px; border-radius:20px; font-weight:bold; font-size:1.1em; }
.badge-uncertain{ background:#9E9E9E; color:white; padding:6px 18px; border-radius:20px; font-weight:bold; font-size:1.1em; }
.agent-card { background:#f8f9fa; border:1px solid #dee2e6; padding:12px; border-radius:8px; margin:6px 0; }
</style>
""", unsafe_allow_html=True)

# ─── CONSTANTS ────────────────────────────────────────────────────────────────
INSPECTION_POLICIES = {
    "surface": """
CIRCUITSENSE SURFACE DEFECT POLICY — Version 3.1

Disposition Rules:
1. PASS: Defect cosmetic only, no functional impact. Scratch < 2mm on non-contact surfaces.
   Dent < 0.1mm on non-critical surfaces. Condition: Severity = LOW
2. REWORK: Scratch on contact surface. Contamination removable without structural risk.
   Colour spot > tolerance but < 5mm. Condition: Severity = MEDIUM
3. SCRAP: Defect compromises integrity after rework. Contamination of active component surfaces.
   Multiple defects (>=3) on same unit. Condition: Severity = HIGH

Override: Any surface defect on Class A (safety-critical) component -> SCRAP regardless.
""",
    "structural": """
CIRCUITSENSE STRUCTURAL DEFECT POLICY — Version 3.1

Disposition Rules:
1. PASS: No structural defect (handled by PassThrough Agent).
2. REWORK: Hairline crack < 1mm on non-load-bearing surface.
   Missing non-critical passive component. Condition: Severity = LOW, non-critical zone.
3. SCRAP: Any crack > 1mm or on load-bearing/connector/seal surface.
   Any burn mark. Missing critical component (IC, connector, power).
   Any short circuit. Hole in substrate. Condition: Severity = MEDIUM or HIGH.

Override: Any structural defect on PCB carrying >5V -> SCRAP.
Any structural defect on pharmaceutical capsule -> SCRAP (patient safety).
""",
    "general": """
CIRCUITSENSE GENERAL INSPECTION POLICY — Version 3.1

Escalation: Vision confidence < 0.60 -> escalate to human inspector.
Audit: Every inspection must produce a complete decision log entry.
SCRAP decisions require secondary confirmation log entry.
"""
}

# ─── CATEGORY-SPECIFIC VISUAL CUES (Point 1) ─────────────────────────────────
# Domain expertise injected into the Vision Agent prompt.
# Describes what each defect type looks like per product — NOT ground truth leakage.
# Balanced instruction works correctly for both normal and defective images.
CATEGORY_VISUAL_CUES = {
    "pcb1": """Category-specific inspection guidance for PCB (pcb1):
  - burn    : Blackened or discoloured traces, scorched substrate, heat damage around components
  - missing : Empty solder pads, unpopulated component footprints, absent ICs or resistors
  - short   : Unintended solder bridges connecting adjacent pins or traces
  - scratch : Linear marks cutting across copper traces or PCB surface coating
  - melt    : Deformed plastic connectors, warped substrate, fused or distorted components

Examine the image carefully. Only report a defect if you can clearly see one of the above.""",

    "capsules": """Category-specific inspection guidance for capsules:
  - scratch : Linear marks or grooves on the smooth capsule shell
  - crack   : Hairline fractures in the casing, especially along the seam or edge
  - leak    : Bubbling, blistering, or discolouration suggesting content seepage
  - dent    : Depressions or flat spots on the otherwise cylindrical surface
  - discolor: Patches of abnormal colour differing from the uniform capsule body

Examine the image carefully. Only report a defect if you can clearly see one of the above.""",

    "cashew": """Category-specific inspection guidance for cashew kernels:
  - colour  : Dark spots, discolouration patches, or abnormal brown/black regions
  - scratch : Surface marks, gouges, or disrupted surface texture
  - hole    : Small cavities or perforations in the kernel surface
  - breakage: Missing chunks, cracked edges, or split kernels
  - contamination: Foreign particles, surface irregularities, or textural anomalies

Examine the image carefully. Only report a defect if you can clearly see one of the above.""",

    "unknown": "Inspect carefully for any visible surface or structural defects. Only report a defect if you can clearly see one."
}

# ─── PIPELINE FIELDS (Fix 2) ──────────────────────────────────────────────────
# Immutable fields set at pipeline entry and carried forward by every node.
# Used by _carry_pipeline_fields() to replace {**state, ...} in node returns.
PIPELINE_FIELDS = ["image_path", "image_b64", "category", "policy_id", "defect_type_gt"]


# ─── UTILITIES ────────────────────────────────────────────────────────────────
def utc_now() -> str:
    return datetime.datetime.now(datetime.timezone.utc).isoformat()


def _carry_pipeline_fields(state: "GlobalState") -> dict:
    """
    Returns immutable pipeline fields from state.
    Every node return is built as:
        {**_carry_pipeline_fields(state), <owned fields>, "decision_log": ...}
    This replaces {**state, ...} and enforces that nodes only write fields they own.
    """
    return {k: state.get(k) for k in PIPELINE_FIELDS if k in state}


def resize_and_encode(image_input, max_size: int = 1024) -> str:
    """Resize and base64-encode an image from file path or PIL Image."""
    if isinstance(image_input, (str, Path)):
        img = Image.open(str(image_input)).convert("RGB")
    else:
        img = image_input.convert("RGB")
    w, h = img.size
    if max(w, h) > max_size:
        ratio = max_size / max(w, h)
        img = img.resize((int(w * ratio), int(h * ratio)), Image.LANCZOS)
    buffer = BytesIO()
    img.save(buffer, format="JPEG", quality=90)
    return base64.standard_b64encode(buffer.getvalue()).decode("utf-8")


def build_vision_message(image_b64: str, text_prompt: str) -> list:
    return [{"role": "user", "content": [
        {"type": "image_url", "image_url": {
            "url": f"data:image/jpeg;base64,{image_b64}", "detail": "high"}},
        {"type": "text", "text": text_prompt}
    ]}]


# ─── DATASET ──────────────────────────────────────────────────────────────────
@st.cache_data
def load_enriched_dataset() -> pd.DataFrame:
    if os.path.exists("df_enriched.csv"):
        return pd.read_csv("df_enriched.csv")
    return pd.DataFrame()


# ─── OPENAI CLIENT ────────────────────────────────────────────────────────────
def get_client():
    api_key  = st.session_state.get("openai_api_key", os.environ.get("OPENAI_API_KEY", ""))
    api_base = st.session_state.get("openai_api_base", "")
    if api_key and api_base:
        return OpenAI(api_key=api_key, base_url=api_base)
    return OpenAI(api_key=api_key)


# ══════════════════════════════════════════════════════════════════════════════
#  LANGGRAPH STATE + NODES
# ══════════════════════════════════════════════════════════════════════════════

class GlobalState(TypedDict, total=False):
    # Pipeline input (immutable)
    image_path:            str
    image_b64:             str
    category:              str
    policy_id:             str
    defect_type_gt:        str   # Ground truth — for post-hoc eval only, NEVER in prompts
    # Vision Agent output
    defect_observed:       bool
    defect_class:          str
    defect_type_observed:  str
    severity:              str
    defect_location:       str
    vision_confidence:     float
    vision_evidence:       str
    # Specialist Agent output
    agent_selected:        str
    specialist_assessment: str
    # Policy / Passthrough Agent output
    disposition:           str
    policy_clause:         str
    policy_justification:  str
    # Audit
    decision_log:          List[dict]
    final_report:          str


# ─── NODE 1: VISION AGENT ────────────────────────────────────────────────────
def vision_agent_node(state: GlobalState) -> GlobalState:
    """
    Fix 1: No ground truth in prompt — pure visual classification.
    Fix 2: Returns only owned fields via _carry_pipeline_fields().
    Point 1: Category-specific visual cues injected as domain guidance.
    """
    oai      = get_client()
    category = state.get("category", "unknown")

    # Point 1: Fetch category-specific visual cues (domain guidance, not GT)
    visual_cues = CATEGORY_VISUAL_CUES.get(category, CATEGORY_VISUAL_CUES["unknown"])

    system_prompt = """You are a precision quality control vision inspector for an electronics manufacturer.
Analyse product images and return structured JSON only — no preamble, no markdown.
You must rely entirely on what you can observe in the image."""

    # Fix 1: No 'Known defect label' line anywhere in this prompt
    user_prompt = f"""Analyse this product image for quality defects.

Product category: {category}

{visual_cues}

Return this exact JSON:
{{
  "defect_observed": true or false,
  "defect_class": "surface" or "structural" or "none",
  "defect_type_observed": "specific defect type you can see",
  "severity": "low" or "medium" or "high",
  "defect_location": "where on the product",
  "vision_confidence": 0.0 to 1.0,
  "vision_evidence": "one sentence of visual evidence"
}}

Classification guide:
  surface    = cosmetic defects: scratches, colour spots, stains, dents, discolouration
  structural = functional defects: cracks, burns, missing components, holes, shorts
  none       = no defect visible

Severity: low=cosmetic only, medium=functional risk possible, high=failure likely.
Set vision_confidence below 0.60 only if the image is too unclear to assess reliably."""

    messages = [{"role": "system", "content": system_prompt}] + \
               build_vision_message(state.get("image_b64", ""), user_prompt)
    try:
        resp = oai.chat.completions.create(
            model="gpt-4o", messages=messages, temperature=0, max_tokens=300)
        raw  = re.sub(r"```json|```", "", resp.choices[0].message.content.strip())
        vo   = json.loads(raw)
    except Exception as e:
        vo = {"defect_observed": True, "defect_class": "surface",
              "defect_type_observed": "unknown", "severity": "medium",
              "defect_location": "undetermined", "vision_confidence": 0.5,
              "vision_evidence": f"API error: {str(e)[:60]}"}

    log = {"timestamp": utc_now(), "node": "VisionAgent",
           "category": category, "output": vo,
           "model": "gpt-4o (vision)", "gt_leak": False}

    # Fix 2: Build return from pipeline fields + owned outputs only
    return {
        **_carry_pipeline_fields(state),
        "defect_observed":      vo.get("defect_observed", True),
        "defect_class":         vo.get("defect_class", "surface"),
        "defect_type_observed": vo.get("defect_type_observed", ""),
        "severity":             vo.get("severity", "medium"),
        "defect_location":      vo.get("defect_location", ""),
        "vision_confidence":    float(vo.get("vision_confidence", 0.5)),
        "vision_evidence":      vo.get("vision_evidence", ""),
        "decision_log":         state.get("decision_log", []) + [log]
    }


# ─── NODE 2: SUPERVISOR AGENT ─────────────────────────────────────────────────
def supervisor_agent_node(state: GlobalState) -> GlobalState:
    """
    Deterministic routing — no LLM call.
    Fix 2: Returns only agent_selected + upstream fields.
    Supervisor fix: confidence check fires BEFORE defect_observed check.
    """
    dc   = state.get("defect_class", "surface")
    conf = state.get("vision_confidence", 0.5)

    # Confidence check FIRST — low confidence escalates regardless of defect_observed
    if conf < 0.60:
        sel    = "uncertain"
        reason = f"low confidence ({conf:.2f} < 0.60) — escalate to human"
    elif not state.get("defect_observed") or dc == "none":
        sel    = "passthrough"
        reason = "no defect detected"
    elif dc == "structural":
        sel    = "structural"
        reason = "defect_class=structural"
    else:
        sel    = "surface"
        reason = "defect_class=surface"

    log = {"timestamp": utc_now(), "node": "SupervisorAgent",
           "output": {"agent_selected": sel}, "routing_reason": reason}

    # Fix 2: Carry all upstream fields explicitly
    return {
        **_carry_pipeline_fields(state),
        "defect_observed":      state.get("defect_observed"),
        "defect_class":         dc,
        "defect_type_observed": state.get("defect_type_observed", ""),
        "severity":             state.get("severity", ""),
        "defect_location":      state.get("defect_location", ""),
        "vision_confidence":    conf,
        "vision_evidence":      state.get("vision_evidence", ""),
        "agent_selected":       sel,
        "decision_log":         state.get("decision_log", []) + [log]
    }


# ─── NODE 3a: SURFACE DEFECT AGENT ───────────────────────────────────────────
def surface_defect_agent_node(state: GlobalState) -> GlobalState:
    """Fix 2: Returns only severity + specialist_assessment + upstream fields."""
    oai           = get_client()
    defect_type   = state.get("defect_type_observed", "")
    severity_in   = state.get("severity", "medium")
    location      = state.get("defect_location", "")
    evidence      = state.get("vision_evidence", "")
    category      = state.get("category", "unknown")

    system_prompt = "You are a surface defect characterisation specialist at CircuitSense."
    user_prompt   = f"""Characterise this surface defect:
Product: {category} | Type: {defect_type}
Severity: {severity_in} | Location: {location}
Evidence: {evidence}

Provide 3-5 sentences on functional impact, rework feasibility, and category-specific considerations.
End with: SEVERITY_CLASSIFICATION: [LOW|MEDIUM|HIGH]"""

    resp = oai.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "system", "content": system_prompt},
                  {"role": "user",   "content": user_prompt}],
        temperature=0.2, max_tokens=300)
    assessment = resp.choices[0].message.content.strip()
    m          = re.search(r"SEVERITY_CLASSIFICATION:\s*(LOW|MEDIUM|HIGH)", assessment, re.IGNORECASE)
    sev        = m.group(1).lower() if m else severity_in

    log = {"timestamp": utc_now(), "node": "SurfaceDefectAgent",
           "output": {"severity_confirmed": sev, "assessment": assessment[:150]},
           "model": "gpt-4o-mini"}

    return {
        **_carry_pipeline_fields(state),
        "defect_observed":      state.get("defect_observed"),
        "defect_class":         state.get("defect_class"),
        "defect_type_observed": defect_type,
        "defect_location":      location,
        "vision_confidence":    state.get("vision_confidence"),
        "vision_evidence":      evidence,
        "agent_selected":       state.get("agent_selected"),
        "severity":             sev,
        "specialist_assessment":assessment,
        "decision_log":         state.get("decision_log", []) + [log]
    }


# ─── NODE 3b: STRUCTURAL DEFECT AGENT ────────────────────────────────────────
def structural_defect_agent_node(state: GlobalState) -> GlobalState:
    """Fix 2: Returns only severity + specialist_assessment + upstream fields."""
    oai         = get_client()
    defect_type = state.get("defect_type_observed", "")
    severity_in = state.get("severity", "high")
    location    = state.get("defect_location", "")
    evidence    = state.get("vision_evidence", "")
    category    = state.get("category", "unknown")

    system_prompt = "You are a structural defect characterisation specialist at CircuitSense."
    user_prompt   = f"""Characterise this structural defect:
Product: {category} | Type: {defect_type}
Severity: {severity_in} | Location: {location}
Evidence: {evidence}

Assess functional/safety impact, structural integrity, rework feasibility.
For capsules: consider patient safety. For PCBs: consider voltage risk.
End with: SEVERITY_CLASSIFICATION: [LOW|MEDIUM|HIGH]"""

    resp = oai.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "system", "content": system_prompt},
                  {"role": "user",   "content": user_prompt}],
        temperature=0.2, max_tokens=300)
    assessment = resp.choices[0].message.content.strip()
    m          = re.search(r"SEVERITY_CLASSIFICATION:\s*(LOW|MEDIUM|HIGH)", assessment, re.IGNORECASE)
    sev        = m.group(1).lower() if m else severity_in

    log = {"timestamp": utc_now(), "node": "StructuralDefectAgent",
           "output": {"severity_confirmed": sev, "assessment": assessment[:150]},
           "model": "gpt-4o-mini"}

    return {
        **_carry_pipeline_fields(state),
        "defect_observed":      state.get("defect_observed"),
        "defect_class":         state.get("defect_class"),
        "defect_type_observed": defect_type,
        "defect_location":      location,
        "vision_confidence":    state.get("vision_confidence"),
        "vision_evidence":      evidence,
        "agent_selected":       state.get("agent_selected"),
        "severity":             sev,
        "specialist_assessment":assessment,
        "decision_log":         state.get("decision_log", []) + [log]
    }


# ─── NODE 3c: PASSTHROUGH AGENT ──────────────────────────────────────────────
def passthrough_agent_node(state: GlobalState) -> GlobalState:
    """Fix 2: Returns only disposition fields + upstream fields."""
    sel  = state.get("agent_selected", "passthrough")
    conf = state.get("vision_confidence", 1.0)

    if sel == "uncertain":
        disp, assess, clause = (
            "UNCERTAIN",
            f"Vision confidence ({conf:.2f}) is below threshold 0.60. "
            "Case escalated to human inspector for review.",
            "General Policy §2: Low-confidence → human escalation."
        )
    else:
        disp, assess, clause = (
            "PASS",
            "No defect detected. Unit cleared for shipment.",
            "General Policy §1: No defect → PASS confirmed."
        )

    log = {"timestamp": utc_now(), "node": "PassThroughAgent",
           "output": {"disposition": disp}}

    return {
        **_carry_pipeline_fields(state),
        "defect_observed":      state.get("defect_observed"),
        "defect_class":         state.get("defect_class"),
        "defect_type_observed": state.get("defect_type_observed", ""),
        "defect_location":      state.get("defect_location", ""),
        "severity":             state.get("severity", ""),
        "vision_confidence":    conf,
        "vision_evidence":      state.get("vision_evidence", ""),
        "agent_selected":       sel,
        "disposition":          disp,
        "specialist_assessment":assess,
        "policy_clause":        clause,
        "policy_justification": assess,
        "decision_log":         state.get("decision_log", []) + [log]
    }


# ─── NODE 4: POLICY REASONING AGENT ──────────────────────────────────────────
def policy_reasoning_agent_node(state: GlobalState) -> GlobalState:
    """Fix 2: Returns only disposition fields + upstream fields."""
    # Skip if PassThrough already set disposition
    if state.get("disposition") in ["PASS", "UNCERTAIN"]:
        log = {"timestamp": utc_now(), "node": "PolicyReasoningAgent",
               "output": {"skipped": True}}
        return {
            **_carry_pipeline_fields(state),
            "defect_observed":      state.get("defect_observed"),
            "defect_class":         state.get("defect_class"),
            "defect_type_observed": state.get("defect_type_observed", ""),
            "defect_location":      state.get("defect_location", ""),
            "severity":             state.get("severity", ""),
            "vision_confidence":    state.get("vision_confidence"),
            "vision_evidence":      state.get("vision_evidence", ""),
            "agent_selected":       state.get("agent_selected"),
            "specialist_assessment":state.get("specialist_assessment", ""),
            "disposition":          state.get("disposition"),
            "policy_clause":        state.get("policy_clause", ""),
            "policy_justification": state.get("policy_justification", ""),
            "decision_log":         state.get("decision_log", []) + [log]
        }

    oai          = get_client()
    defect_class = state.get("defect_class", "surface")
    policy_text  = INSPECTION_POLICIES.get(defect_class, INSPECTION_POLICIES["general"])

    user_prompt = f"""Determine inspection disposition.

Product: {state.get('category')} | Defect: {state.get('defect_type_observed')}
Class: {defect_class} | Severity: {state.get('severity')} | Confidence: {state.get('vision_confidence',0):.2f}

Specialist assessment: {state.get('specialist_assessment','')}

Policy:
{policy_text}

Return JSON:
{{"disposition":"PASS|REWORK|SCRAP","policy_clause":"exact clause","justification":"2-3 sentences"}}"""

    try:
        resp   = oai.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "system", "content":
                       "You are the Policy Adjudication Agent at CircuitSense. Return valid JSON only."},
                      {"role": "user", "content": user_prompt}],
            temperature=0, max_tokens=300)
        raw    = re.sub(r"```json|```", "", resp.choices[0].message.content.strip())
        result = json.loads(raw)
    except Exception as e:
        result = {"disposition": "SCRAP", "policy_clause": "Error fallback",
                  "justification": str(e)[:100]}

    log = {"timestamp": utc_now(), "node": "PolicyReasoningAgent",
           "output": result, "policy_used": f"{defect_class} policy",
           "model": "gpt-4o-mini"}

    return {
        **_carry_pipeline_fields(state),
        "defect_observed":      state.get("defect_observed"),
        "defect_class":         defect_class,
        "defect_type_observed": state.get("defect_type_observed", ""),
        "defect_location":      state.get("defect_location", ""),
        "severity":             state.get("severity", ""),
        "vision_confidence":    state.get("vision_confidence"),
        "vision_evidence":      state.get("vision_evidence", ""),
        "agent_selected":       state.get("agent_selected"),
        "specialist_assessment":state.get("specialist_assessment", ""),
        "disposition":          result.get("disposition", "SCRAP"),
        "policy_clause":        result.get("policy_clause", ""),
        "policy_justification": result.get("justification", ""),
        "decision_log":         state.get("decision_log", []) + [log]
    }


# ─── NODE 5: RESPONSE NODE ────────────────────────────────────────────────────
def response_node(state: GlobalState) -> GlobalState:
    """Fix 2: Reads all fields explicitly. Returns only final_report + full state."""
    d_emoji = {"PASS": "✅", "REWORK": "🔧", "SCRAP": "❌", "UNCERTAIN": "⚠️"
               }.get(state.get("disposition", ""), "❓")
    report  = (
        f"CIRCUITSENSE INSPECTION REPORT\n"
        f"{'─'*45}\n"
        f"Policy ID    : {state.get('policy_id','N/A')}\n"
        f"Category     : {state.get('category','N/A').upper()}\n"
        f"Timestamp    : {utc_now()[:19].replace('T',' ')} UTC\n\n"
        f"VISION FINDINGS\n"
        f"Defect Class : {state.get('defect_class','N/A').upper()}\n"
        f"Defect Type  : {state.get('defect_type_observed','N/A')}\n"
        f"Severity     : {state.get('severity','N/A').upper()}\n"
        f"Confidence   : {state.get('vision_confidence',0):.0%}\n"
        f"Evidence     : {state.get('vision_evidence','N/A')}\n\n"
        f"DISPOSITION  : {d_emoji} {state.get('disposition','N/A')}\n"
        f"Clause       : {state.get('policy_clause','N/A')}\n"
        f"Justification: {state.get('policy_justification','N/A')}\n"
        f"GT Leak      : No\n"
        f"{'─'*45}"
    )
    log = {"timestamp": utc_now(), "node": "ResponseNode",
           "output": {"disposition": state.get("disposition"), "report_generated": True}}

    return {
        **_carry_pipeline_fields(state),
        "defect_observed":      state.get("defect_observed"),
        "defect_class":         state.get("defect_class"),
        "defect_type_observed": state.get("defect_type_observed", ""),
        "defect_location":      state.get("defect_location", ""),
        "severity":             state.get("severity", ""),
        "vision_confidence":    state.get("vision_confidence"),
        "vision_evidence":      state.get("vision_evidence", ""),
        "agent_selected":       state.get("agent_selected"),
        "specialist_assessment":state.get("specialist_assessment", ""),
        "disposition":          state.get("disposition"),
        "policy_clause":        state.get("policy_clause", ""),
        "policy_justification": state.get("policy_justification", ""),
        "final_report":         report,
        "decision_log":         state.get("decision_log", []) + [log]
    }


# ─── ROUTING ──────────────────────────────────────────────────────────────────
def route_after_supervisor(state: GlobalState) -> str:
    return {
        "surface":     "surface_agent",
        "structural":  "structural_agent",
        "passthrough": "passthrough_agent",
        "uncertain":   "passthrough_agent"
    }.get(state.get("agent_selected", "passthrough"), "passthrough_agent")


# ─── GRAPH ────────────────────────────────────────────────────────────────────
@st.cache_resource
def build_graph():
    wf = StateGraph(GlobalState)
    wf.add_node("vision_agent",           vision_agent_node)
    wf.add_node("supervisor_agent",       supervisor_agent_node)
    wf.add_node("surface_agent",          surface_defect_agent_node)
    wf.add_node("structural_agent",       structural_defect_agent_node)
    wf.add_node("passthrough_agent",      passthrough_agent_node)
    wf.add_node("policy_reasoning_agent", policy_reasoning_agent_node)
    wf.add_node("response_node",          response_node)

    wf.set_entry_point("vision_agent")
    wf.add_edge("vision_agent", "supervisor_agent")
    wf.add_conditional_edges(
        "supervisor_agent", route_after_supervisor,
        {"surface_agent":    "surface_agent",
         "structural_agent": "structural_agent",
         "passthrough_agent":"passthrough_agent"}
    )
    wf.add_edge("surface_agent",          "policy_reasoning_agent")
    wf.add_edge("structural_agent",       "policy_reasoning_agent")
    wf.add_edge("passthrough_agent",      "response_node")
    wf.add_edge("policy_reasoning_agent", "response_node")
    wf.add_edge("response_node",          END)
    return wf.compile()


# ─── SESSION STATE ────────────────────────────────────────────────────────────
def init_session():
    for k, v in {
        "inspection_history": [], "openai_configured": False,
        "openai_api_key": "", "openai_api_base": ""
    }.items():
        if k not in st.session_state:
            st.session_state[k] = v

init_session()

# Auto-load key from HF Spaces secret
if "OPENAI_API_KEY" in os.environ and not st.session_state.get("openai_api_key"):
    st.session_state.openai_api_key  = os.environ["OPENAI_API_KEY"]
    st.session_state.openai_configured = True

df_enriched = load_enriched_dataset()
app         = build_graph()


# ══════════════════════════════════════════════════════════════════════════════
# STREAMLIT UI — unchanged from v1 except matplotlib import added
# ══════════════════════════════════════════════════════════════════════════════

st.markdown("""
<div class="main-header">
  <h1 style="margin:0;font-size:1.8em;">🔬 CircuitSense — AI Quality Inspection</h1>
  <p style="margin:5px 0 0 0;opacity:.85;">
    GPT-4o Vision · 5-Node LangGraph · Vision → Specialist → Policy Reasoning
  </p>
</div>
""", unsafe_allow_html=True)

# ── SIDEBAR ───────────────────────────────────────────────────────────────────
with st.sidebar:
    st.markdown("## ⚙️ Configuration")
    with st.expander("🔑 API Keys", expanded=not st.session_state.openai_configured):
        okey  = st.text_input("OpenAI API Key", type="password",
                               value=st.session_state.get("openai_api_key", ""),
                               placeholder="sk-...")
        obase = st.text_input("API Base URL (optional)",
                               value=st.session_state.get("openai_api_base", ""),
                               placeholder="Azure/proxy endpoint")
        lskey = st.text_input("LangSmith Key (optional)", type="password",
                               placeholder="ls__...")
        if okey:
            st.session_state.openai_api_key  = okey
            st.session_state.openai_configured = True
        if obase:
            st.session_state.openai_api_base = obase
        if lskey:
            os.environ.update({"LANGCHAIN_TRACING_V2": "true",
                                "LANGCHAIN_API_KEY":    lskey,
                                "LANGCHAIN_PROJECT":    "MLS1-CircuitSense-Inspection"})

    st.divider()
    st.markdown("## 📋 Inspection Pipeline")
    st.markdown("""
```
Image Input
    ↓
Vision Agent (gpt-4o)
    ↓
Supervisor Agent (routing)
    ↓
Surface / Structural /
PassThrough Agent
    ↓
Policy Reasoning Agent
    ↓
Response Node
```
""")
    st.divider()
    if st.button("🗑️ Clear History", use_container_width=True):
        st.session_state.inspection_history = []
        st.rerun()


# ── MAIN AREA ─────────────────────────────────────────────────────────────────
tab_inspect, tab_history, tab_log = st.tabs([
    "🔬 Run Inspection", "📊 Session History", "📋 Audit Trail"
])

with tab_inspect:
    col_input, col_result = st.columns([1, 1])

    with col_input:
        st.markdown("### 📥 Product Image Input")
        input_method = st.radio("Image source",
                                ["Upload image", "Select from dataset"],
                                horizontal=True)

        image_pil      = None
        category       = "unknown"
        defect_type_gt = "unknown"
        policy_id      = f"MANUAL-{datetime.datetime.now().strftime('%H%M%S')}"

        if input_method == "Upload image":
            uploaded = st.file_uploader("Upload product image (JPEG/PNG)",
                                         type=["jpg", "jpeg", "png"])
            if uploaded:
                image_pil = Image.open(uploaded)
                st.image(image_pil, caption="Uploaded image", width="stretch")
                category       = st.selectbox("Product category",
                                               ["pcb1", "capsules", "cashew", "other"])
                defect_type_gt = st.text_input("Known defect label (optional)",
                                                placeholder="e.g. scratch")

        else:
            if df_enriched.empty:
                st.warning("df_enriched.csv not found. Run the notebook first.")
            else:
                cat_filter  = st.selectbox("Filter by category",
                                            ["all"] + list(df_enriched['category'].unique()))
                df_filtered = df_enriched if cat_filter == "all" \
                              else df_enriched[df_enriched['category'] == cat_filter]
                options = [f"{r['policy_id']} — {r['category']} / {r['defect_type']}"
                           for _, r in df_filtered.iterrows()]
                sel = st.selectbox("Select inspection record", options)
                if sel:
                    pid            = sel.split(" — ")[0]
                    row            = df_enriched[df_enriched['policy_id'] == pid].iloc[0]
                    policy_id      = row['policy_id']
                    category       = row['category']
                    defect_type_gt = row['defect_type']
                    try:
                        image_pil = Image.open(row['image_path'])
                        st.image(image_pil,
                                 caption=f"{category} / {defect_type_gt}",
                                 width="stretch")
                        st.caption(f"**Description:** {row.get('defect_description','N/A')}")
                    except Exception:
                        st.error("Image file not found. Run the notebook to download VisA.")

        run_btn = st.button("🚀 Run Inspection", type="primary",
                            use_container_width=True, disabled=(image_pil is None))

    with col_result:
        st.markdown("### 📊 Inspection Result")

        if run_btn and image_pil is not None:
            if not st.session_state.get("openai_api_key") \
                    and "OPENAI_API_KEY" not in os.environ:
                st.error("⚠️ Please enter your OpenAI API key in the sidebar.")
            else:
                with st.spinner("Running 5-node inspection pipeline..."):
                    try:
                        image_b64 = resize_and_encode(image_pil)
                        initial   = GlobalState(
                            image_b64=image_b64, category=category,
                            defect_type_gt=defect_type_gt, policy_id=policy_id,
                            decision_log=[]
                        )
                        result = app.invoke(initial)
                        st.session_state.inspection_history.append(result)
                        st.success("✅ Inspection complete")

                        disp        = result.get("disposition", "N/A")
                        badge_class = {"PASS":      "badge-pass",
                                       "REWORK":    "badge-rework",
                                       "SCRAP":     "badge-scrap",
                                       "UNCERTAIN": "badge-uncertain"}.get(disp, "")
                        st.markdown(f'<br><span class="{badge_class}">⬤ {disp}</span><br><br>',
                                    unsafe_allow_html=True)

                        m1, m2, m3, m4 = st.columns(4)
                        m1.metric("Defect Class", result.get("defect_class", "N/A").upper())
                        m2.metric("Severity",     result.get("severity", "N/A").upper())
                        m3.metric("Confidence",   f"{result.get('vision_confidence',0):.0%}")
                        m4.metric("Nodes Run",    len(result.get("decision_log", [])))

                        with st.expander("🔍 Vision Findings", expanded=True):
                            st.write(f"**Defect type:** {result.get('defect_type_observed','N/A')}")
                            st.write(f"**Location:** {result.get('defect_location','N/A')}")
                            st.write(f"**Evidence:** {result.get('vision_evidence','N/A')}")

                        with st.expander("⚙️ Specialist Assessment", expanded=True):
                            st.write(f"**Agent:** "
                                     f"{result.get('agent_selected','N/A').upper()} DEFECT AGENT")
                            st.write(result.get('specialist_assessment', 'N/A'))

                        with st.expander("📋 Policy Decision", expanded=True):
                            st.write(f"**Policy clause:** {result.get('policy_clause','N/A')}")
                            st.write(f"**Justification:** {result.get('policy_justification','N/A')}")

                    except Exception as e:
                        st.error(f"Inspection failed: {str(e)}")

        elif not run_btn:
            st.info("Select or upload a product image and click Run Inspection.")


with tab_history:
    st.markdown("### 📊 Session Inspection History")
    if st.session_state.inspection_history:
        rows = []
        for r in st.session_state.inspection_history:
            rows.append({
                "Policy ID":   r.get("policy_id", ""),
                "Category":    r.get("category", ""),
                "Defect GT":   r.get("defect_type_gt", ""),
                "Defect Seen": r.get("defect_type_observed", ""),
                "Class":       r.get("defect_class", ""),
                "Severity":    r.get("severity", ""),
                "Agent":       r.get("agent_selected", ""),
                "Disposition": r.get("disposition", ""),
                "Confidence":  f"{r.get('vision_confidence',0):.0%}",
                "Log Entries": len(r.get("decision_log", []))
            })
        df_hist = pd.DataFrame(rows)

        def color_disposition(val):
            colors = {"PASS": "#e8f5e9", "REWORK": "#fff3e0",
                      "SCRAP": "#ffebee", "UNCERTAIN": "#f5f5f5"}
            return f"background-color:{colors.get(val,'white')}"

        st.dataframe(
            df_hist.style.map(color_disposition, subset=["Disposition"]),
            use_container_width=True
        )

        if len(df_hist) > 1:
            disp_colors = {"PASS": "#4CAF50", "REWORK": "#FF9800",
                           "SCRAP": "#F44336", "UNCERTAIN": "#9E9E9E"}
            disp_counts = df_hist["Disposition"].value_counts()
            bar_colors  = [disp_colors.get(d, "#333333") for d in disp_counts.index]
            fig, ax = plt.subplots(figsize=(6, 3))
            disp_counts.plot(kind='bar', ax=ax, color=bar_colors, edgecolor='white')
            ax.set_title("Disposition Distribution — This Session")
            ax.tick_params(axis='x', rotation=0)
            plt.tight_layout()
            st.pyplot(fig)

        st.download_button(
            "⬇️ Download Session Results (CSV)",
            data=df_hist.to_csv(index=False),
            file_name=f"circuitsense_results_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
            mime="text/csv"
        )
    else:
        st.info("No inspections run yet in this session.")


with tab_log:
    st.markdown("### 📋 Full Audit Trail")
    st.caption("Every node execution is logged here.")

    if st.session_state.inspection_history:
        for i, result in enumerate(reversed(st.session_state.inspection_history), 1):
            disp  = result.get("disposition", "N/A")
            emoji = {"PASS": "✅", "REWORK": "🔧",
                     "SCRAP": "❌", "UNCERTAIN": "⚠️"}.get(disp, "❓")
            with st.expander(
                f"Inspection {len(st.session_state.inspection_history)-i+1} — "
                f"{result.get('policy_id','')} | {result.get('category','')} | "
                f"{emoji} {disp}",
                expanded=(i == 1)
            ):
                COLORS = {
                    "VisionAgent":          "#E3F2FD",
                    "SupervisorAgent":      "#F3E5F5",
                    "SurfaceDefectAgent":   "#E8F5E9",
                    "StructuralDefectAgent":"#FFF3E0",
                    "PassThroughAgent":     "#F5F5F5",
                    "PolicyReasoningAgent": "#FCE4EC",
                    "ResponseNode":         "#E0F2F1"
                }
                for entry in result.get("decision_log", []):
                    node  = entry.get("node", "")
                    color = COLORS.get(node, "#FAFAFA")
                    st.markdown(
                        f'<div style="background:{color};padding:8px;'
                        f'border-radius:6px;margin:4px 0;">'
                        f'<b>{node}</b> &nbsp;|&nbsp; '
                        f'<small>{entry.get("timestamp","")[:19].replace("T"," ")}</small><br>'
                        f'<small>{str(entry.get("output",""))[:200]}</small>'
                        f'</div>',
                        unsafe_allow_html=True
                    )
                st.download_button(
                    f"⬇️ Download Audit Log — {result.get('policy_id','')}",
                    data=json.dumps(result.get("decision_log", []), indent=2),
                    file_name=f"audit_{result.get('policy_id','result')}.json",
                    mime="application/json"
                )
    else:
        st.info("No inspections run yet. Go to 'Run Inspection' to start.")


# ── FOOTER ────────────────────────────────────────────────────────────────────
st.divider()
st.markdown("""
<div style="text-align:center;color:#888;font-size:.8em;padding:10px">
  CircuitSense AI Inspection — MLS-1 v4 | GPT-4o Vision · LangGraph 5-Node Pipeline · VisA Dataset (CC BY 4.0)
  <br>⚠️ Demonstration system. Not for production use without human oversight.
</div>
""", unsafe_allow_html=True)