api.py

"""
api.py
======
FastAPI backend for the DeFi Risk Simulation Lab.

Endpoints:
  GET  /                         -> Serve frontend index.html
  GET  /static/*                 -> Serve frontend static files
  POST /api/run                  -> Start a new simulation (returns run_id)
  GET  /api/status/{run_id}      -> Check run status / summary
  GET  /api/results              -> Get latest timeseries CSV data
  GET  /api/blockchain/status    -> Hardhat node connection info
  GET  /api/blockchain/tx/{hash} -> On-chain tx receipt by hash
  GET  /api/audit_trail/{run_id} -> All blockchain tx hashes for a run
  WS   /ws/{run_id}              -> WebSocket: stream step events in real-time

WebSocket event schema (JSON):
  {
    "type":       "STEP" | "CRASH" | "WHALE" | "ARB" | "MEV" | "DONE" | "ERROR",
    "step":       int,
    "total_steps": int,
    "progress":   float,
    "pool_usdc":  float,
    "pool_eth":   float,
    "amm_price":  float,
    "oracle_price": float,
    "spread_pct": float,
    "price_delta": float,
    "events":     [ { "agent_type", "action", "amount", "token", "reasoning" } ],
    "crash_detail": { ... },          // only on CRASH / WHALE
    "blockchain": { "tx_count": int } // count of on-chain txs this step
  }
"""

import asyncio
import json
import logging
import os
import re
import subprocess
import tempfile
import uuid
from pathlib import Path
from typing import Dict, Any, Optional

import uvicorn
import httpx
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, UploadFile, File, Form, Request
from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
from pydantic import BaseModel
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware

# ── Stress-testing engine imports ───────────────────────────────────
from contract_analyzer import analyze_source
from nlp_engine        import parse_scenario
from script_generator  import generate_hardhat_script, generate_from_scenario_id
from script_validator  import validate_and_sanitize
from risk_analyzer     import full_risk_report
from scenarios         import list_scenarios, get_scenario

import config
from database import SimulationLogger
from agent_config_transformer import generate_agent_configs
from model import DeFiSimulationModel

# ── Conditional blockchain imports ─────────────────────────────────
if config.BLOCKCHAIN_ENABLED:
    try:
        from blockchain_adapter import BlockchainAdapter
        from blockchain_client import BlockchainClient
        _BLOCKCHAIN_READY = True
    except ImportError:
        _BLOCKCHAIN_READY = False
        logging.warning("blockchain_adapter not found — running in mock-only mode")
else:
    _BLOCKCHAIN_READY = False

# Shared BlockchainClient (singleton) for status endpoint
_bc_singleton = None
if _BLOCKCHAIN_READY:
    try:
        _bc_singleton = BlockchainClient(rpc_url=config.HARDHAT_RPC_URL)
    except Exception:
        pass

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ─────────────────────────────────────────────────────────────────── #
# APP                                                                  #
# ─────────────────────────────────────────────────────────────────── #
app = FastAPI(title="DeFi Risk Simulation Lab", version="2.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

FRONTEND_DIR = Path(__file__).parent / "frontend"
app.mount("/static", StaticFiles(directory=str(FRONTEND_DIR)), name="static")

# In-memory run store
_runs: Dict[str, Dict[str, Any]] = {}

# Stress-test in-memory store: {job_id -> {status, script, analysis, results}}
_stress_jobs: Dict[str, Dict[str, Any]] = {}

# Stores the latest simulation completion data for Ask Arpit context
_last_sim_result: Dict[str, Any] = {}


# ─────────────────────────────────────────────────────────────────── #
# ROUTES — static / basic                                              #
# ─────────────────────────────────────────────────────────────────── #

@app.get("/", response_class=HTMLResponse)
@app.get("/home", response_class=HTMLResponse)
async def root():
    """Serve the Home Dashboard page."""
    page_path = FRONTEND_DIR / "home.html"
    return HTMLResponse(content=page_path.read_text(encoding="utf-8"))


@app.get("/playground", response_class=HTMLResponse)
async def playground_page():
    """Serve the Agent Simulation (Playground) page."""
    index_path = FRONTEND_DIR / "index.html"
    return HTMLResponse(content=index_path.read_text(encoding="utf-8"))


@app.get("/stress", response_class=HTMLResponse)
async def stress_test_page():
    """Serve the Smart Contract Stress Testing Engine page."""
    page_path = FRONTEND_DIR / "stress_test.html"
    return HTMLResponse(content=page_path.read_text(encoding="utf-8"))


@app.get("/runs", response_class=HTMLResponse)
async def runs_page():
    """Serve the Simulation Runs history page."""
    page_path = FRONTEND_DIR / "runs.html"
    return HTMLResponse(content=page_path.read_text(encoding="utf-8"))


@app.post("/api/run")
async def start_run(
    steps: int = config.SIMULATION_STEPS, 
    groq_key: str = "", 
    prompt: str = "", 
    scenario_id: str = "",
    amm_config: str = "",
    liquidity_config: str = ""
):
    run_id = str(uuid.uuid4())[:8]
    if groq_key:
        os.environ["GROQ_API_KEY"] = groq_key

    amm_cfg = None
    if amm_config:
        try:
            amm_cfg = json.loads(amm_config)
        except Exception as e:
            logger.warning("Failed to parse amm_config JSON: %s", e)

    liq_cfg = None
    if liquidity_config:
        try:
            liq_cfg = json.loads(liquidity_config)
        except Exception as e:
            logger.warning("Failed to parse liquidity_config JSON: %s", e)

    _runs[run_id] = {
        "status": "pending", 
        "steps": steps, 
        "events": [], 
        "audit_log": [],
        "prompt": prompt,
        "scenario_id": scenario_id,
        "amm_config": amm_cfg,
        "liquidity_config": liq_cfg
    }
    logger.info("New run created: %s (%d steps)", run_id, steps)
    return {"run_id": run_id, "steps": steps}


@app.get("/api/status/{run_id}")
async def get_status(run_id: str):
    run = _runs.get(run_id)
    if not run:
        raise HTTPException(status_code=404, detail="Run not found")
    return run


@app.get("/api/results")
async def get_results():
    csv_path = Path(config.CSV_PATH)
    if not csv_path.exists():
        return {"rows": []}
    import csv
    with open(csv_path, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        rows = list(reader)
    return {"rows": rows[:500]}


# ─────────────────────────────────────────────────────────────────── #
# ROUTES — blockchain                                                   #
# ─────────────────────────────────────────────────────────────────── #

@app.get("/api/blockchain/status")
async def blockchain_status():
    """Return Hardhat node connection info and deployed contract addresses."""
    if _bc_singleton:
        return _bc_singleton.get_status()
    return {"connected": False, "reason": "Blockchain integration disabled or web3 not installed"}


@app.get("/api/blockchain/tx/{tx_hash}")
async def blockchain_tx(tx_hash: str):
    """Return on-chain receipt for a given transaction hash."""
    if not _bc_singleton:
        raise HTTPException(status_code=503, detail="Blockchain not available")
    receipt = _bc_singleton.get_transaction(tx_hash)
    if not receipt:
        raise HTTPException(status_code=404, detail="Transaction not found")
    return receipt


@app.get("/api/audit_trail/{run_id}")
async def audit_trail(run_id: str):
    """Return all blockchain audit entries recorded during a simulation run."""
    run = _runs.get(run_id)
    if not run:
        raise HTTPException(status_code=404, detail="Run not found")
    return {
        "run_id":      run_id,
        "total":       len(run.get("audit_log", [])),
        "on_chain":    sum(1 for e in run.get("audit_log", []) if e.get("on_chain")),
        "entries":     run.get("audit_log", []),
    }


# ─────────────────────────────────────────────────────────────────── #
# ROUTES — SMART CONTRACT STRESS TESTING ENGINE                        #
# ─────────────────────────────────────────────────────────────────── #

@app.post("/api/contract/upload")
async def contract_upload(file: UploadFile = File(...)):
    """
    Upload a .sol file and return its static analysis.
    Returns functions, state vars, events, vulnerabilities, interaction map.
    """
    if not file.filename or not file.filename.endswith(".sol"):
        raise HTTPException(status_code=400, detail="Only .sol files are accepted")
    try:
        source  = (await file.read()).decode("utf-8")
        analysis = analyze_source(source, filename=file.filename)
        job_id  = str(uuid.uuid4())[:8]
        _stress_jobs[job_id] = {
            "status":   "analyzed",
            "filename": file.filename,
            "analysis": analysis,
            "script":   None,
            "results":  None,
        }
        return {"job_id": job_id, "analysis": analysis}
    except Exception as e:
        logger.exception("Contract upload failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/contract/from-url")
async def contract_from_url(url: str = Form(...)):
    """
    Fetch a Solidity contract from a raw GitHub URL and analyze it.
    E.g.: https://raw.githubusercontent.com/owner/repo/main/contracts/Token.sol
    """
    if not url.startswith("https://"):
        raise HTTPException(status_code=400, detail="URL must start with https://")
    try:
        async with httpx.AsyncClient(timeout=15) as client:
            resp = await client.get(url)
        if resp.status_code != 200:
            raise HTTPException(status_code=400, detail=f"Failed to fetch URL: HTTP {resp.status_code}")
        source   = resp.text
        filename = url.split("/")[-1] or "contract.sol"
        analysis = analyze_source(source, filename=filename)
        job_id   = str(uuid.uuid4())[:8]
        _stress_jobs[job_id] = {
            "status":   "analyzed",
            "filename": filename,
            "source":   source,
            "analysis": analysis,
            "script":   None,
            "results":  None,
        }
        return {"job_id": job_id, "analysis": analysis}
    except HTTPException:
        raise
    except Exception as e:
        logger.exception("Contract from URL failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/script/scenarios")
async def get_scenarios():
    """Return the list of prebuilt historical crisis scenarios."""
    return {"scenarios": list_scenarios()}


# ─────────────────────────────────────────────────────────────────── #
# AI METRIC EXPLANATION                                                #
# ─────────────────────────────────────────────────────────────────── #
class ExplainMetricRequest(BaseModel):
    metric_id: str
    metric_title: str
    metric_data: dict = {}


class AskArpitRequest(BaseModel):
    question: str
    history: list = []  # [{"role": "user"|"assistant", "content": str}]


@app.post("/api/ask-arpit")
async def ask_arpit(req: AskArpitRequest):
    """
    Conversational AI endpoint. Uses the latest simulation audit as context.
    Maintains conversation history for follow-up questions.
    """
    audit_ctx = _last_sim_result.get("markdown_audit", "")
    summary_ctx = _last_sim_result.get("summary", {})

    import json as _json
    summary_str = _json.dumps(summary_ctx, indent=2)[:600] if summary_ctx else ""
    audit_str   = audit_ctx[:2400] if audit_ctx else ""

    system_prompt = (
        "You are Arpit, an expert DeFi security analyst and smart contract auditor. "
        "You just completed a stress simulation on a DeFi protocol and have the full audit report. "
        "Answer questions about the simulation results in clear, concise language. "
        "Be direct and technical but avoid unnecessary jargon. "
        "When the user asks about specific risks, vulnerabilities, or numbers, refer to the audit data provided.\n\n"
        f"=== SIMULATION SUMMARY ===\n{summary_str}\n\n"
        f"=== AUDIT REPORT ===\n{audit_str}"
    )

    try:
        from groq import Groq  # type: ignore[import]
        if not config.GROQ_API_KEY:
            raise ValueError("GROQ_API_KEY not set")

        client = Groq(api_key=config.GROQ_API_KEY)

        # Build messages: system + prior history + new question
        messages = [{"role": "system", "content": system_prompt}]
        for h in (req.history or [])[-8:]:  # keep last 8 turns
            if h.get("role") in ("user", "assistant") and h.get("content"):
                messages.append({"role": h["role"], "content": h["content"]})
        messages.append({"role": "user", "content": req.question})

        response = client.chat.completions.create(
            model=config.GROQ_MODEL,
            messages=messages,
            temperature=0.5,
            max_tokens=512,
        )
        answer = response.choices[0].message.content.strip()
        return {"answer": answer}

    except Exception as e:
        logger.warning(f"ask-arpit Groq call failed: {e}")
        if not audit_ctx:
            return {"answer": "No simulation results are available yet. Run a simulation first, then ask me anything about the results."}
        return {"answer": "I ran into a connection issue with the AI engine. Please check your Groq API key and try again."}


_SUPA_URL  = config.SUPABASE_URL
_SUPA_KEY  = config.SUPABASE_KEY
_SUPA_BUCKET = "audit_reports"


@app.post("/api/upload-audit")
async def upload_audit(file: UploadFile = File(...)):
    """
    Proxy PDF upload to Supabase Storage.
    Called from the browser to avoid CORS issues with direct Supabase Storage requests.
    Returns the public URL of the uploaded report.
    """
    pdf_bytes = await file.read()
    filename  = file.filename or f"DeFiLab_Audit_{uuid.uuid4().hex[:8]}.pdf"
    upload_url = f"{_SUPA_URL}/storage/v1/object/{_SUPA_BUCKET}/{filename}"

    try:
        async with httpx.AsyncClient(timeout=30) as client:
            resp = await client.put(
                upload_url,
                content=pdf_bytes,
                headers={
                    "Authorization": f"Bearer {_SUPA_KEY}",
                    "apikey": _SUPA_KEY,
                    "Content-Type": "application/pdf",
                    "x-upsert": "true",
                },
            )
    except httpx.ConnectError as exc:
        logger.error("Supabase DNS/connection error: %s", exc)
        raise HTTPException(
            status_code=503,
            detail=(
                f"Cannot reach Supabase at {_SUPA_URL}. "
                "The project subdomain may not have propagated yet (new projects can take a few minutes). "
                "Verify your project URL in Supabase Dashboard → Settings → API."
            ),
        )

    if resp.status_code not in (200, 201):
        logger.error("Supabase upload failed: %s %s", resp.status_code, resp.text)
        raise HTTPException(status_code=502, detail=f"Supabase upload failed: {resp.text}")

    public_url = f"{_SUPA_URL}/storage/v1/object/public/{_SUPA_BUCKET}/{filename}"
    return {"publicUrl": public_url, "filename": filename}


@app.post("/api/explain-metric")
async def explain_metric(req: ExplainMetricRequest):
    _FALLBACKS = {
        "protocol_health": "This chart tracks the Total Value Locked (TVL) and oracle price over time. TVL shows how much capital is in the protocol — a sharp drop indicates funds are being drained or withdrawn. The dashed oracle price line shows the reference price used for collateral valuations; large deviations from the AMM price create arbitrage windows and can enable manipulation attacks.",
        "solvency": "The solvency meter shows the protocol's current financial health as a percentage of its starting capital. 100% means fully healthy with no value lost. Below 50% signals significant stress — roughly half the protocol's value has been moved out or drained. Below 20% indicates near-insolvency, where the protocol may not be able to repay all depositors.",
        "gatekeeper": "The Gatekeeper Matrix shows which transactions were allowed (green checkmarks) versus blocked and reverted (red). The contract's IdentityGatekeeper verifies each caller before allowing sensitive operations. A high ratio of red REVERTs means attackers are being stopped at the gate — the access control is working. If all transactions are green, either the simulation is benign or the gatekeeper isn't filtering effectively.",
        "attack_flow": "This diagram traces how capital flows through the protocol during an attack. Lit red edges show active attack paths: from a Flash Lender (instant uncollateralised loan), to the AMM pool (price manipulation via large swap), to the Oracle (which is now reading a manipulated price), and finally to the Lending Pool where undercollateralised positions get liquidated. Each lit edge is an attack vector currently being exploited.",
        "agent_activity": "This stacked bar chart counts how many on-chain transactions each agent class makes per simulation step. Tall red/attacker bars indicate concentrated malicious activity in those steps. If retail and mid-tier bars dominate, the protocol is experiencing normal trading pressure. A sudden spike in MEV/Attacker bars often precedes a solvency drop — watch for that correlation.",
        "abm": "This intelligence matrix profiles every agent type's on-chain behaviour. TXNS and TOTAL SPENT show activity volume. VELOCITY measures transactions per step — a burst above 5 triggers a flag. The ANOMALY SCORE (0-10) summarises detected risk signals. Red DANGER rows are agents whose actions match known attack patterns like flash-loan exploits, liquidation cascades, or volume dominance. Yellow WARNING rows need monitoring.",
        "lending": "This panel monitors the health of the lending pool in real time. COLLATERAL is the ETH backing loans. TOTAL DEBT is outstanding USDC borrowed. The HEALTH FACTOR (HF) is the collateral-to-debt ratio — below 1.0 means a position is under water and can be liquidated. BAD DEBT is debt that can no longer be fully recovered even after liquidation. UTILIZATION above 80% means the pool is nearly fully borrowed, leaving little liquidity buffer.",
    }

    try:
        from groq import Groq  # type: ignore[import]
        if not config.GROQ_API_KEY:
            raise ValueError("GROQ_API_KEY not configured")

        client = Groq(api_key=config.GROQ_API_KEY)

        import json as _json
        data_summary = _json.dumps(req.metric_data, indent=2)[:900]

        prompt = (
            f"You are a DeFi risk analyst assistant explaining a live dashboard metric to a non-technical user.\n\n"
            f"Metric: {req.metric_title}\n"
            f"Current Data:\n{data_summary}\n\n"
            f"Instructions:\n"
            f"- Use simple, plain English — no heavy jargon\n"
            f"- In 3-5 sentences explain WHAT this metric shows AND what the current values mean for the protocol's health\n"
            f"- If values look concerning (high anomaly scores, low solvency, many blocked transactions, low health factor, etc.) say so clearly\n"
            f"- End with one short, actionable takeaway\n"
            f"- Reply with only the explanation text, no headers or bullet points"
        )

        response = client.chat.completions.create(
            model=config.GROQ_MODEL,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.45,
            max_tokens=320,
        )
        explanation = response.choices[0].message.content.strip()
        return {"explanation": explanation}

    except Exception as e:
        logger.warning(f"explain-metric Groq call failed ({e}), using fallback")
        fallback = _FALLBACKS.get(
            req.metric_id,
            "This metric provides real-time insight into the DeFi protocol behaviour during the stress simulation. Higher anomaly values and lower solvency scores indicate increased systemic risk."
        )
        return {"explanation": fallback}


@app.post("/api/script/generate")
async def generate_script(
    prompt:      str = Form(...),
    job_id:      str = Form(""),
    scenario_id: str = Form(""),
    amm_config:  str = Form(""),
    liquidity_config: str = Form(""),
):
    """
    Generate a Hardhat stress-test script.
    Either from NLP prompt (+ optional job_id for contract context)
    or from a prebuilt scenario_id.
    Returns: {job_id, script, validation, scenario}
    """
    try:
        # Resolve contract analysis context
        analysis = None
        if job_id and job_id in _stress_jobs:
            analysis = _stress_jobs[job_id].get("analysis")

        # Resolve scenario
        if scenario_id and not prompt.strip():
            scenario = get_scenario(scenario_id)
        else:
            scenario = parse_scenario(prompt, contract_context=analysis)
            if scenario_id:
                # merge historical parameters with NLP actor/description
                base = get_scenario(scenario_id)
                scenario["parameters"]     = base.get("parameters", {})
                scenario["market_events"]  = base.get("parameters", {}).get("market_events", scenario.get("market_events", []))

        # Inject dynamic configurations if present
        if amm_config:
            try:
                scenario["amm"] = json.loads(amm_config)
            except Exception as e:
                logger.warning(f"Failed to parse amm_config JSON: {e}")
        if liquidity_config:
            try:
                scenario["liquidity_pool"] = json.loads(liquidity_config)
            except Exception as e:
                logger.warning(f"Failed to parse liquidity_config JSON: {e}")

        # Generate script
        script = generate_hardhat_script(scenario, contract_analysis=analysis)
        clean_script, validation = validate_and_sanitize(script)

        # Store/update job — ensure entry always exists regardless of source
        if not job_id:
            job_id = str(uuid.uuid4())[:8]

        # setdefault: creates the record if missing (e.g. after server restart)
        _stress_jobs.setdefault(job_id, {"status": "scripted", "analysis": analysis})

        _stress_jobs[job_id]["script"]   = clean_script
        _stress_jobs[job_id]["scenario"] = scenario
        _stress_jobs[job_id]["status"]   = "scripted"

        return {
            "job_id":     job_id,
            "script":     clean_script,
            "validation": validation,
            "scenario":   scenario,
        }
    except Exception as e:
        logger.exception("Script generation failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/script/run")
async def run_script(
    job_id: str = Form(...),
    script: str = Form(""),
):
    """
    Execute the stress-test script in the local Hardhat sandbox.
    Streams stdout/stderr. Returns risk report.
    """
    job = _stress_jobs.get(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")

    # Use provided script override or stored script
    exec_script = script.strip() or job.get("script", "")
    if not exec_script:
        raise HTTPException(status_code=400, detail="No script to execute")

    # Re-validate before execution
    clean_script, validation = validate_and_sanitize(exec_script)
    if not validation["valid"]:
        return JSONResponse(
            status_code=422,
            content={"error": "Script failed safety validation", "details": validation}
        )

    # Write script to a temp file inside the blockchain directory
    blockchain_dir = Path(__file__).parent / "blockchain"
    scripts_dir    = blockchain_dir / "scripts"
    scripts_dir.mkdir(exist_ok=True)

    script_path = scripts_dir / f"stress_{job_id}.js"
    script_path.write_text(clean_script, encoding="utf-8")

    _stress_jobs[job_id]["status"] = "running"

    import sys as _sys
    import subprocess
    import functools

    npx_cmd = "npx.cmd" if _sys.platform == "win32" else "npx"

    def _run_hardhat():
        """Blocking subprocess call — runs in a thread pool so the event loop stays free."""
        return subprocess.run(
            [npx_cmd, "hardhat", "run", str(script_path), "--network", "localhost"],
            cwd=str(blockchain_dir),
            capture_output=True,
            timeout=90,
        )

    try:
        loop = asyncio.get_event_loop()
        try:
            proc_result = await asyncio.wait_for(
                loop.run_in_executor(None, _run_hardhat),
                timeout=95,
            )
        except asyncio.TimeoutError:
            raise HTTPException(status_code=504, detail="Script execution timed out (90 s)")

        stdout    = proc_result.stdout.decode("utf-8", errors="replace")
        stderr    = proc_result.stderr.decode("utf-8", errors="replace")
        exit_code = proc_result.returncode

        # Build risk report
        scenario = job.get("scenario", {})
        analysis = job.get("analysis")
        report   = full_risk_report(stdout, scenario, analysis)

        result = {
            "job_id":    job_id,
            "exit_code": exit_code,
            "stdout":    stdout[-8000:],   # cap for payload size
            "stderr":    stderr[-2000:],
            "report":    report,
            "success":   exit_code == 0,
        }
        _stress_jobs[job_id]["results"] = result
        _stress_jobs[job_id]["status"]  = "done" if exit_code == 0 else "error"
        return result

    except HTTPException:
        raise
    except Exception as e:
        _stress_jobs[job_id]["status"] = "error"
        logger.exception("Hardhat execution failed")
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        # Clean up temp script
        try:
            script_path.unlink(missing_ok=True)
        except Exception:
            pass


@app.get("/api/script/results/{job_id}")
async def script_results(job_id: str):
    """Return stored risk results for a completed stress-test job."""
    job = _stress_jobs.get(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")
    return {
        "job_id":   job_id,
        "status":   job.get("status"),
        "results":  job.get("results"),
        "scenario": job.get("scenario"),
        "analysis": job.get("analysis"),
    }


# ─────────────────────────────────────────────────────────────────── #
# WEBSOCKET — STREAMING SIMULATION                                      #
# ─────────────────────────────────────────────────────────────────── #

@app.websocket("/ws/{run_id}")
async def websocket_run(websocket: WebSocket, run_id: str):
    await websocket.accept()
    run = _runs.get(run_id)
    if not run:
        await websocket.send_text(json.dumps({"type": "ERROR", "msg": "Run not found"}))
        await websocket.close()
        return

    total_steps = run["steps"]
    _runs[run_id]["status"] = "running"

    try:
        db_path = f"sim_{run_id}.db"

        # ── Parse Scenario -> Dynamic Configs ──────────────────────────────
        scenario_desc = run.get("prompt", "")
        amm_config = run.get("amm_config")
        liquidity_config = run.get("liquidity_config")

        if run.get("scenario_id"):
            base_scenario = get_scenario(run.get("scenario_id"))
            scenario_desc = base_scenario.get("description", "")
            
        # Parse text into scenario JSON if it's raw text
        if scenario_desc and not run.get("scenario_id"):
            parsed = parse_scenario(scenario_desc)
            scenario_desc = parsed.get("description", scenario_desc)
            if not amm_config and "amm" in parsed:
                amm_config = parsed["amm"]
            if not liquidity_config and "liquidity_pool" in parsed:
                liquidity_config = parsed["liquidity_pool"]
            
        dynamic_configs = {}
        if scenario_desc:
            dynamic_configs = generate_agent_configs({"description": scenario_desc})
            logger.info(f"Loaded dynamic configs for {len(dynamic_configs)} agents from NLP.")

        # ── Build protocol (blockchain-aware if enabled) ─────────────
        if _BLOCKCHAIN_READY:
            from MockDeFiProtocol import MockDeFiProtocol as _Mock
            from blockchain_client import BlockchainClient as _BC
            _mock = _Mock(amm_config=amm_config, liquidity_config=liquidity_config)
            _bc   = _BC(rpc_url=config.HARDHAT_RPC_URL)
            protocol = BlockchainAdapter(mock=_mock, blockchain=_bc)
            logger.info("Run %s: using BlockchainAdapter (on_chain=%s)", run_id, _bc.available)
        else:
            from MockDeFiProtocol import MockDeFiProtocol
            protocol = MockDeFiProtocol(amm_config=amm_config, liquidity_config=liquidity_config)
            logger.info("Run %s: using MockDeFiProtocol (no blockchain)", run_id)

        from database import SimulationLogger
        db = SimulationLogger(db_path=db_path)

        # ── Instantiate New Architecture Model ───────────────────────
        import model
        # Override the agent spawning with dynamic configs passed down
        # We need to temporarily mock protocol just for `DeFiSimulationModel`
        # However, `DeFiSimulationModel` instantiates its own `MockDeFiProtocol` and `SimulationLogger` inside `__init__`.
        # So we just instantiate it, and swap out the db and protocol.
        
        sim_model = DeFiSimulationModel(seed=config.RANDOM_SEED, db_path=db_path)
        
        # Override internally generated DB to use the one tied to this run_id
        sim_model.db.close() # close the default one created in __init__
        sim_model.db = db
        
        # Swap out protocol unconditionally
        sim_model.protocol = protocol
                
        # Inject dynamic configs and correct protocol/db references into the already spawned agents
        for agent in sim_model.schedule.agents:
            # Correct the pointers
            agent.db = db
            agent.protocol = protocol
                
            agent_class_name = agent.__class__.__name__
            if agent_class_name in dynamic_configs:
                agent.dynamic_config = dynamic_configs[agent_class_name]
                logger.info(f"Injected config into {agent_class_name}: {agent.dynamic_config}")

        # ── Step loop ────────────────────────────────────────────────
        for step_num in range(1, total_steps + 1):
            # Tell adapter what step we're on (for audit log)
            if hasattr(protocol, "set_step"):
                protocol.set_step(step_num)

            protocol.advance_oracle_price()
            pre_price = protocol.get_token_price()
            sim_model.schedule.step()

            liq          = protocol.get_pool_liquidity()
            amm_price    = protocol.get_token_price()
            oracle_price = protocol.get_oracle_price()
            spread_pct   = (amm_price - oracle_price) / oracle_price * 100 if oracle_price else 0.0
            price_delta  = (amm_price - pre_price) / pre_price * 100 if pre_price else 0.0

            step_events  = _get_step_events(db, step_num)
            event_type   = _classify_event(step_events, abs(spread_pct), abs(price_delta))

            crash_detail = None
            if event_type in ("CRASH", "WHALE"):
                crash_detail = _build_crash_detail(
                    step_events, liq, amm_price, oracle_price, spread_pct, price_delta
                )

            # Collect blockchain audit entries for this step
            bc_step_entries = []
            if hasattr(protocol, "audit_log"):
                bc_step_entries = [e for e in protocol.audit_log if e.get("step") == step_num]

            payload = {
                "type":         event_type,
                "step":         step_num,
                "total_steps":  total_steps,
                "progress":     step_num / total_steps,
                "pool_usdc":    round(liq["reserve_usdc"], 2),
                "pool_eth":     round(liq["reserve_eth"], 4),
                "amm_price":    round(amm_price, 2),
                "oracle_price": round(oracle_price, 2),
                "spread_pct":   round(spread_pct, 3),
                "price_delta":  round(price_delta, 3),
                "events":       step_events,
                "crash_detail": crash_detail,
                "blockchain": {
                    "tx_count": len(bc_step_entries),
                    "on_chain": sum(1 for e in bc_step_entries if e.get("on_chain")),
                    "entries":  bc_step_entries[:10],  # cap for payload size
                },
            }

            await websocket.send_text(json.dumps(payload))
            _runs[run_id]["events"].append(payload)

            # Collect audit trail entries into run store
            if bc_step_entries:
                _runs[run_id]["audit_log"].extend(bc_step_entries)

            await asyncio.sleep(0.18)

        # ── DONE ─────────────────────────────────────────────────────
        summary = db.get_summary()
        db.export_csv(config.CSV_PATH)
        db.close()

        # Final audit log snapshot
        if hasattr(protocol, "audit_log"):
            _runs[run_id]["audit_log"] = protocol.get_audit_log()

        await websocket.send_text(json.dumps({
            "type":       "DONE",
            "step":       total_steps,
            "summary":    summary,
            "pool_usdc":  round(protocol.get_pool_liquidity()["reserve_usdc"], 2),
            "pool_eth":   round(protocol.get_pool_liquidity()["reserve_eth"], 4),
            "amm_price":  round(protocol.get_token_price(), 2),
            "audit_count": len(_runs[run_id]["audit_log"]),
            "on_chain_count": sum(1 for e in _runs[run_id]["audit_log"] if e.get("on_chain")),
            "blockchain_enabled": _BLOCKCHAIN_READY,
        }))
        _runs[run_id]["status"] = "done"
        await websocket.close()

    except WebSocketDisconnect:
        logger.info("WebSocket disconnected for run %s", run_id)
        _runs[run_id]["status"] = "disconnected"
    except Exception as exc:
        logger.exception("Simulation error in run %s", run_id)
        _runs[run_id]["status"] = "error"
        try:
            await websocket.send_text(json.dumps({"type": "ERROR", "msg": str(exc)}))
        except Exception:
            pass


@app.websocket("/ws/stress/{run_id}")
async def websocket_stress_run(websocket: WebSocket, run_id: str):
    await websocket.accept()
    run = _runs.get(run_id)
    if not run:
        await websocket.send_text(json.dumps({"type": "ERROR", "msg": "Run not found"}))
        await websocket.close()
        return

    total_steps = run["steps"]
    _runs[run_id]["status"] = "running"
    
    flagged_agents = {}    # agent_type -> dict of issues
    lending_history: list = []   # one snapshot per step

    try:
        db_path = f"stress_{run_id}.db"

        scenario_desc = run.get("prompt", "")
        amm_config = run.get("amm_config")
        liquidity_config = run.get("liquidity_config")

        if run.get("scenario_id"):
            base_scenario = get_scenario(run.get("scenario_id"))
            scenario_desc = base_scenario.get("description", "")
            
        if scenario_desc and not run.get("scenario_id"):
            parsed = parse_scenario(scenario_desc)
            scenario_desc = parsed.get("description", scenario_desc)
            if not amm_config and "amm" in parsed:
                amm_config = parsed["amm"]
            if not liquidity_config and "liquidity_pool" in parsed:
                liquidity_config = parsed["liquidity_pool"]
            
        dynamic_configs = {}
        if scenario_desc:
            dynamic_configs = generate_agent_configs({"description": scenario_desc})
            logger.info("Loaded dynamic configs for stress test.")

        if _BLOCKCHAIN_READY:
            from MockDeFiProtocol import MockDeFiProtocol as _Mock
            from blockchain_client import BlockchainClient as _BC
            from blockchain_adapter import BlockchainAdapter
            _mock = _Mock(amm_config=amm_config, liquidity_config=liquidity_config)
            _bc   = _BC(rpc_url=config.HARDHAT_RPC_URL)
            protocol = BlockchainAdapter(mock=_mock, blockchain=_bc)
        else:
            from MockDeFiProtocol import MockDeFiProtocol
            protocol = MockDeFiProtocol(amm_config=amm_config, liquidity_config=liquidity_config)

        db = SimulationLogger(db_path=db_path)
        
        from model import DeFiSimulationModel
        sim_model = DeFiSimulationModel(seed=config.RANDOM_SEED, db_path=db_path)
        sim_model.db.close()
        sim_model.db = db
        sim_model.protocol = protocol
                
        for agent in sim_model.schedule.agents:
            agent.db = db
            agent.protocol = protocol
            agent_class_name = agent.__class__.__name__
            if agent_class_name in dynamic_configs:
                agent.dynamic_config = dynamic_configs[agent_class_name]

        for step_num in range(1, total_steps + 1):
            if hasattr(protocol, "set_step"):
                protocol.set_step(step_num)

            protocol.advance_oracle_price()
            pre_price = protocol.get_token_price()
            sim_model.schedule.step()

            liq          = protocol.get_pool_liquidity()
            amm_price    = protocol.get_token_price()
            oracle_price = protocol.get_oracle_price()
            spread_pct   = (amm_price - oracle_price) / oracle_price * 100 if oracle_price else 0.0
            price_delta  = (amm_price - pre_price) / pre_price * 100 if pre_price else 0.0
            lending_state = protocol.get_lending_state()
            lending_state["step"] = step_num
            lending_history.append(lending_state)

            # Mesa agents log their action BEFORE self.steps increments.
            step_events  = _get_step_events(db, step_num - 1)
            event_type   = _classify_event(step_events, abs(spread_pct), abs(price_delta))
            
            # --- Normalize ETH amounts to USDC ---
            for e in step_events:
                if e["token"] == "ETH":
                    e["amount"] = e["amount"] * oracle_price
                    e["token"] = "USDC (converted)"
            
            
            # --- FLAGGING LOGIC ---
            for e in step_events:
                a_type = e["agent_type"]
                act = e["action"]
                is_breaking = False
                reason = ""
                
                if act in ("LIQUIDATION", "MEV_LIQUIDATE_AND_DUMP", "BANK_RUN_WITHDRAWAL"):
                    is_breaking = True
                    reason = f"Executed critical protocol action: {act}."
                elif abs(price_delta) > 5.0 and act in ["WHALE_SWAP", "MEV_SANDWICH", "WHALE_LOOP_LEVERAGE"]:
                    is_breaking = True
                    reason = f"Caused massive price delta of {price_delta:.2f}% via {act}."
                elif abs(spread_pct) > 10.0:
                    is_breaking = True
                    reason = f"Drove AMM/Oracle spread to unsafe {spread_pct:.2f}%."
                
                if is_breaking:
                    if a_type not in flagged_agents:
                        flagged_agents[a_type] = {"type": a_type, "incidents": 1, "reasons": set([reason])}
                    else:
                        flagged_agents[a_type]["incidents"] += 1
                        flagged_agents[a_type]["reasons"].add(reason)

            # --- PER-STEP AGENT FLAGS (streamed to frontend for real-time matrix) ---
            pool_usdc_now = liq.get("reserve_usdc", 1) or 1
            agent_vel: dict = {}
            for e in step_events:
                agent_vel[e["agent_type"]] = agent_vel.get(e["agent_type"], 0) + 1

            step_agent_flags: dict = {}
            for a_type, vel in agent_vel.items():
                flags: list = []
                reasons: list = []
                if vel > 5:
                    flags.append("VELOCITY_BURST")
                    reasons.append(f"{vel} txns in step {step_num}")
                for ev in step_events:
                    if ev["agent_type"] != a_type:
                        continue
                    amt = ev.get("amount", 0) or 0
                    act = ev.get("action", "")
                    if amt > 0.20 * pool_usdc_now:
                        if "VOLUME_SPIKE" not in flags:
                            flags.append("VOLUME_SPIKE")
                            reasons.append(f"${amt:,.0f} > 20% of pool (${pool_usdc_now:,.0f})")
                    if act in ("LIQUIDATION", "MEV_LIQUIDATE_AND_DUMP", "BANK_RUN_WITHDRAWAL"):
                        if "DANGEROUS_ACTION" not in flags:
                            flags.append("DANGEROUS_ACTION")
                            reasons.append(f"Executed {act} at step {step_num}")
                    if act in ("WHALE_LOOP_LEVERAGE", "WHALE_SWAP", "MEV_SANDWICH") and abs(price_delta) > 5.0:
                        if "PRICE_MANIP" not in flags:
                            flags.append("PRICE_MANIP")
                            reasons.append(f"Price delta {price_delta:+.2f}% via {act}")
                if abs(spread_pct) > 10.0:
                    if "SPREAD_MANIP" not in flags:
                        flags.append("SPREAD_MANIP")
                        reasons.append(f"Spread {spread_pct:.2f}% at step {step_num}")
                step_agent_flags[a_type] = {"velocity": vel, "flags": flags, "reasons": reasons}

            crash_detail = None
            if event_type in ("CRASH", "WHALE"):
                crash_detail = _build_crash_detail(step_events, liq, amm_price, oracle_price, spread_pct, price_delta)

            bc_step_entries = []
            if hasattr(protocol, "audit_log"):
                bc_step_entries = [e for e in protocol.audit_log if e.get("step") == step_num]

            payload = {
                "type":         event_type,
                "step":         step_num,
                "total_steps":  total_steps,
                "progress":     step_num / total_steps,
                "pool_usdc":    round(liq["reserve_usdc"], 2),
                "pool_eth":     round(liq["reserve_eth"], 4),
                "amm_price":    round(amm_price, 2),
                "oracle_price": round(oracle_price, 2),
                "spread_pct":   round(spread_pct, 3),
                "price_delta":  round(price_delta, 3),
                "events":       step_events,
                "agent_flags":  step_agent_flags,
                "lending":      lending_state,
                "crash_detail": crash_detail,
                "blockchain": {
                    "tx_count": len(bc_step_entries),
                    "on_chain": sum(1 for e in bc_step_entries if e.get("on_chain")),
                    "entries":  bc_step_entries[:10],
                },
            }

            await websocket.send_text(json.dumps(payload))
            _runs[run_id]["events"].append(payload)
            if bc_step_entries:
                _runs[run_id]["audit_log"].extend(bc_step_entries)

            await asyncio.sleep(0.18)

        # ── DONE WITH STEPS — GENERATE AUDIT ──────────────────────────────────────────────
        summary = db.get_summary()
        db.export_csv(config.CSV_PATH)
        db.close()

        if hasattr(protocol, "audit_log"):
            _runs[run_id]["audit_log"] = protocol.get_audit_log()
            
        for k in flagged_agents:
            flagged_agents[k]["reasons"] = list(flagged_agents[k]["reasons"])
        flagged_list = list(flagged_agents.values())

        # ── LENDING ANALYTICS ────────────────────────────────────────────────────────
        lending_analytics: dict = {}
        if lending_history:
            lending_analytics = {
                "peak_debt_usdc":            max(s.get("total_debt_usdc", 0)      for s in lending_history),
                "peak_collateral_eth":        max(s.get("total_collateral_eth", 0) for s in lending_history),
                "min_health_factor_seen":     min(s.get("min_health_factor", 9999) for s in lending_history),
                "avg_health_factor_final":    lending_history[-1].get("avg_health_factor", 9999),
                "peak_liquidatable_count":    max(s.get("liquidatable_count", 0)   for s in lending_history),
                "peak_bad_debt_usdc":         max(s.get("bad_debt_usdc", 0)        for s in lending_history),
                "bad_debt_steps":             sum(1 for s in lending_history if s.get("has_bad_debt")),
                "peak_utilization_rate":      round(max(s.get("utilization_rate", 0) for s in lending_history), 4),
                "active_borrowers_peak":      max(s.get("active_borrower_count", 0) for s in lending_history),
                "total_liquidations":         sum(
                    1 for ev in _runs[run_id]["events"]
                    for e in ev.get("events", [])
                    if e.get("action") in ("MEV_LIQUIDATE_AND_DUMP", "LIQUIDATION")
                ),
                "bank_run_count":             sum(
                    1 for ev in _runs[run_id]["events"]
                    for e in ev.get("events", [])
                    if e.get("action") == "BANK_RUN_WITHDRAWAL"
                ),
                "total_whale_leverage_loops": sum(
                    1 for ev in _runs[run_id]["events"]
                    for e in ev.get("events", [])
                    if e.get("action") == "WHALE_LOOP_LEVERAGE"
                ),
                "ltv_max":          0.80,
                "liq_threshold":    0.85,
                "liq_bonus":        0.05,
            }

        from audit_generator import generate_audit_report

        # Read config.py for LLM context instead of user theoretical inputs
        config_content = ""
        try:
            config_path = Path(__file__).parent / "config.py"
            with open(config_path, "r", encoding="utf-8") as f:
                # Strip comments and empty lines to save tokens
                config_content = "\n".join([line.split("#")[0].strip() for line in f if line.strip() and not line.strip().startswith("#")])
        except Exception as e:
            logger.warning("Failed to read config.py for audit generation: %s", e)

        # Compile a timeline of events for context, highlighting key events
        sim_events = []
        for e in _runs[run_id]["events"]:
            # Only include steps that had notable interactions or events
            if e.get("events") or e.get("type", "STEP") != "STEP":
                step_data = {
                    "step": e["step"],
                    "event_type": e.get("type"),
                    "amm_price": e.get("amm_price"),
                    "oracle_price": e.get("oracle_price"),
                    "spread_pct": e.get("spread_pct"),
                    "price_delta": e.get("price_delta"),
                }
                
                # Aggregate actions by agent type and action to save tokens
                from collections import Counter
                act_counts = Counter(f"{act.get('agent_type')}_{act.get('action')}" for act in e.get("events", []))
                step_data["actions_summary"] = dict(act_counts)
                    
                sim_events.append(step_data)
                
        # Limit to most volatile steps if too large
        if len(sim_events) > 12:
            sim_events = sorted(sim_events, key=lambda x: abs(x["price_delta"]) + abs(x["spread_pct"]), reverse=True)[:12]
            sim_events.sort(key=lambda x: x["step"])

        markdown_audit = generate_audit_report(
            config_content=config_content,
            agent_configs=dynamic_configs,
            flagged_agents=flagged_list,
            simulation_summary=summary,
            simulation_events=sim_events,
            lending_analytics=lending_analytics,
        )

        done_payload = {
            "type":       "DONE",
            "step":       total_steps,
            "summary":    summary,
            "pool_usdc":  round(protocol.get_pool_liquidity()["reserve_usdc"], 2),
            "pool_eth":   round(protocol.get_pool_liquidity()["reserve_eth"], 4),
            "amm_price":  round(protocol.get_token_price(), 2),
            "audit_count": len(_runs[run_id]["audit_log"]),
            "on_chain_count": sum(1 for e in _runs[run_id]["audit_log"] if e.get("on_chain")),
            "blockchain_enabled": _BLOCKCHAIN_READY,
            "markdown_audit": markdown_audit
        }
        await websocket.send_text(json.dumps(done_payload))
        # Persist for Ask Arpit context
        _last_sim_result["markdown_audit"] = markdown_audit
        _last_sim_result["summary"] = summary
        _runs[run_id]["status"] = "done"
        await websocket.close()

    except WebSocketDisconnect:
        logger.info("WebSocket disconnected for stress run %s", run_id)
        _runs[run_id]["status"] = "disconnected"
    except Exception as exc:
        logger.exception("Simulation error in stress run %s", run_id)
        _runs[run_id]["status"] = "error"
        try:
            await websocket.send_text(json.dumps({"type": "ERROR", "msg": str(exc)}))
        except Exception:
            pass


# ─────────────────────────────────────────────────────────────────── #
# HELPERS                                                              #
# ─────────────────────────────────────────────────────────────────── #

def _get_step_events(db: SimulationLogger, step: int):
    import sqlite3
    conn = sqlite3.connect(db._db_path)
    cur  = conn.execute(
        "SELECT agent_type, action, amount, token, reasoning "
        "FROM agent_actions WHERE step=? ORDER BY id",
        (step,)
    )
    rows = [
        {"agent_type": r[0], "action": r[1],
         "amount": r[2], "token": r[3], "reasoning": r[4]}
        for r in cur.fetchall()
    ]
    conn.close()
    return rows


def _classify_event(events, abs_spread_pct: float, abs_price_delta: float) -> str:
    actions = [e["action"] for e in events]
    if "LIQUIDATION" in actions or "MEV_LIQUIDATE_AND_DUMP" in actions or "BANK_RUN_WITHDRAWAL" in actions:
        return "CRASH"
    if "WHALE_SWAP" in actions or "WHALE_LOOP_LEVERAGE" in actions or abs_price_delta > 8.0:
        return "WHALE"
    if any(a.startswith("MEV_") for a in actions):
        return "MEV"
    if "ARBITRAGE_SWAP" in actions or abs_spread_pct > 2.0:
        return "ARB"
    return "STEP"


def _build_crash_detail(events, liq, amm_price, oracle_price, spread_pct, price_delta) -> dict:
    liquidations = [e for e in events if e["action"] in ("LIQUIDATION", "MEV_LIQUIDATE_AND_DUMP")]
    whale_swaps  = [e for e in events if e["action"] in ("WHALE_SWAP", "WHALE_LOOP_LEVERAGE")]
    mev_events   = [e for e in events if e.get("action", "").startswith("MEV_")]
    bank_runs    = [e for e in events if e.get("action") == "BANK_RUN_WITHDRAWAL"]

    cause = "Unknown"
    if bank_runs:
        cause = f"PROTOCOL BANK RUN — {len(bank_runs)} major withdrawals occurred"
    elif liquidations:
        cause = f"HEALTH FACTOR COLLAPSE — {len(liquidations)} position(s) liquidated"
    elif whale_swaps:
        cause = f"WHALE MARKET MANIPULATION — price moved {price_delta:+.2f}%"

    mev_note = f" | MEV extractions: {len(mev_events)}" if mev_events else ""

    return {
        "cause":        cause + mev_note,
        "price_impact": f"{price_delta:+.2f}%",
        "spread":       f"{spread_pct:+.3f}%",
        "pool_usdc":    f"{liq['reserve_usdc']:,.2f}",
        "pool_eth":     f"{liq['reserve_eth']:.4f}",
        "amm_price":    f"${amm_price:,.2f}",
        "oracle_price": f"${oracle_price:,.2f}",
        "liquidations": len(liquidations),
        "reasoning":    bank_runs[0]["reasoning"] if bank_runs else (
                        liquidations[0]["reasoning"] if liquidations else (
                        whale_swaps[0]["reasoning"] if whale_swaps else (
                        mev_events[0]["reasoning"] if mev_events else ""))),
    }


# ─────────────────────────────────────────────────────────────────── #
# ENTRY POINT                                                          #
# ─────────────────────────────────────────────────────────────────── #
if __name__ == "__main__":
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=True)