"""
RecallTrace OpenEnv — Hugging Face Spaces Entry Point
Space: https://huggingface.co/spaces/Majen/recalltrace-openenv
Team: Shamanth MS | P G Ayush Rai | Shreya B J
"""

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))

import gradio as gr
from recall_env.env import RecallTraceEnv
from grader.grader import grade, compute_reward
from scenario.scenario import get_scenario, list_levels

# ── helpers ──────────────────────────────────────────────────────────────────

def run_episode(level: str) -> str:
    scenario = get_scenario(level)
    env = RecallTraceEnv(level=level)
    obs = env.reset()

    lot_catalog   = scenario.get("lot_catalog", {})
    transformations = scenario.get("transformations", {})
    contaminated_lot = scenario.get("contaminated_lot", "")

    contaminated_lots = {l for l, m in lot_catalog.items() if m.get("contaminated", False)}
    contaminated_lots.add(contaminated_lot)
    if contaminated_lot in transformations:
        contaminated_lots.add(transformations[contaminated_lot])

    actions = []
    for node_id in scenario["nodes"]:
        actions.append({"type": "inspect_node", "node_id": node_id})
    for lot_id in sorted(contaminated_lots):
        actions.append({"type": "trace_lot", "lot_id": lot_id})
    for node_id, node_data in scenario["nodes"].items():
        for lot_id, qty in node_data["inventory"].items():
            if lot_id in contaminated_lots and qty > 0:
                actions.append({"type": "quarantine", "node_id": node_id, "lot_id": lot_id, "quantity": qty})
    notified = set()
    for node_id, node_data in scenario["nodes"].items():
        if any(l in contaminated_lots for l in node_data["inventory"]):
            if node_id not in notified:
                actions.append({"type": "notify", "node_id": node_id})
                notified.add(node_id)
    actions.append({"type": "finalize"})

    lines = []
    lines.append(f"📋 Recall Notice: {obs['recall_notice']}\n")
    lines.append(f"🗺️  Nodes: {list(scenario['nodes'].keys())}\n")
    lines.append("─" * 55)

    final_info = {}
    env2 = RecallTraceEnv(level=level)
    env2.reset()
    for action in actions:
        obs, reward, done, info = env2.step(action)
        lines.append(f"▶ {action['type']:15s} → reward={reward:+.2f}  done={done}")
        if done:
            final_info = info
            break

    score = final_info.get("score", 0.0)
    bd    = final_info.get("breakdown", {})
    lines.append("─" * 55)
    lines.append(f"✅ Final Score        : {score:.4f}")
    lines.append(f"   Quarantine Score  : {bd.get('quarantine_score', '-')}")
    lines.append(f"   Notification Score: {bd.get('notification_score', '-')}")
    lines.append(f"   All Stock Clear   : {bd.get('all_stock_quarantined', '-')}")
    lines.append(f"   All Nodes Notified: {bd.get('all_affected_notified', '-')}")

    return "\n".join(lines)


def run_all():
    results = []
    for level in list_levels():
        results.append(f"{'='*55}")
        results.append(f"  LEVEL: {level.upper()}")
        results.append(f"{'='*55}")
        results.append(run_episode(level))
        results.append("")
    return "\n".join(results)


# ── Gradio UI ─────────────────────────────────────────────────────────────────

with gr.Blocks(title="RecallTrace OpenEnv") as demo:
    gr.Markdown("""
# 🔍 RecallTrace OpenEnv
**Supply-chain recall traceability — OpenEnv compliant**

An AI agent traces contaminated lots through a shipment graph,
quarantines affected inventory, and notifies relevant nodes.

**Team:** Shamanth MS · P G Ayush Rai · Shreya B J
""")

    with gr.Row():
        level_dropdown = gr.Dropdown(
            choices=["easy", "medium", "hard"],
            value="easy",
            label="Select Task Level"
        )
        run_btn  = gr.Button("▶ Run Episode", variant="primary")
        run_all_btn = gr.Button("⚡ Run All Levels")

    output = gr.Textbox(label="Episode Output", lines=30, max_lines=50)

    run_btn.click(fn=run_episode, inputs=level_dropdown, outputs=output)
    run_all_btn.click(fn=run_all, inputs=None, outputs=output)

    gr.Markdown("""
---
### Action Space
| Action | Description |
|--------|-------------|
| `inspect_node` | Reveal inventory at a node |
| `trace_lot` | Trace a lot across the network |
| `quarantine` | Isolate contaminated stock |
| `notify` | Send recall alert to a node |
| `finalize` | Submit plan → get final score |

### Scoring: 0.0 → 1.0
- Quarantine accuracy + notification coverage
- Penalised for unnecessary quarantines
""")

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