from __future__ import annotations

import re
from collections import defaultdict
from pathlib import Path
from typing import Any

from common import read_jsonl
from gpu_strategy import wilson_lower_bound


class SafeConfigOptimizer:
    """Apply only repeatedly successful, sanitized parameter vectors."""

    def __init__(self, *, intents_path: Path | str, outcomes_path: Path | str) -> None:
        self.intents_path = Path(intents_path)
        self.outcomes_path = Path(outcomes_path)
        self._qualified = self._learn()

    def _learn(self) -> dict[tuple[str, str, str], list[dict[str, Any]]]:
        outcomes = read_jsonl(self.outcomes_path)
        by_task = {
            str(item.get("taskId")): item
            for item in outcomes
            if item.get("taskId") is not None and item.get("outcome") in {"success", "failed"}
        }
        groups: dict[tuple[str, str, str, str], dict[str, Any]] = defaultdict(
            lambda: {"success": 0, "failure": 0, "models": set(), "vector": {}}
        )
        for intent in read_jsonl(self.intents_path):
            task_id = str(intent.get("taskId") or "")
            outcome = by_task.get(task_id)
            vector = intent.get("safeConfigVector")
            if outcome is None or not isinstance(vector, dict) or not vector:
                continue
            route = (
                str(intent.get("taskType") or ""),
                str(intent.get("targetGpu") or ""),
                str(intent.get("framework") or ""),
                str(intent.get("configFingerprint") or ""),
            )
            group = groups[route]
            group["vector"] = dict(vector)
            group["models"].add(str(intent.get("repoId") or ""))
            if outcome.get("outcome") == "success":
                group["success"] += 1
            elif not outcome.get("platformFailure") and not outcome.get("policyCancelled"):
                group["failure"] += 1

        qualified: dict[tuple[str, str, str], list[dict[str, Any]]] = defaultdict(list)
        for (task_type, gpu, framework, fingerprint), group in groups.items():
            success = int(group["success"])
            failure = int(group["failure"])
            if success < 5 or failure > 0 or len(group["models"]) < 2:
                continue
            qualified[(task_type, gpu, framework)].append(
                {
                    "fingerprint": fingerprint,
                    "success": success,
                    "total": success + failure,
                    "lowerBound": wilson_lower_bound(success, success + failure),
                    "vector": group["vector"],
                }
            )
        return dict(qualified)

    def optimize(
        self,
        *,
        task_type: str,
        target_gpu: str,
        framework: str,
        official_config: str,
        official_lower_bound: float,
    ) -> tuple[str, dict[str, Any]]:
        choices = self._qualified.get((task_type, target_gpu, framework)) or []
        choices = [
            item
            for item in choices
            if float(item.get("lowerBound") or 0.0) >= float(official_lower_bound) + 0.05
        ]
        if not choices:
            return official_config, {"source": "official", "applied": False}
        choice = max(choices, key=lambda item: (float(item["lowerBound"]), int(item["success"])))
        patched = official_config
        vector = choice["vector"]
        substitutions = {
            "gpuNum": (r"(\bgpu_num\s*:\s*['\"]?)\d+", r"\g<1>{}"),
            "tensorParallel": (
                r"((?:--tensor-parallel-size|-tp)\s*[, ]?\s*['\"]?)\d+",
                r"\g<1>{}",
            ),
            "maxModelLen": (
                r"((?:--max-model-len|max_model_len|max_seq_len)\s*[: ,]+\s*['\"]?)\d+",
                r"\g<1>{}",
            ),
            "gpuMemoryUtilization": (
                r"((?:--gpu-memory-utilization|gpu_memory_utilization)\s*[: ,]+\s*['\"]?)[0-9.]+",
                r"\g<1>{}",
            ),
        }
        for key, (pattern, replacement) in substitutions.items():
            if key in vector:
                patched = re.sub(pattern, replacement.format(vector[key]), patched, flags=re.IGNORECASE)
        for key, option in (("dtype", "dtype"), ("quantization", "quantization"), ("loadFormat", "load[_-]format")):
            if key in vector:
                patched = re.sub(
                    rf"((?:--{option}|{option})\s*[: ,]+\s*['\"]?)[A-Za-z0-9_-]+",
                    rf"\g<1>{vector[key]}",
                    patched,
                    flags=re.IGNORECASE,
                )
        return patched, {
            "source": "official_plus_learned_patch",
            "applied": patched != official_config,
            "evidenceSuccess": choice["success"],
            "evidenceLowerBound": choice["lowerBound"],
            "fingerprint": choice["fingerprint"],
        }