from __future__ import annotations

import os
import sys
import time
import unittest
from pathlib import Path
from uuid import uuid4

if os.environ.get("RC_RUN_INTEGRATION_TESTS") == "1":
    from supabase import create_client

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT / "src") not in sys.path:
    sys.path.insert(0, str(ROOT / "src"))

from rc.auth import generate_node_secret, hash_node_secret, node_auth_payload
from rc.config import NodeConfig, generate_node_id
from rc.credits import parse_credit_amount_to_subunits

if os.environ.get("RC_RUN_INTEGRATION_TESTS") == "1":
    if sys.version_info >= (3, 11):
        import tomllib
    else:
        import tomli as tomllib


class RuntimeIntegrationTests(unittest.TestCase):
    INTEGRATION_MODEL = "integration/test-model"
    KNOWN_MODEL = "mlx-community/DeepSeek-R1-Distill-Qwen-1.5B-4bit"
    TEST_TASK_BUDGET = parse_credit_amount_to_subunits("2")
    SMALL_RUNTIME_PARAMS = 1_000_000
    SMALL_RUNTIME_PROMPT_TOKENS = 1
    SMALL_RUNTIME_GENERATED_TOKENS = 0
    SMALL_RUNTIME_FLOPS = 2 * SMALL_RUNTIME_PARAMS * (
        SMALL_RUNTIME_PROMPT_TOKENS + SMALL_RUNTIME_GENERATED_TOKENS
    )
    SMALL_RUNTIME_MFLOPS = SMALL_RUNTIME_FLOPS // 1_000_000

    @classmethod
    def setUpClass(cls) -> None:
        if os.environ.get("RC_RUN_INTEGRATION_TESTS") != "1":
            raise unittest.SkipTest("Set RC_RUN_INTEGRATION_TESTS=1 to run live integration tests.")

        config_path = Path(
            os.environ.get(
                "RC_INTEGRATION_REQUESTER_CONFIG",
                str(Path.home() / ".rc" / "config.toml"),
            )
        ).expanduser()
        if not config_path.exists():
            raise unittest.SkipTest(f"Requester config not found at {config_path}")

        with config_path.open("rb") as f:
            cls.requester_cfg = NodeConfig.from_dict(tomllib.load(f))

        cls.client = create_client(cls.requester_cfg.supabase_url, cls.requester_cfg.supabase_key)
        cls._assert_clean_queue()
        cls._assert_requester_has_credits()

    @classmethod
    def _assert_clean_queue(cls) -> None:
        summary = cls.client.rpc("get_public_network_summary", {}).execute().data
        if summary["tasks"]["pending"] or summary["tasks"]["running"] or summary["tasks"]["assigned"]:
            raise unittest.SkipTest(
                "Integration tests require an empty queue. "
                f"Found pending={summary['tasks']['pending']} running={summary['tasks']['running']} "
                f"assigned={summary['tasks']['assigned']}."
            )

    @classmethod
    def _assert_requester_has_credits(cls) -> None:
        if os.environ.get("RC_INTEGRATION_ALLOW_ZERO_BALANCE") == "1":
            return
        balance = cls._get_balance(cls.requester_cfg)
        if balance["balance"] < cls.TEST_TASK_BUDGET * 2:
            raise unittest.SkipTest(
                f"Requester balance is too low for integration tests: {balance['balance']}"
            )

    @classmethod
    def _register_node(cls, role: str, name_prefix: str) -> NodeConfig:
        node_secret = generate_node_secret()
        cfg = NodeConfig(
            node_id=generate_node_id(),
            node_name=f"{name_prefix}-{uuid4().hex[:8]}",
            role=role,
            node_secret=node_secret,
            supabase_url=cls.requester_cfg.supabase_url,
            supabase_key=cls.requester_cfg.supabase_key,
            model=cls.INTEGRATION_MODEL,
        )
        cls.client.rpc(
            "register_node",
            {
                "p_id": cfg.node_id,
                "p_name": cfg.node_name,
                "p_role": cfg.role,
                "p_public_key": "",
                "p_auth_token_hash": hash_node_secret(node_secret),
                "p_gpu_info": {},
            },
        ).execute()
        return cfg

    @classmethod
    def _rpc(cls, name: str, params: dict) -> object:
        return cls.client.rpc(name, params).execute().data

    @classmethod
    def _submit_task(
        cls,
        prompt: str,
        budget: int,
        *,
        model: str | None = None,
    ) -> str:
        task_id = cls._rpc(
            "submit_task",
            {
                **node_auth_payload(cls.requester_cfg),
                "p_prompt": prompt,
                "p_model": model or cls.INTEGRATION_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": {},
            },
        )
        if not task_id:
            raise AssertionError("submit_task did not return a task id")
        return task_id

    @classmethod
    def _get_task(cls, cfg: NodeConfig, task_id: str) -> dict:
        task = cls._rpc(
            "get_task",
            {
                **node_auth_payload(cfg),
                "p_task_id": task_id,
            },
        )
        if task is None:
            raise AssertionError(f"Task {task_id} was not visible to node {cfg.node_id}")
        return task

    @classmethod
    def _get_balance(cls, cfg: NodeConfig) -> dict:
        rows = cls._rpc("get_balance", node_auth_payload(cfg)) or []
        if not rows:
            raise AssertionError(f"No balance row found for node {cfg.node_id}")
        return rows[0]

    def test_01_claim_renew_complete_flow(self) -> None:
        provider_cfg = self._register_node("provider", "itest-provider")
        requester_before = self._get_balance(self.requester_cfg)
        provider_before = self._get_balance(provider_cfg)

        task_id = self._submit_task(
            "Integration test: prove that True implies True.",
            budget=self.TEST_TASK_BUDGET,
        )

        claimed_task_id = self._rpc(
            "claim_task",
            {
                **node_auth_payload(provider_cfg),
                "p_lease_seconds": 10,
            },
        )
        self.assertEqual(claimed_task_id, task_id)

        marked_running = self._rpc(
            "mark_task_running",
            {
                **node_auth_payload(provider_cfg),
                "p_task_id": task_id,
                "p_lease_seconds": 10,
            },
        )
        self.assertTrue(marked_running)

        first_task = self._get_task(provider_cfg, task_id)
        first_lease = first_task["lease_expires_at"]
        self.assertEqual(first_task["status"], "running")

        time.sleep(2)
        renewed = self._rpc(
            "renew_task_lease",
            {
                **node_auth_payload(provider_cfg),
                "p_task_id": task_id,
                "p_lease_seconds": 10,
            },
        )
        self.assertTrue(renewed)

        second_task = self._get_task(provider_cfg, task_id)
        self.assertEqual(second_task["status"], "running")
        self.assertNotEqual(second_task["lease_expires_at"], first_lease)

        result_id = self._rpc(
            "complete_task",
            {
                **node_auth_payload(provider_cfg),
                "p_task_id": task_id,
                "p_lean_code": "theorem integration_true : True := by trivial",
                "p_raw_output": "theorem integration_true : True := by trivial",
                "p_mflops": 0,
                "p_prompt_tokens": self.SMALL_RUNTIME_PROMPT_TOKENS,
                "p_generated_tokens": self.SMALL_RUNTIME_GENERATED_TOKENS,
                "p_model_params": self.SMALL_RUNTIME_PARAMS,
                "p_prefill_eval_ms": 0.5,
                "p_decode_eval_ms": 0.5,
                "p_total_eval_ms": 1.0,
            },
        )
        self.assertTrue(result_id)

        completed_task = self._get_task(self.requester_cfg, task_id)
        self.assertEqual(completed_task["status"], "completed")
        self.assertEqual(completed_task["result"]["flops_used"], self.SMALL_RUNTIME_FLOPS)
        self.assertEqual(completed_task["result"]["mflops_used"], self.SMALL_RUNTIME_MFLOPS)
        self.assertEqual(completed_task["result"]["billed_mflops"], self.SMALL_RUNTIME_MFLOPS)

        requester_after = self._get_balance(self.requester_cfg)
        provider_after = self._get_balance(provider_cfg)
        self.assertEqual(requester_after["balance"], requester_before["balance"] - self.SMALL_RUNTIME_MFLOPS)
        self.assertEqual(provider_after["balance"], provider_before["balance"] + self.SMALL_RUNTIME_MFLOPS)

    def test_02_expiry_requeue_flow(self) -> None:
        provider_cfg = self._register_node("provider", "itest-expiry")
        task_id = self._submit_task(
            "Integration test: prove that P implies P.",
            budget=self.TEST_TASK_BUDGET,
        )

        claimed_task_id = self._rpc(
            "claim_task",
            {
                **node_auth_payload(provider_cfg),
                "p_lease_seconds": 5,
            },
        )
        self.assertEqual(claimed_task_id, task_id)

        marked_running = self._rpc(
            "mark_task_running",
            {
                **node_auth_payload(provider_cfg),
                "p_task_id": task_id,
                "p_lease_seconds": 5,
            },
        )
        self.assertTrue(marked_running)

        time.sleep(6)
        requeued = self._rpc("requeue_expired_tasks", {})
        self.assertGreaterEqual(requeued, 0)

        requeued_task = self._get_task(self.requester_cfg, task_id)
        self.assertEqual(requeued_task["status"], "pending")
        self.assertIn("expired", requeued_task["last_error"])

        reclaimed_task_id = self._rpc(
            "claim_task",
            {
                **node_auth_payload(provider_cfg),
                "p_lease_seconds": 10,
            },
        )
        self.assertEqual(reclaimed_task_id, task_id)

        remarked_running = self._rpc(
            "mark_task_running",
            {
                **node_auth_payload(provider_cfg),
                "p_task_id": task_id,
                "p_lease_seconds": 10,
            },
        )
        self.assertTrue(remarked_running)

        completed = self._rpc(
            "complete_task",
            {
                **node_auth_payload(provider_cfg),
                "p_task_id": task_id,
                "p_lean_code": "theorem integration_id (P : Prop) : P -> P := by intro h; exact h",
                "p_raw_output": "theorem integration_id (P : Prop) : P -> P := by intro h; exact h",
                "p_mflops": 0,
                "p_prompt_tokens": self.SMALL_RUNTIME_PROMPT_TOKENS,
                "p_generated_tokens": self.SMALL_RUNTIME_GENERATED_TOKENS,
                "p_model_params": self.SMALL_RUNTIME_PARAMS,
                "p_prefill_eval_ms": 0.5,
                "p_decode_eval_ms": 0.5,
                "p_total_eval_ms": 1.0,
            },
        )
        self.assertTrue(completed)

        final_task = self._get_task(self.requester_cfg, task_id)
        self.assertEqual(final_task["status"], "completed")
        self.assertEqual(final_task["attempt_count"], 2)

    def test_03_secret_rotation_flow(self) -> None:
        provider_cfg = self._register_node("provider", "itest-rotate")
        old_secret_hash = hash_node_secret(provider_cfg.node_secret)
        new_secret = generate_node_secret()
        new_secret_hash = hash_node_secret(new_secret)

        try:
            rotated = self._rpc(
                "rotate_node_secret",
                {
                    "p_node_id": provider_cfg.node_id,
                    "p_current_secret_hash": old_secret_hash,
                    "p_new_secret_hash": new_secret_hash,
                },
            )
        except Exception as exc:
            if "Could not find the function public.rotate_node_secret" in str(exc):
                self.skipTest(
                    "rotate_node_secret is not deployed in the live project yet. "
                    "Reapply supabase/hardening.sql before rerunning integration tests."
                )
            raise
        self.assertTrue(rotated)

        old_heartbeat = self._rpc(
            "update_node_heartbeat",
            {
                "p_node_id": provider_cfg.node_id,
                "p_node_secret_hash": old_secret_hash,
                "p_online": True,
                "p_gpu_info": {},
            },
        )
        self.assertFalse(old_heartbeat)

        provider_cfg.node_secret = new_secret
        new_heartbeat = self._rpc(
            "update_node_heartbeat",
            {
                **node_auth_payload(provider_cfg),
                "p_online": True,
                "p_gpu_info": {},
            },
        )
        self.assertTrue(new_heartbeat)

    def test_04_complete_task_prefers_runtime_params_over_known_model_catalog(self) -> None:
        provider_cfg = self._register_node("provider", "itest-runtime-params")
        provider_before = self._get_balance(provider_cfg)

        task_id = self._submit_task(
            "Integration test: verify billed compute uses provider runtime params.",
            budget=parse_credit_amount_to_subunits("0.01"),
            model=self.KNOWN_MODEL,
        )

        claimed_task_id = self._rpc(
            "claim_task",
            {
                **node_auth_payload(provider_cfg),
                "p_lease_seconds": 10,
            },
        )
        self.assertEqual(claimed_task_id, task_id)

        marked_running = self._rpc(
            "mark_task_running",
            {
                **node_auth_payload(provider_cfg),
                "p_task_id": task_id,
                "p_lease_seconds": 10,
            },
        )
        self.assertTrue(marked_running)

        result_id = self._rpc(
            "complete_task",
            {
                **node_auth_payload(provider_cfg),
                "p_task_id": task_id,
                "p_lean_code": "theorem integration_runtime_params : True := by trivial",
                "p_raw_output": "theorem integration_runtime_params : True := by trivial",
                "p_mflops": 0,
                "p_prompt_tokens": 2,
                "p_generated_tokens": 0,
                "p_model_params": self.SMALL_RUNTIME_PARAMS,
                "p_prefill_eval_ms": 0.5,
                "p_decode_eval_ms": 0.5,
                "p_total_eval_ms": 1.0,
            },
        )
        self.assertTrue(result_id)

        completed_task = self._get_task(self.requester_cfg, task_id)
        self.assertEqual(completed_task["status"], "completed")
        self.assertEqual(completed_task["result"]["flops_used"], 4_000_000)
        self.assertEqual(completed_task["result"]["mflops_used"], 4)
        self.assertEqual(completed_task["result"]["billed_mflops"], 4)

        provider_after = self._get_balance(provider_cfg)
        self.assertEqual(provider_after["balance"], provider_before["balance"] + 4)


if __name__ == "__main__":
    unittest.main()
