#!/usr/bin/env python3
from __future__ import annotations

import argparse
import hashlib
import json
import os
import random
import signal
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
import platform

import tomli_w
from supabase import create_client
from rc.credits import parse_credit_amount_to_subunits
from rc.inference.provision import ensure_model_downloaded

if sys.version_info >= (3, 11):
    import tomllib
else:
    import tomli as tomllib


PROMPTS = (
    "Prove in Lean 4 that for every proposition P, P implies P.",
    "Prove in Lean 4 that addition on natural numbers is associative.",
    "Prove in Lean 4 that addition on natural numbers is commutative.",
    "Show in Lean 4 that reversing a list twice returns the original list.",
    "Prove in Lean 4 that mapping the identity function over a list returns the same list.",
    "Show in Lean 4 that length (xs ++ ys) = length xs + length ys.",
    "Prove in Lean 4 that if a <= b and b <= c then a <= c for natural numbers.",
    "Show in Lean 4 that n < n + 1 for every natural number n.",
    "Prove in Lean 4 that multiplication distributes over addition on the right.",
    "Prove in Lean 4 that multiplication distributes over addition on the left.",
    "Show in Lean 4 that True implies True.",
    "Prove in Lean 4 that Nat.succ is injective.",
)


def default_provider_model() -> str:
    """Use the smallest supported DeepSeek R1 Distill Qwen tier for soak runs."""
    if platform.system() == "Darwin" and platform.machine() == "arm64":
        return "mlx-community/DeepSeek-R1-Distill-Qwen-1.5B-4bit"
    return "bartowski/DeepSeek-R1-Distill-Qwen-1.5B-GGUF"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run an overnight Supabase-backed soak test.")
    parser.add_argument(
        "--repo",
        default=str(Path(__file__).resolve().parents[1]),
        help="Repository root containing .venv and rc sources.",
    )
    parser.add_argument(
        "--requester-config",
        default=str(Path.home() / ".rc" / "config.toml"),
        help="Path to the funded requester config.",
    )
    parser.add_argument(
        "--run-dir",
        default=f"/tmp/rc-overnight-soak-{datetime.now(timezone.utc).strftime('%Y%m%d-%H%M%S')}",
        help="Directory for provider homes and soak logs.",
    )
    parser.add_argument("--provider-count", type=int, default=10)
    parser.add_argument("--duration-hours", type=float, default=8.0)
    parser.add_argument("--queue-depth", type=int, default=12)
    parser.add_argument("--task-budget", default="1")
    parser.add_argument("--submit-interval", type=int, default=30)
    parser.add_argument("--monitor-interval", type=int, default=60)
    parser.add_argument("--poll-interval", type=int, default=1)
    parser.add_argument("--max-tokens", type=int, default=48)
    parser.add_argument("--temperature", type=float, default=0.2)
    parser.add_argument("--threads", type=int, default=1)
    parser.add_argument("--ctx-size", type=int, default=2048)
    parser.add_argument("--lease-seconds", type=int, default=60)
    parser.add_argument("--keepalive-interval", type=int, default=10)
    parser.add_argument("--gpu-layers", type=int, default=0)
    parser.add_argument("--provider-model", default=default_provider_model())
    parser.add_argument(
        "--research-mode",
        action="store_true",
        help="Submit autonomous research episodes plus proof-attempt jobs instead of plain queue tasks.",
    )
    parser.add_argument(
        "--research-domain",
        default="naturals",
        help="Domain label stored on research episodes when --research-mode is enabled.",
    )
    parser.add_argument(
        "--provider-model-path",
        default="",
        help="Optional local model path. Defaults to the requester config model_path or the managed local cache.",
    )
    return parser.parse_args()


def load_toml(path: Path) -> dict:
    with path.open("rb") as f:
        return tomllib.load(f)


def dump_toml(path: Path, data: dict) -> None:
    with path.open("wb") as f:
        tomli_w.dump(data, f)


def sha256_text(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


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


def append_jsonl(path: Path, row: dict) -> None:
    with path.open("a", encoding="utf-8") as f:
        f.write(json.dumps(row, sort_keys=True) + "\n")


class SoakRunner:
    def __init__(self, args: argparse.Namespace) -> None:
        self.args = args
        self.repo = Path(args.repo).resolve()
        self.requester_config_path = Path(args.requester_config).expanduser().resolve()
        self.requester_cfg = load_toml(self.requester_config_path)
        self.client = create_client(
            self.requester_cfg["supabase"]["url"],
            self.requester_cfg["supabase"]["key"],
        )
        self.run_dir = Path(args.run_dir).expanduser().resolve()
        self.logs_dir = self.run_dir / "logs"
        self.logs_dir.mkdir(parents=True, exist_ok=True)
        self.metrics_path = self.run_dir / "metrics.jsonl"
        self.submissions_path = self.run_dir / "submissions.jsonl"
        self.summary_path = self.run_dir / "summary.json"
        self.provider_processes: list[subprocess.Popen[bytes]] = []
        self.stop_requested = False
        self.last_submit_at = 0.0
        self.last_monitor_at = 0.0
        self.requester_name = self.requester_cfg["node"]["name"]
        self.requester_node_id = self.requester_cfg["node"]["id"]
        self.requester_secret = self.requester_cfg["node"]["secret"]
        self.task_budget = parse_credit_amount_to_subunits(args.task_budget)
        configured_model_path = self.requester_cfg.get("inference", {}).get("model_path", "")
        self.provider_model_path = str(
            ensure_model_downloaded(
                args.provider_model,
                args.provider_model_path or configured_model_path,
            )
        )

    def install_signal_handlers(self) -> None:
        def _handle_signal(signum: int, frame: object) -> None:
            self.stop_requested = True
            self.log_event("signal", {"signum": signum})

        signal.signal(signal.SIGINT, _handle_signal)
        signal.signal(signal.SIGTERM, _handle_signal)

    def log_event(self, event_type: str, payload: dict) -> None:
        append_jsonl(
            self.metrics_path,
            {
                "ts": utc_now(),
                "event": event_type,
                **payload,
            },
        )

    def requester_auth_payload(self) -> dict[str, str]:
        return {
            "p_node_id": self.requester_node_id,
            "p_node_secret_hash": sha256_text(self.requester_secret),
        }

    def get_summary(self) -> dict:
        return self.client.rpc("get_public_network_summary", {}).execute().data

    def get_public_nodes(self) -> list[dict]:
        return self.client.rpc("get_public_nodes", {"p_limit": 50}).execute().data or []

    def get_balance(self) -> dict:
        rows = self.client.rpc(
            "get_balance",
            self.requester_auth_payload(),
        ).execute().data or []
        if not rows:
            self.ensure_requester_registered()
            rows = self.client.rpc(
                "get_balance",
                self.requester_auth_payload(),
            ).execute().data or []
        if not rows:
            raise RuntimeError(
                "No requester balance row found after registration. "
                f"Check requester config {self.requester_config_path} and Supabase runtime state."
            )
        return rows[0]

    def ensure_requester_registered(self) -> None:
        self.client.rpc(
            "register_node",
            {
                "p_id": self.requester_node_id,
                "p_name": self.requester_name,
                "p_role": self.requester_cfg["node"]["role"],
                "p_public_key": "",
                "p_auth_token_hash": sha256_text(self.requester_secret),
                "p_gpu_info": self.requester_cfg.get("gpu", {}) or {},
            },
        ).execute()

    def submit_task(self, prompt: str, budget: int) -> str:
        if self.args.research_mode:
            return self.submit_research_attempt(prompt, budget)
        task_id = self.client.rpc(
            "submit_task",
            {
                **self.requester_auth_payload(),
                "p_prompt": prompt,
                "p_model": self.args.provider_model,
                "p_budget": budget,
                "p_task_type": "lean_proof_gen",
                "p_research_episode_id": None,
                "p_attempt_index": 0,
                "p_parent_attempt_id": None,
                "p_task_metadata": {},
            },
        ).execute().data
        if not task_id:
            raise RuntimeError("submit_task did not return a task id")
        append_jsonl(
            self.submissions_path,
            {
                "ts": utc_now(),
                "task_id": task_id,
                "budget": budget,
                "prompt": prompt,
            },
        )
        return task_id

    def submit_research_attempt(self, prompt: str, budget: int) -> str:
        episode_id = self.client.rpc(
            "create_research_episode",
            {
                **self.requester_auth_payload(),
                "p_conjecture_text": prompt,
                "p_formal_statement": "",
                "p_domain": self.args.research_domain,
                "p_strategy_tag": "soak",
                "p_metadata": {"source": "overnight_soak"},
                "p_interestingness_score": 5.0,
                "p_novelty_score": 5.0,
                "p_max_attempts": 3,
                "p_budget_per_attempt": budget,
            },
        ).execute().data
        payload = self.client.rpc(
            "enqueue_research_attempt",
            {
                "p_episode_id": episode_id,
                **self.requester_auth_payload(),
                "p_prompt": prompt,
                "p_model": self.args.provider_model,
                "p_budget": budget,
                "p_task_type": "proof_attempt",
                "p_strategy_tag": "soak",
                "p_metadata": {
                    "domain": self.args.research_domain,
                    "conjecture_family": "soak",
                    "interestingness_score": 5.0,
                    "novelty_score": 5.0,
                },
            },
        ).execute().data
        task_id = payload["task_id"]
        append_jsonl(
            self.submissions_path,
            {
                "ts": utc_now(),
                "episode_id": episode_id,
                "task_id": task_id,
                "budget": budget,
                "prompt": prompt,
                "research_mode": True,
            },
        )
        return task_id

    def ensure_clean_queue(self) -> None:
        summary = self.get_summary()["tasks"]
        active = summary["pending"] + summary["running"] + summary["assigned"]
        if active:
            raise RuntimeError(
                f"Queue must be empty before starting soak test. "
                f"Found pending={summary['pending']} running={summary['running']} assigned={summary['assigned']}."
            )

    def prepare_providers(self) -> None:
        rc_bin = self.repo / ".venv" / "bin" / "rc"
        python_bin = self.repo / ".venv" / "bin" / "python"
        for idx in range(1, self.args.provider_count + 1):
            home = self.run_dir / f"provider-{idx:02d}"
            env = os.environ.copy()
            env["HOME"] = str(home)
            cmd = [
                str(rc_bin),
                "init",
                "--name",
                f"soak-provider-{idx:02d}",
                "--role",
                "provider",
                "--supabase-url",
                self.requester_cfg["supabase"]["url"],
                "--supabase-key",
                self.requester_cfg["supabase"]["key"],
                "--model",
                self.args.provider_model,
                "--skip-model-download",
            ]
            result = subprocess.run(
                cmd,
                cwd=self.repo,
                env=env,
                capture_output=True,
                text=True,
            )
            if result.returncode != 0:
                raise RuntimeError(
                    f"Provider init failed for {home}:\n{result.stdout}\n{result.stderr}"
                )

            cfg_path = home / ".rc" / "config.toml"
            cfg = load_toml(cfg_path)
            cfg["inference"]["model"] = self.args.provider_model
            cfg["inference"]["model_path"] = self.provider_model_path
            cfg["inference"]["model_downloaded"] = True
            dump_toml(cfg_path, cfg)

            subprocess.run(
                [str(python_bin), "-c", "print('provider-ready')"],
                cwd=self.repo,
                env=env,
                check=True,
                capture_output=True,
                text=True,
            )

    def start_providers(self) -> None:
        rc_bin = self.repo / ".venv" / "bin" / "rc"
        for idx in range(1, self.args.provider_count + 1):
            home = self.run_dir / f"provider-{idx:02d}"
            env = os.environ.copy()
            env["HOME"] = str(home)
            env["RC_N_GPU_LAYERS"] = str(self.args.gpu_layers)
            log_path = self.logs_dir / f"provider-{idx:02d}.log"
            log_file = log_path.open("ab")
            proc = subprocess.Popen(
                [
                    str(rc_bin),
                    "serve",
                    "--poll-interval",
                    str(self.args.poll_interval),
                    "--max-tokens",
                    str(self.args.max_tokens),
                    "--temperature",
                    str(self.args.temperature),
                    "--threads",
                    str(self.args.threads),
                    "--ctx-size",
                    str(self.args.ctx_size),
                    "--lease-seconds",
                    str(self.args.lease_seconds),
                    "--keepalive-interval",
                    str(self.args.keepalive_interval),
                ],
                cwd=self.repo,
                env=env,
                stdout=log_file,
                stderr=subprocess.STDOUT,
            )
            self.provider_processes.append(proc)
            self.log_event(
                "provider_started",
                {
                    "provider_index": idx,
                    "pid": proc.pid,
                    "log_path": str(log_path),
                },
            )

    def monitor(self) -> None:
        summary = self.get_summary()
        balance = self.get_balance()
        nodes = self.get_public_nodes()
        soak_nodes = [
            node for node in nodes if str(node.get("name", "")).startswith("soak-provider-")
        ]
        healthy = sum(1 for node in soak_nodes if node.get("health") == "healthy")
        online = sum(1 for node in soak_nodes if node.get("status") == "online")
        snapshot = {
            "summary": summary,
            "balance": balance,
            "soak_provider_count": len(soak_nodes),
            "soak_provider_healthy": healthy,
            "soak_provider_online": online,
        }
        self.summary_path.write_text(json.dumps(snapshot, indent=2), encoding="utf-8")
        self.log_event("monitor", snapshot)

    def maybe_submit_work(self) -> None:
        summary = self.get_summary()["tasks"]
        active = summary["pending"] + summary["running"] + summary["assigned"]
        balance = self.get_balance()
        available = int(balance["balance"])
        target = min(self.args.queue_depth, max(available // self.task_budget, 0))

        while active < target and available >= self.task_budget and not self.stop_requested:
            prompt = random.choice(PROMPTS)
            task_id = self.submit_task(prompt, self.task_budget)
            self.log_event(
                "task_submitted",
                {
                    "task_id": task_id,
                    "budget": self.task_budget,
                    "active_before_submit": active,
                },
            )
            active += 1
            available -= self.task_budget

    def sleep_interruptible(self, seconds: int) -> None:
        deadline = time.time() + seconds
        while time.time() < deadline and not self.stop_requested:
            time.sleep(1)

    def shutdown(self) -> None:
        for proc in self.provider_processes:
            if proc.poll() is None:
                proc.send_signal(signal.SIGINT)
        deadline = time.time() + 30
        for proc in self.provider_processes:
            remaining = max(0.0, deadline - time.time())
            try:
                proc.wait(timeout=remaining)
            except subprocess.TimeoutExpired:
                proc.kill()
        self.log_event("shutdown", {"provider_count": len(self.provider_processes)})

    def run(self) -> None:
        self.install_signal_handlers()
        self.ensure_clean_queue()
        self.prepare_providers()
        self.start_providers()
        self.log_event(
            "soak_started",
            {
                "duration_hours": self.args.duration_hours,
                "provider_count": self.args.provider_count,
                "queue_depth": self.args.queue_depth,
                "task_budget": self.task_budget,
                "run_dir": str(self.run_dir),
            },
        )

        deadline = time.time() + (self.args.duration_hours * 3600)
        while time.time() < deadline and not self.stop_requested:
            now = time.time()
            if now - self.last_submit_at >= self.args.submit_interval:
                self.maybe_submit_work()
                self.last_submit_at = now
            if now - self.last_monitor_at >= self.args.monitor_interval:
                self.monitor()
                self.last_monitor_at = now
            self.sleep_interruptible(5)

        self.monitor()
        self.shutdown()


def main() -> int:
    args = parse_args()
    runner = SoakRunner(args)
    runner.run()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
