from __future__ import annotations

import json
import os
import re
import threading
from dataclasses import dataclass
from datetime import timedelta
from typing import Any

from llm_classifier import LLMAssistedClassifier
from models import ModelInspection
from common import parse_datetime, utc_now


# ModelHub allocations are preferred over product-card capacities. Nine values
# were repeated consistently across 299 structured PREFLIGHT_OOM reports pulled
# on 2026-08-10. The remaining currently verifiable devices use published card
# specifications and are replaced automatically when ModelHub returns its own
# observed allocation in a future failure log.
GPU_MEMORY_EVIDENCE: dict[str, dict[str, Any]] = {
    "Ascend_910-b4": {"memoryGiB": 32.0, "source": "modelhub_preflight_oom:8"},
    "Cambricon_mlu-370-x4": {"memoryGiB": 24.0, "source": "modelhub_preflight_oom:49"},
    "Iluvatar_bi-100": {"memoryGiB": 32.0, "source": "modelhub_preflight_oom:48"},
    "Iluvatar_bi-150": {"memoryGiB": 32.0, "source": "modelhub_preflight_oom:24"},
    "Iluvatar_mrv-100": {"memoryGiB": 32.0, "source": "modelhub_preflight_oom:50"},
    "MetaX_c-500": {"memoryGiB": 64.0, "source": "modelhub_preflight_oom:29"},
    "Sunrise_pt-200-x1": {"memoryGiB": 64.0, "source": "modelhub_preflight_oom:36"},
    "Vastai_va16": {"memoryGiB": 32.0, "source": "modelhub_preflight_oom:40"},
    "hygon_k100-ai": {"memoryGiB": 64.0, "source": "modelhub_preflight_oom:15"},
    "Ascend_910-b3": {
        "memoryGiB": 64.0,
        "source": "published_card_spec",
        "sourceUrl": "https://aclanthology.org/2025.emnlp-main.1630.pdf",
    },
    "Biren_166m": {
        "memoryGiB": 64.0,
        "source": "manufacturer_spec",
        "sourceUrl": "https://www.birentech.com/news/id6rz98v3obczy77cmxzfgk3/",
    },
    "Cambricon_mlu-370-x8": {
        "memoryGiB": 48.0,
        "source": "manufacturer_spec",
        "sourceUrl": "https://cambricon.com/index.php?a=lists&c=index&catid=406&m=content",
    },
    "Kunlunxin_p-800": {
        "memoryGiB": 96.0,
        "source": "published_procurement_spec",
        "sourceUrl": "https://pms2g.shrcb.com/cms/cmscaigougg/1d93b4b8166041e096c65d739073ade1.html",
    },
    "Mthreads_s4000": {
        "memoryGiB": 48.0,
        "source": "manufacturer_spec",
        "sourceUrl": "https://docs.mthreads.com/s4000/s4000-doc-online/product_specifications/",
    },
}
OBSERVED_GPU_MEMORY_GIB = {
    gpu: float(evidence["memoryGiB"])
    for gpu, evidence in GPU_MEMORY_EVIDENCE.items()
}
MODEL_LOAD_OVERHEAD = 1.20
TEXT_TOKENIZER_TASKS = {
    "text-generation",
    "visual-multi-modal",
    "reinforcement_learning",
    "question_answering",
    "feature_emb",
    "text_classification",
}
# A small, deliberately conservative baseline used only to decide whether an
# architecture needs semantic review. It is not an allowlist: unknown values
# remain eligible when the optional LLM is disabled or abstains.
MATURE_MODEL_TYPES = {
    "baichuan",
    "bert",
    "bloom",
    "chatglm",
    "cohere",
    "deepseek_v2",
    "deepseek_v3",
    "falcon",
    "gemma",
    "gemma2",
    "gemma3",
    "glm",
    "glm4",
    "gpt2",
    "gpt_bigcode",
    "internlm",
    "internlm2",
    "llama",
    "mistral",
    "mixtral",
    "mpt",
    "opt",
    "phi",
    "phi3",
    "qwen2",
    "qwen2_moe",
    "qwen3",
    "qwen3_moe",
    "roberta",
    "t5",
    "whisper",
    "xlm-roberta",
}


@dataclass(frozen=True)
class PreflightAssessment:
    allowed: bool
    config_params: str
    reason: str | None
    warnings: tuple[str, ...]
    ambiguous_reasons: tuple[str, ...]
    metadata: dict[str, Any]


class CandidatePreflightAdvisor:
    def __init__(
        self,
        *,
        llm_classifier: LLMAssistedClassifier | None = None,
        gpu_memory_gib: dict[str, float] | None = None,
    ) -> None:
        self.llm_classifier = llm_classifier
        environment_memory = _load_gpu_memory_overrides()
        self.gpu_memory_gib = {
            **OBSERVED_GPU_MEMORY_GIB,
            **environment_memory,
            **(gpu_memory_gib or {}),
        }
        self.gpu_memory_evidence = {
            gpu: dict(GPU_MEMORY_EVIDENCE.get(gpu) or {"memoryGiB": memory, "source": "constructor_override"})
            for gpu, memory in self.gpu_memory_gib.items()
        }
        for gpu in environment_memory:
            self.gpu_memory_evidence[gpu] = {
                "memoryGiB": self.gpu_memory_gib[gpu],
                "source": "environment_override",
            }
        for gpu in (gpu_memory_gib or {}):
            self.gpu_memory_evidence[gpu] = {
                "memoryGiB": self.gpu_memory_gib[gpu],
                "source": "constructor_override",
            }
        self._lock = threading.Lock()
        self._assessed = 0
        self._hard_blocks = 0
        self._llm_blocks = 0
        self._context_clamps = 0
        self._ambiguous = 0
        self._feedback_stats: dict[str, Any] = {}

    def set_feedback_stats(self, report: dict[str, Any] | None) -> None:
        self._feedback_stats = report if isinstance(report, dict) else {}
        for gpu, value in (self._feedback_stats.get("observedGpuMemoryGiB") or {}).items():
            try:
                memory_gib = float(value)
            except (TypeError, ValueError):
                continue
            if not 0 < memory_gib <= 1024:
                continue
            self.gpu_memory_gib[str(gpu)] = memory_gib
            self.gpu_memory_evidence[str(gpu)] = {
                "memoryGiB": memory_gib,
                "source": "local_modelhub_preflight_oom",
            }

    def assess(
        self,
        *,
        inspection: ModelInspection,
        task_type: str,
        target_gpu: str,
        framework: str,
        config_params: str,
    ) -> PreflightAssessment:
        warnings: list[str] = []
        ambiguous: list[str] = []
        metadata: dict[str, Any] = {
            "modelType": inspection.model_type,
            "architectures": inspection.architectures,
            "quantizationMethod": inspection.quantization_method,
            "estimatedLoadBytes": inspection.estimated_load_bytes(framework),
            "repositoryOnDiskBytes": inspection.repository_size_bytes,
            "gpuMemoryGiB": self.gpu_memory_gib.get(target_gpu),
            "gpuMemoryEvidence": self.gpu_memory_evidence.get(target_gpu),
        }
        with self._lock:
            self._assessed += 1

        # Empty file_paths means an injected/test inspection lacks structural
        # metadata. Real discoveries with an empty tree already fail the weight
        # compatibility gate, so do not make this test/fallback state a blocker.
        has_structure_metadata = bool(inspection.file_paths)
        applies_text_structure_rules = (
            task_type in TEXT_TOKENIZER_TASKS and framework != "llamacpp"
        )
        if has_structure_metadata and applies_text_structure_rules:
            if not inspection.has_root_config:
                return self._hard_block(
                    config_params,
                    "preflight_missing_root_config",
                    warnings,
                    ambiguous,
                    metadata,
                )
            if not inspection.has_root_standard_weights:
                return self._hard_block(
                    config_params,
                    "preflight_missing_root_weights",
                    warnings,
                    ambiguous,
                    metadata,
                )
            if not inspection.has_root_tokenizer:
                return self._hard_block(
                    config_params,
                    "preflight_missing_root_tokenizer",
                    warnings,
                    ambiguous,
                    metadata,
                )

        estimated_bytes = inspection.estimated_load_bytes(framework)
        repository_bytes = inspection.repository_size_bytes
        memory_gib = self.gpu_memory_gib.get(target_gpu)
        if has_structure_metadata and memory_gib is None:
            return self._hard_block(
                config_params,
                "preflight_gpu_memory_unknown",
                warnings,
                ambiguous,
                metadata,
            )
        if memory_gib and has_structure_metadata and repository_bytes is None:
            return self._hard_block(
                config_params,
                "preflight_model_size_unknown",
                warnings,
                ambiguous,
                metadata,
            )
        memory_sizing_bytes = repository_bytes or estimated_bytes
        if memory_sizing_bytes and memory_gib:
            required_gib = memory_sizing_bytes / (1024**3) * MODEL_LOAD_OVERHEAD
            metadata["estimatedRequiredGiB"] = round(required_gib, 3)
            metadata["memorySizingBasis"] = (
                "recursive_repository_on_disk" if repository_bytes else "selected_weights_fallback"
            )
            metadata["maximumRepositorySizeGiB"] = round(memory_gib / MODEL_LOAD_OVERHEAD, 3)
            if required_gib > memory_gib:
                return self._hard_block(
                    config_params,
                    "preflight_predicted_oom",
                    warnings,
                    ambiguous,
                    metadata,
                )

        profile_key = "|".join(
            (
                target_gpu,
                framework,
                task_type,
                inspection.model_type or "unknown",
                inspection.quantization_method or "none",
            )
        )
        profile_feedback = (
            (self._feedback_stats.get("recentProfileCombinationStats") or {}).get(profile_key)
            or {}
        )
        if profile_feedback:
            metadata["recentProfileFeedback"] = profile_feedback
            last_terminal_at = parse_datetime(profile_feedback.get("lastTerminalAt"))
            consecutive_failures = int(profile_feedback.get("consecutiveFailures") or 0)
            circuit_open = bool(
                last_terminal_at is not None
                and consecutive_failures >= 5
                and last_terminal_at + timedelta(hours=12) > utc_now()
            )
            if circuit_open:
                return self._hard_block(
                    config_params,
                    "preflight_recent_profile_failure_circuit",
                    warnings,
                    ambiguous,
                    metadata,
                )

        rendered_config, context_clamped = clamp_context_length(
            config_params,
            inspection.max_context_length,
        )
        if context_clamped:
            warnings.append("preflight_context_length_clamped")
            with self._lock:
                self._context_clamps += 1

        if inspection.config_fetch_error:
            warnings.append("model_config_metadata_unavailable")
        elif applies_text_structure_rules:
            if not inspection.model_type:
                ambiguous.append("model_type_missing_or_unknown")
            elif inspection.model_type.lower() not in MATURE_MODEL_TYPES:
                ambiguous.append(f"architecture_not_in_mature_baseline:{inspection.model_type.lower()}")
            auto_map = inspection.model_config.get("auto_map")
            if isinstance(auto_map, dict) and auto_map:
                ambiguous.append("custom_remote_code_architecture")
            if inspection.quantization_method:
                ambiguous.append(f"quantization_compatibility:{inspection.quantization_method}")
            if inspection.architectures and any(
                not architecture.endswith(
                    (
                        "ForCausalLM",
                        "ForConditionalGeneration",
                        "ForSequenceClassification",
                        "Model",
                    )
                )
                for architecture in inspection.architectures
            ):
                ambiguous.append("nonstandard_architecture_name")

        if ambiguous:
            with self._lock:
                self._ambiguous += 1
            llm = self.llm_classifier
            review_gate = getattr(llm, "should_review_candidate", None) if llm is not None else None
            should_review = bool(
                llm is not None
                and llm.enabled
                and (review_gate(ambiguous) if callable(review_gate) else True)
            )
            if should_review:
                decision = llm.classify(
                    inspection=inspection,
                    task_type=task_type,
                    target_gpu=target_gpu,
                    framework=framework,
                    ambiguous_reasons=ambiguous,
                )
                metadata["llmDecision"] = decision
                if llm.blocks(decision):
                    with self._lock:
                        self._llm_blocks += 1
                    reason = str(decision.get("reason") or "incompatible")
                    return PreflightAssessment(
                        allowed=False,
                        config_params=rendered_config,
                        reason=f"llm_high_confidence_incompatible:{reason}",
                        warnings=tuple(warnings),
                        ambiguous_reasons=tuple(ambiguous),
                        metadata=metadata,
                    )
                warnings.append(f"llm_review_{decision.get('decision') or 'abstain'}")
            elif llm is not None and llm.enabled:
                warnings.append("llm_review_not_needed")
            else:
                warnings.append("llm_review_unavailable")

        return PreflightAssessment(
            allowed=True,
            config_params=rendered_config,
            reason=None,
            warnings=tuple(warnings),
            ambiguous_reasons=tuple(ambiguous),
            metadata=metadata,
        )

    def _hard_block(
        self,
        config_params: str,
        reason: str,
        warnings: list[str],
        ambiguous: list[str],
        metadata: dict[str, Any],
    ) -> PreflightAssessment:
        with self._lock:
            self._hard_blocks += 1
        return PreflightAssessment(
            allowed=False,
            config_params=config_params,
            reason=reason,
            warnings=tuple(warnings),
            ambiguous_reasons=tuple(ambiguous),
            metadata=metadata,
        )

    def summary(self) -> dict[str, Any]:
        with self._lock:
            summary = {
                "enabled": True,
                "assessedCandidates": self._assessed,
                "hardBlocks": self._hard_blocks,
                "llmBlocks": self._llm_blocks,
                "ambiguousCandidates": self._ambiguous,
                "contextLengthClamps": self._context_clamps,
                "knownGpuMemoryGiB": dict(self.gpu_memory_gib),
                "gpuMemoryEvidence": dict(self.gpu_memory_evidence),
            }
        summary["llm"] = self.llm_classifier.summary() if self.llm_classifier else {"enabled": False}
        return summary


def clamp_context_length(config: str, maximum: int | None) -> tuple[str, bool]:
    if not maximum or maximum <= 0:
        return config, False
    changed = False

    def replace_value(match: re.Match[str]) -> str:
        nonlocal changed
        current = int(match.group("value"))
        if current <= maximum:
            return match.group(0)
        changed = True
        return f"{match.group('prefix')}{maximum}{match.group('suffix')}"

    patterns = (
        r"(?P<prefix>(?:max_model_len|max_seq_len)\s*[:=]\s*['\"]?)(?P<value>\d+)(?P<suffix>['\"]?)",
        r"(?P<prefix>['\"]max_seq_len['\"]\s*:\s*['\"]?)(?P<value>\d+)(?P<suffix>['\"]?)",
        r"(?P<prefix>MAX_MODEL_LEN\s*,?\s*value\s*:\s*['\"]?)(?P<value>\d+)(?P<suffix>['\"]?)",
        r"(?P<prefix>--max-model-len(?:\s+|\s*,\s*(?:\n\s*)?|\s*\n\s*-\s*)['\"]?)(?P<value>\d+)(?P<suffix>['\"]?)",
    )
    rendered = config
    for pattern in patterns:
        rendered = re.sub(pattern, replace_value, rendered, flags=re.IGNORECASE)
    return rendered, changed


def _load_gpu_memory_overrides() -> dict[str, float]:
    raw = os.getenv("MODELHUB_GPU_MEMORY_GIB_JSON")
    if not raw:
        return {}
    try:
        payload = json.loads(raw)
    except json.JSONDecodeError:
        return {}
    if not isinstance(payload, dict):
        return {}
    result: dict[str, float] = {}
    for gpu, value in payload.items():
        try:
            parsed = float(value)
        except (TypeError, ValueError):
            continue
        if parsed > 0:
            result[str(gpu)] = parsed
    return result