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

import importlib.util
import time
import unittest
from pathlib import Path


MODULE_PATH = Path(__file__).with_name("clinic_api.py")
SPEC = importlib.util.spec_from_file_location("clinic_api", MODULE_PATH)
clinic_api = importlib.util.module_from_spec(SPEC)
assert SPEC.loader is not None
SPEC.loader.exec_module(clinic_api)


def response(value: float) -> dict:
    return {
        "status": "success",
        "data": {
            "resultType": "matrix",
            "result": [{
                "metric": {"instance": "tikv-0", "unsafe": "discarded"},
                "values": [[1, str(value)], [2, str(value + 1)]],
            }],
        },
    }


def plan(queries: list[dict]) -> dict:
    return {
        "schemaVersion": "clinic-metrics-batch-plan-v1",
        "scenarioId": "memory-oom",
        "timeWindow": "historical:absolute;start=1;end=2;timezone=UTC;source=test",
        "target": {"orgId": "org-1", "clusterId": "cluster-1"},
        "start": 1,
        "end": 2,
        "step": "60",
        "queries": queries,
    }


class FakeClinicAPI(clinic_api.ClinicAPI):
    def __init__(self, responses: dict[str, dict], delay: float = 0):
        self.responses = responses
        self.delay = delay
        self.calls = []

    def query_range_for_target(self, org_id, cluster_id, query, start, end, step="60"):
        self.calls.append((org_id, cluster_id, query, start, end, step))
        if self.delay:
            time.sleep(self.delay)
        return self.responses[query]

    def get_cluster_by_id(self, cluster_id):
        raise AssertionError("batch execution must not rediscover an exact target")


class ClinicMetricsBatchTest(unittest.TestCase):
    def test_runs_independent_queries_concurrently_and_summarizes(self):
        api = FakeClinicAPI({"metric_a": response(1), "metric_b": response(3)}, delay=0.1)
        started = time.monotonic()
        result = api.query_range_batch(plan([
            {"queryId": "alignment", "source": "data_plane", "query": "metric_a"},
            {"queryId": "mechanism", "source": "cluster_store", "query": "metric_b"},
        ]))
        elapsed = time.monotonic() - started

        self.assertLess(elapsed, 0.18)
        self.assertEqual(result["completion"], "complete")
        self.assertEqual([item["queryId"] for item in result["queries"]], [
            "alignment", "mechanism"
        ])
        self.assertEqual(result["queries"][0]["summary"]["sampleCount"], 2)
        self.assertEqual(result["queries"][0]["summary"]["series"][0]["labels"], {
            "instance": "tikv-0"
        })
        self.assertTrue(all(call[0:2] == ("org-1", "cluster-1") for call in api.calls))

    def test_recovers_cross_source_with_predeclared_parallel_split(self):
        api = FakeClinicAPI({
            "mixed": {"error": "HTTP 400: cross_source_metrics_rejected"},
            "data_metric": response(2),
            "store_metric": response(4),
        })
        result = api.query_range_batch(plan([{
            "queryId": "mechanism",
            "source": "mixed",
            "query": "mixed",
            "fallbackQueries": [
                {"source": "data_plane", "query": "data_metric"},
                {"source": "cluster_store", "query": "store_metric"},
            ],
        }]))

        query = result["queries"][0]
        self.assertEqual(result["completion"], "complete")
        self.assertEqual(query["status"], "complete")
        self.assertEqual(query["summary"]["seriesCount"], 2)
        self.assertEqual(query["recovery"]["code"], "cross_source_metrics_rejected")
        self.assertTrue(query["recovery"]["succeeded"])
        self.assertEqual({call[2] for call in api.calls}, {
            "mixed", "data_metric", "store_metric"
        })

    def test_recovers_top_level_or_without_agent_retry(self):
        api = FakeClinicAPI({
            "metric_a or metric_b": {"error": "cross_source_metrics_rejected"},
            "metric_a": response(1),
            "metric_b": response(5),
        })
        result = api.query_range_batch(plan([{
            "queryId": "strongest-alternative",
            "source": "mixed",
            "query": "metric_a or metric_b",
        }]))

        self.assertEqual(result["queries"][0]["status"], "complete")
        self.assertEqual(result["queries"][0]["recovery"]["method"],
                         "source_homogeneous_parallel_split")

    def test_rejects_a_plan_that_can_exceed_six_provider_calls(self):
        with self.assertRaisesRegex(ValueError, "provider_query_budget_exceeded"):
            clinic_api._validate_metrics_batch_plan(plan([{
                "queryId": "mechanism",
                "source": "mixed",
                "query": "mixed",
                "fallbackQueries": [
                    {"source": "data_plane", "query": f"metric_{index}"}
                    for index in range(6)
                ],
            }]))

    def test_trusted_context_overrides_agent_target_window_and_bounds(self):
        api = FakeClinicAPI({"metric_a": response(1)})
        untrusted_plan = plan([{
            "queryId": "alignment",
            "source": "data_plane",
            "query": "metric_a",
        }])
        untrusted_plan.update({
            "scenarioId": "wrong",
            "timeWindow": "shortened",
            "target": {"orgId": "wrong", "clusterId": "wrong"},
            "start": 999,
            "end": 1000,
        })
        context = {
            "schemaVersion": "auto-oncall-clinic-metrics-context-v1",
            "scenarioId": "memory-oom",
            "timeWindow": "historical:absolute;start=1;end=2;timezone=UTC;source=test",
            "target": {"orgId": "org-1", "clusterId": "cluster-1"},
        }

        result = api.query_range_batch(untrusted_plan, context)

        self.assertEqual(result["scenarioId"], "memory-oom")
        self.assertEqual(result["timeWindow"], context["timeWindow"])
        self.assertEqual(api.calls[0][0:5], ("org-1", "cluster-1", "metric_a", 1, 2))

    def test_memory_recipe_requires_no_agent_promql(self):
        recipe_queries = clinic_api._METRICS_BATCH_RECIPES["memory-oom"]
        api = FakeClinicAPI({item["query"]: response(index + 1)
                             for index, item in enumerate(recipe_queries)})
        context = {
            "schemaVersion": "auto-oncall-clinic-metrics-context-v1",
            "scenarioId": "memory-oom",
            "timeWindow": "historical:absolute;start=1;end=2;timezone=UTC;source=test",
            "target": {"orgId": "org-1", "clusterId": "cluster-1"},
        }

        result = api.query_range_batch({
            "schemaVersion": "clinic-metrics-batch-plan-v1",
            "recipe": "memory-oom",
        }, context)

        self.assertEqual(result["completion"], "complete")
        self.assertEqual(len(api.calls), 3)
        self.assertEqual([item["queryId"] for item in result["queries"]], [
            "alignment", "mechanism", "strongest-alternative"
        ])


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