#!/usr/bin/env python3
import argparse
import csv
import json
import math
import os
import random
import signal
import sqlite3
import statistics
import subprocess
import tempfile
import time
import urllib.error
import urllib.request


AGENTS = [f"agent_{i}" for i in range(8)]

WORLD_CONFIG = {
    "num_agents": 8,
    "num_cores": 4,
    "genesis_tokens_per_agent": 1000,
    "commons_threshold_per_turn": 100,
    "base_inference_rate": 1,
    "thinking_layer_discount": 0.1,
    "mine_base_weight": 10.0,
    "stake_weight_per_token": 0.01,
    "burn_weight_per_token": 0.05,
    "burn_decay_rate": 0.02,
    "burn_maturity_turns": 3,
    "unstake_delay_turns": 5,
    "interest_rate_per_turn": 0.01,
    "signing_bonus": 50,
    "block_threshold": 20.0,
    "attested_confirmation_window": 3,
    "slash_both_on_timeout": True,
}


def http_json(url: str, method: str = "GET", payload=None, timeout=30):
    data = None
    headers = {"Content-Type": "application/json"}
    if payload is not None:
        data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read().decode("utf-8"))


def wait_engine(base_url: str, timeout_s: float = 30.0):
    start = time.time()
    while time.time() - start < timeout_s:
        try:
            http_json(f"{base_url}/config")
            return
        except Exception:
            time.sleep(0.2)
    raise RuntimeError("engine did not become ready")


def gini(values):
    xs = sorted(values)
    n = len(xs)
    if n == 0:
        return 0.0
    total = sum(xs)
    if total == 0:
        return 0.0
    weighted = 0.0
    for i, x in enumerate(xs, start=1):
        weighted += i * x
    return (2.0 * weighted) / (n * total) - (n + 1) / n


def shifted_nonnegative(values):
    m = min(values)
    if m <= 0:
        shift = 1 - m
        return [x + shift for x in values], shift
    return list(values), 0


def choose_actions(turn, state, rng):
    agents = state["agents"]
    by_id = {a["agent_id"]: a for a in agents}
    poorest = min(agents, key=lambda a: a["balance"])["agent_id"]

    inputs = []
    for aid in AGENTS:
        a = by_id[aid]
        balance = int(a.get("balance", 0))
        staked = int(a.get("staked", 0))

        action = {"action": "mine"}
        idx = int(aid.split("_")[1])

        if balance < -180:
            action = {"action": "job"}
        elif staked > 0 and turn % 15 == 0:
            action = {"action": "unstake"}
        elif balance > 360 and staked < 220 and (turn + idx) % 7 == 0:
            action = {"action": "stake", "amount": 80}
        elif balance > 600 and (turn + idx) % 11 == 0:
            action = {"action": "burn", "amount": 50}
        elif balance > 450 and turn % 13 == 0 and poorest != aid:
            action = {"action": "transfer", "to": poorest, "amount": 20, "fee": 1}
        elif balance > 250 and turn % 10 == 0:
            action = {"action": "study"}

        thinking_units = rng.randint(50, 350)
        output_units = rng.randint(40, 280)

        inputs.append(
            {
                "agent_id": aid,
                "thinking": "",
                "action": action,
                "speech": None,
                "thinking_units": thinking_units,
                "output_units": output_units,
            }
        )

    return inputs


def query_sqlite_metrics(db_path: str):
    conn = sqlite3.connect(db_path)
    cur = conn.cursor()

    cur.execute("SELECT COUNT(*) FROM blocks WHERE turn > 0")
    block_count = cur.fetchone()[0]

    cur.execute(
        """
        SELECT json_extract(data, '$.validator_type'), COUNT(*)
        FROM blocks
        WHERE turn > 0
        GROUP BY json_extract(data, '$.validator_type')
        """
    )
    validator_counts = {row[0]: row[1] for row in cur.fetchall()}

    cur.execute(
        """
        SELECT json_extract(data, '$.tx_type'), COUNT(*)
        FROM transactions
        GROUP BY json_extract(data, '$.tx_type')
        """
    )
    tx_counts = {row[0]: row[1] for row in cur.fetchall()}

    conn.close()
    return block_count, validator_counts, tx_counts


def run_one(run_id: int, turns: int, engine_bin: str):
    rng = random.Random(1000 + run_id)
    port = 3100 + run_id
    base = f"http://127.0.0.1:{port}"

    with tempfile.TemporaryDirectory(prefix=f"sim_run_{run_id}_") as tmpd:
        db_path = os.path.join(tmpd, "sim.db")
        env = os.environ.copy()
        env["DB_PATH"] = db_path
        env["PORT"] = str(port)

        proc = subprocess.Popen(
            [engine_bin],
            env=env,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        try:
            wait_engine(base)

            http_json(
                f"{base}/init",
                method="POST",
                payload={"config": WORLD_CONFIG, "agent_ids": AGENTS},
            )

            total_fees = 0
            turns_with_block = 0
            contracts_settled = 0
            contracts_defaulted = 0
            action_counts = {
                "mine": 0,
                "stake": 0,
                "unstake": 0,
                "burn": 0,
                "study": 0,
                "job": 0,
                "transfer": 0,
            }

            for turn in range(1, turns + 1):
                state = http_json(f"{base}/state")
                inputs = choose_actions(turn, state, rng)
                for it in inputs:
                    name = it["action"]["action"]
                    if name in action_counts:
                        action_counts[name] += 1
                out = http_json(f"{base}/turn", method="POST", payload={"inputs": inputs})
                data = out.get("data", {}) if isinstance(out, dict) else {}

                total_fees += int(data.get("inference_fees_collected", 0) or 0)
                if data.get("block_winner"):
                    turns_with_block += 1
                contracts_settled += len(data.get("contracts_settled", []) or [])
                contracts_defaulted += len(data.get("contracts_defaulted", []) or [])

            final_state = http_json(f"{base}/state")
            balances = [int(a["balance"]) for a in final_state["agents"]]
            staked = [int(a["staked"]) for a in final_state["agents"]]
            wealth = [b + s for b, s in zip(balances, staked)]
            total_wealth = sum(wealth)
            top1_share = max(wealth) / total_wealth if total_wealth > 0 else 0.0
            wealth_shifted, shift_used = shifted_nonnegative(wealth)
            top1_shifted_share = max(wealth_shifted) / sum(wealth_shifted)

            block_count_db, validator_counts, tx_counts = query_sqlite_metrics(db_path)

            return {
                "run_id": run_id,
                "turns": turns,
                "token_supply_final": int(final_state.get("token_supply", 0)),
                "fees_total": int(total_fees),
                "turns_with_block": int(turns_with_block),
                "block_count_db": int(block_count_db),
                "settled_contracts": int(contracts_settled),
                "defaulted_contracts": int(contracts_defaulted),
                "mean_balance_final": statistics.mean(balances),
                "median_balance_final": statistics.median(balances),
                "min_balance_final": min(balances),
                "max_balance_final": max(balances),
                "negative_balance_agents": sum(1 for b in balances if b < 0),
                "gini_wealth_final": gini(wealth),
                "top1_wealth_share_final": top1_share,
                "wealth_shift_used": float(shift_used),
                "gini_wealth_shifted": gini(wealth_shifted),
                "top1_wealth_shifted_share_final": top1_shifted_share,
                "validator_mine": int(validator_counts.get("mine", 0)),
                "validator_stake": int(validator_counts.get("stake", 0)),
                "validator_burn": int(validator_counts.get("burn", 0)),
                "tx_transfer": int(tx_counts.get("transfer", 0)),
                "tx_stake": int(tx_counts.get("stake", 0)),
                "tx_unstake": int(tx_counts.get("unstake", 0)),
                "tx_burn": int(tx_counts.get("burn", 0)),
                "tx_study": int(tx_counts.get("study", 0)),
                "tx_job": int(tx_counts.get("job", 0)),
                "actions_mine": int(action_counts["mine"]),
                "actions_stake": int(action_counts["stake"]),
                "actions_unstake": int(action_counts["unstake"]),
                "actions_burn": int(action_counts["burn"]),
                "actions_study": int(action_counts["study"]),
                "actions_job": int(action_counts["job"]),
                "actions_transfer": int(action_counts["transfer"]),
            }
        finally:
            try:
                proc.send_signal(signal.SIGTERM)
                proc.wait(timeout=3)
            except Exception:
                proc.kill()


def summarize(rows):
    numeric_keys = [
        k
        for k in rows[0].keys()
        if k not in {"run_id"}
        and isinstance(rows[0][k], (int, float))
    ]
    out = {}
    for k in numeric_keys:
        xs = [float(r[k]) for r in rows]
        mean = statistics.mean(xs)
        sd = statistics.stdev(xs) if len(xs) > 1 else 0.0
        se = sd / math.sqrt(len(xs)) if len(xs) > 0 else 0.0
        ci95 = 1.96 * se
        out[k] = {
            "mean": mean,
            "std": sd,
            "ci95": ci95,
            "min": min(xs),
            "max": max(xs),
        }
    return out


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--runs", type=int, default=20)
    parser.add_argument("--turns", type=int, default=50)
    parser.add_argument(
        "--engine-bin",
        default="sim-engine/target/release/sim-engine",
    )
    parser.add_argument("--out-dir", default="paper/results")
    args = parser.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)

    rows = []
    for run_id in range(1, args.runs + 1):
        row = run_one(run_id, args.turns, args.engine_bin)
        rows.append(row)
        print(
            f"run {run_id:02d}/{args.runs} | supply={row['token_supply_final']} | "
            f"gini={row['gini_wealth_final']:.3f} | blocks={row['turns_with_block']}"
        )

    csv_path = os.path.join(args.out_dir, "run_metrics.csv")
    with open(csv_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)

    summary = {
        "runs": args.runs,
        "turns": args.turns,
        "world_config": WORLD_CONFIG,
        "aggregate": summarize(rows),
    }
    summary_path = os.path.join(args.out_dir, "summary.json")
    with open(summary_path, "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)

    print(f"wrote {csv_path}")
    print(f"wrote {summary_path}")


if __name__ == "__main__":
    main()