from __future__ import annotations

import sys
import unittest
from pathlib import Path
from unittest.mock import patch


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

from rc.config import DEFAULT_BACKEND_URL, NodeConfig, resolve_backend_url
from rc.inference.catalog import SUPPORTED_MODELS, choose_best_model
from rc.inference.flops import compute_credit_subunits, compute_flops, compute_mflops, compute_tflops
from rc.inference.engine import suggest_thread_count


class ComputeFlopsTests(unittest.TestCase):
    def test_matches_expected_formula(self) -> None:
        self.assertEqual(
            compute_flops(
                model_params=13_000_000_000,
                prompt_tokens=100,
                generated_tokens=20,
            ),
            3_120_000_000_000,
        )

    def test_preserves_zero_without_flooring(self) -> None:
        self.assertEqual(
            compute_flops(
                model_params=0,
                prompt_tokens=0,
                generated_tokens=0,
            ),
            0,
        )


class ComputeMflopsTests(unittest.TestCase):
    def test_matches_expected_formula(self) -> None:
        self.assertEqual(
            compute_mflops(
                model_params=13_000_000_000,
                prompt_tokens=100,
                generated_tokens=20,
            ),
            3_120_000,
        )

    def test_uses_exact_integer_subunit_floor(self) -> None:
        model_params = 1_500_000_123
        prompt_tokens = 17
        generated_tokens = 19
        expected = (2 * model_params * (prompt_tokens + generated_tokens)) // 1_000_000
        self.assertEqual(
            compute_mflops(
                model_params=model_params,
                prompt_tokens=prompt_tokens,
                generated_tokens=generated_tokens,
            ),
            expected,
        )

    def test_small_runs_can_floor_to_zero(self) -> None:
        self.assertEqual(
            compute_mflops(
                model_params=0,
                prompt_tokens=0,
                generated_tokens=0,
            ),
            0,
        )


class ComputeTflopsTests(unittest.TestCase):
    def test_matches_expected_formula(self) -> None:
        self.assertAlmostEqual(
            compute_tflops(
                model_params=13_000_000_000,
                prompt_tokens=100,
                generated_tokens=20,
            ),
            3.12,
        )

    def test_zero_flops_stays_zero(self) -> None:
        self.assertAlmostEqual(
            compute_tflops(
                model_params=0,
                prompt_tokens=0,
                generated_tokens=0,
            ),
            0.0,
        )

    def test_matches_credit_subunit_scale_exactly(self) -> None:
        credit_subunits = compute_credit_subunits(
            model_params=1_500_000_000,
            prompt_tokens=128,
            generated_tokens=64,
        )
        self.assertAlmostEqual(
            compute_tflops(
                model_params=1_500_000_000,
                prompt_tokens=128,
                generated_tokens=64,
            ),
            credit_subunits / 1_000_000,
        )


class BackendConfigTests(unittest.TestCase):
    def test_node_config_defaults_to_deployed_backend(self) -> None:
        cfg = NodeConfig()
        self.assertEqual(cfg.backend_url, DEFAULT_BACKEND_URL)

    def test_from_dict_uses_deployed_backend_when_missing(self) -> None:
        cfg = NodeConfig.from_dict({})
        self.assertEqual(cfg.backend_url, DEFAULT_BACKEND_URL)

    def test_explicit_backend_url_overrides_config(self) -> None:
        cfg = NodeConfig(backend_url="https://configured.example.com")
        self.assertEqual(
            resolve_backend_url("https://flag.example.com/", cfg),
            "https://flag.example.com",
        )
        self.assertEqual(
            resolve_backend_url(None, cfg),
            "https://configured.example.com",
        )


class SqlHardeningContractTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.runtime_sql = (ROOT / "supabase" / "reservation_runtime.sql").read_text()
        cls.hardening_sql = (ROOT / "supabase" / "hardening.sql").read_text()
        cls.config_py = (ROOT / "src" / "rc" / "config.py").read_text()

    def test_complete_task_requires_running_state(self) -> None:
        for sql in (self.runtime_sql, self.hardening_sql):
            self.assertIn(
                "IF v_task_status <> 'running' THEN",
                sql,
            )
            self.assertNotIn(
                "IF v_task_status NOT IN ('assigned', 'running') THEN",
                sql,
            )

    def test_complete_task_recomputes_precise_credit_subunits_server_side(self) -> None:
        for sql in (self.runtime_sql, self.hardening_sql):
            self.assertIn(
                "v_runtime_flops := compute_runtime_flops(",
                sql,
            )
            self.assertIn(
                "v_runtime_mflops := compute_runtime_mflops(",
                sql,
            )
            self.assertIn(
                "v_runtime_credit_subunits := compute_runtime_credit_subunits(",
                sql,
            )
            self.assertIn(
                "v_billed_credits := LEAST(v_runtime_credit_subunits, v_credit_budget);",
                sql,
            )
            self.assertNotIn(
                "v_billed_credits := LEAST(v_runtime_gflops, v_credit_budget);",
                sql,
            )
            self.assertNotIn(
                "v_billed_credits := LEAST(v_runtime_tflops, v_credit_budget);",
                sql,
            )

    def test_runtime_flops_are_stored_without_information_loss(self) -> None:
        for sql in (self.runtime_sql, self.hardening_sql):
            self.assertIn("ADD COLUMN IF NOT EXISTS flops_used bigint NOT NULL DEFAULT 0", sql)
            self.assertIn("INSERT INTO task_results (", sql)
            self.assertIn("flops_used,", sql)
            self.assertIn("v_runtime_flops,", sql)
            self.assertIn("CREATE OR REPLACE FUNCTION compute_runtime_flops(", sql)

    def test_runtime_accounting_prefers_provider_reported_model_params(self) -> None:
        expected_formula = (
            "COALESCE(NULLIF(GREATEST(COALESCE(p_model_params, 0), 0), 0), "
            "get_known_model_param_count(p_model), 0)"
        )
        expected_completion = (
            "v_model_params := COALESCE(\n"
            "        NULLIF(GREATEST(COALESCE(p_model_params, 0), 0), 0),\n"
            "        get_known_model_param_count(v_task_model)\n"
            "    );"
        )
        old_formula = (
            "COALESCE(get_known_model_param_count(p_model), "
            "NULLIF(GREATEST(COALESCE(p_model_params, 0), 0), 0), 0)"
        )
        old_completion = (
            "v_model_params := COALESCE(\n"
            "        get_known_model_param_count(v_task_model),\n"
            "        NULLIF(GREATEST(COALESCE(p_model_params, 0), 0), 0)\n"
            "    );"
        )
        for sql in (self.runtime_sql, self.hardening_sql):
            self.assertIn(expected_formula, sql)
            self.assertIn(expected_completion, sql)
            self.assertNotIn(old_formula, sql)
            self.assertNotIn(old_completion, sql)

    def test_runtime_mflops_no_longer_clamps_to_one(self) -> None:
        for sql in (self.runtime_sql, self.hardening_sql):
            self.assertNotIn(") / 1000000,\n        1\n    )::bigint;", sql)

    def test_runtime_tflops_no_longer_rounds_away_precision(self) -> None:
        for sql in (self.runtime_sql, self.hardening_sql):
            self.assertIn("SELECT compute_runtime_flops(", sql)
            self.assertNotIn("SELECT ROUND(", sql)

    def test_complete_task_updates_research_attempts_for_verification(self) -> None:
        expected_fragments = (
            "UPDATE proof_attempts",
            "SET provider_id = p_node_id,",
            "task_result_id = v_result_id,",
            "generated_lean_code = COALESCE(p_lean_code, '')",
            "raw_output = COALESCE(p_raw_output, '')",
            "ELSE 'generated'",
            "ELSE 'Provider generated candidate proof; verification pending'",
            "WHERE task_id = p_task_id;",
        )
        for sql in (self.runtime_sql, self.hardening_sql):
            for fragment in expected_fragments:
                self.assertIn(fragment, sql)

    def test_known_model_catalog_is_mirrored_in_sql(self) -> None:
        for model in SUPPORTED_MODELS:
            expected = f"WHEN '{model.repo_id}' THEN {model.param_count}::bigint"
            self.assertIn(expected, self.runtime_sql)
            self.assertIn(expected, self.hardening_sql)

    def test_running_tasks_can_renew_their_lease(self) -> None:
        for sql in (self.runtime_sql, self.hardening_sql):
            self.assertIn("CREATE OR REPLACE FUNCTION renew_task_lease(", sql)
            self.assertIn("SET lease_expires_at = now() + make_interval(secs => v_lease_seconds)", sql)
            self.assertIn("AND status = 'running';", sql)

    def test_hardening_defines_runtime_mode_for_fresh_installs(self) -> None:
        self.assertIn("CREATE OR REPLACE FUNCTION get_runtime_mode()", self.hardening_sql)
        self.assertIn("'runtime_mode', get_runtime_mode(),", self.hardening_sql)
        self.assertIn("GRANT EXECUTE ON FUNCTION get_runtime_mode() TO anon, authenticated;", self.hardening_sql)

    def test_hardening_supports_secret_rotation(self) -> None:
        self.assertIn("CREATE OR REPLACE FUNCTION rotate_node_secret(", self.hardening_sql)
        self.assertIn("SET auth_token_hash = p_new_secret_hash", self.hardening_sql)
        self.assertIn(
            "GRANT EXECUTE ON FUNCTION rotate_node_secret(uuid, text, text) TO anon, authenticated;",
            self.hardening_sql,
        )

    def test_hardening_heartbeat_refreshes_gpu_info(self) -> None:
        self.assertIn(
            "CREATE OR REPLACE FUNCTION update_node_heartbeat(",
            self.hardening_sql,
        )
        self.assertIn("p_gpu_info jsonb DEFAULT '{}'::jsonb", self.hardening_sql)
        self.assertIn("gpu_info = CASE", self.hardening_sql)
        self.assertIn("COALESCE(nodes.gpu_info, '{}'::jsonb) || p_gpu_info", self.hardening_sql)
        self.assertIn(
            "GRANT EXECUTE ON FUNCTION update_node_heartbeat(uuid, text, boolean, jsonb) TO anon, authenticated;",
            self.hardening_sql,
        )

    def test_hardening_network_summary_only_counts_online_nodes_as_healthy(self) -> None:
        self.assertIn("COUNT(*) FILTER (WHERE status = 'online') AS online_nodes,", self.hardening_sql)
        self.assertIn("COUNT(*) FILTER (WHERE status = 'offline') AS offline_nodes,", self.hardening_sql)
        self.assertIn(
            "WHERE status = 'online'\n                  AND last_heartbeat IS NOT NULL\n                  AND last_heartbeat >= now() - (SELECT healthy_window FROM freshness)\n            ) AS healthy_nodes,",
            self.hardening_sql,
        )
        self.assertIn(
            "WHERE role IN ('provider', 'both')\n                  AND status = 'online'\n                  AND last_heartbeat IS NOT NULL",
            self.hardening_sql,
        )
        self.assertIn(
            "WHERE role IN ('requester', 'both')\n                  AND status = 'online'\n                  AND last_heartbeat IS NOT NULL",
            self.hardening_sql,
        )
        self.assertIn("'stale', GREATEST(ns.online_nodes - ns.healthy_nodes, 0),", self.hardening_sql)

    def test_verified_lemmas_table_and_rpc_exist(self) -> None:
        for sql in (self.runtime_sql, self.hardening_sql):
            self.assertIn("CREATE TABLE IF NOT EXISTS verified_lemmas (", sql)
            self.assertIn("CREATE OR REPLACE FUNCTION store_verified_lemma(", sql)
            self.assertIn(
                "GRANT EXECUTE ON FUNCTION store_verified_lemma(text, text, jsonb, text, text, text, text, text, text, text, boolean) TO service_role;",
                sql,
            )
        self.assertIn("CREATE OR REPLACE FUNCTION get_public_verified_lemmas(", self.hardening_sql)
        self.assertIn(
            "GRANT EXECUTE ON FUNCTION get_public_verified_lemmas(integer) TO anon, authenticated;",
            self.hardening_sql,
        )

    def test_backend_owned_runtime_rpcs_are_service_role_only(self) -> None:
        expected_grants = (
            "GRANT EXECUTE ON FUNCTION backend_list_active_research_jobs() TO service_role;",
            "GRANT EXECUTE ON FUNCTION backend_get_research_job_state(uuid) TO service_role;",
            "GRANT EXECUTE ON FUNCTION backend_create_research_episode(uuid, text, text, text, text, jsonb, double precision, double precision) TO service_role;",
            "GRANT EXECUTE ON FUNCTION backend_enqueue_research_attempt(uuid, uuid, text, text, text, text, uuid, text, text, jsonb) TO service_role;",
            "GRANT EXECUTE ON FUNCTION backend_verify_research_attempt(uuid, text, text, text, text, text, jsonb) TO service_role;",
            "GRANT EXECUTE ON FUNCTION backend_update_research_job_status(uuid, text, text) TO service_role;",
        )
        unexpected_grants = (
            "GRANT EXECUTE ON FUNCTION backend_list_active_research_jobs() TO anon, authenticated;",
            "GRANT EXECUTE ON FUNCTION backend_get_research_job_state(uuid) TO anon, authenticated;",
            "GRANT EXECUTE ON FUNCTION backend_create_research_episode(uuid, text, text, text, text, jsonb, double precision, double precision) TO anon, authenticated;",
            "GRANT EXECUTE ON FUNCTION backend_enqueue_research_attempt(uuid, uuid, text, text, text, text, uuid, text, text, jsonb) TO anon, authenticated;",
            "GRANT EXECUTE ON FUNCTION backend_verify_research_attempt(uuid, text, text, text, text, text, jsonb) TO anon, authenticated;",
            "GRANT EXECUTE ON FUNCTION backend_update_research_job_status(uuid, text, text) TO anon, authenticated;",
        )
        for fragment in expected_grants:
            self.assertIn(fragment, self.runtime_sql)
        for fragment in unexpected_grants:
            self.assertNotIn(fragment, self.runtime_sql)

    def test_config_gpu_payload_tracks_total_memory(self) -> None:
        self.assertIn("total_memory_gb: Optional[float] = None", self.config_py)
        self.assertIn('"total_memory_gb": self.total_memory_gb', self.config_py)
        self.assertIn("def gpu_info_payload(self) -> dict[str, object]:", self.config_py)


class CatalogSelectionTests(unittest.TestCase):
    def test_choose_best_model_prefers_largest_tier_that_fits(self) -> None:
        with patch("rc.inference.catalog.platform.system", return_value="Linux"):
            with patch("rc.inference.catalog.platform.machine", return_value="x86_64"):
                self.assertEqual(
                    choose_best_model(available_memory_gb=18).repo_id,
                    "bartowski/DeepSeek-R1-Distill-Qwen-14B-GGUF",
                )

    def test_choose_best_model_prefers_mlx_on_apple_for_same_size(self) -> None:
        with patch("rc.inference.catalog.platform.system", return_value="Darwin"):
            with patch("rc.inference.catalog.platform.machine", return_value="arm64"):
                self.assertEqual(
                    choose_best_model(available_memory_gb=18).repo_id,
                    "mlx-community/DeepSeek-R1-Distill-Qwen-14B-4bit",
                )


class ThreadSizingTests(unittest.TestCase):
    def test_splits_threads_across_local_provider_processes(self) -> None:
        self.assertEqual(
            suggest_thread_count(active_provider_count=10, cpu_count=12),
            1,
        )

    def test_preserves_parallelism_for_single_provider(self) -> None:
        self.assertEqual(
            suggest_thread_count(active_provider_count=1, cpu_count=12),
            11,
        )


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