from __future__ import annotations

import io
import json
import re
import urllib.request
import zipfile
from typing import Any
from urllib.parse import urlparse

from failure_taxonomy import classify_failure_report
from llm_classifier import LLMAssistedClassifier


MAX_LOG_ARCHIVE_BYTES = 20_000_000
MAX_RUNTIME_LOG_BYTES = 20_000_000
MAX_ERROR_REPORT_BYTES = 1_000_000
ERROR_LINE_PATTERN = re.compile(
    r"(?:\b(?:[A-Za-z_]*(?:Error|Exception)|PREFLIGHT_[A-Z_]+|OOM)\b|"
    r"out of memory|not supported|unsupported|does not recognize|cannot|can.t|"
    r"not found|no such file|failed to|invalid|traceback|"
    r"找不到空闲卡|不支持|暂不支持|不兼容|请换用|请更换)",
    re.IGNORECASE,
)
MODEL_TYPE_NOT_RECOGNIZED_PATTERN = re.compile(
    r"model\s+type\s+[`'\"](?P<model_type>[A-Za-z0-9_.-]+)[`'\"]\s+but\s+"
    r"(?:Transformers\s+)?does\s+not\s+recognize\s+this\s+architecture",
    re.IGNORECASE,
)
MODEL_ARCHITECTURES_NOT_SUPPORTED_PATTERN = re.compile(
    r"Model\s+architectures?\s*(?P<architectures>\[[^\]\n]{1,500}\])\s+"
    r"(?:are|is)\s+not\s+supported\s+for\s+now",
    re.IGNORECASE,
)
SUBMITTED_FRAMEWORK_PATTERN = re.compile(
    r"\[submit\]\s*framework\s*:\s*(?P<framework>[A-Za-z0-9_.-]{1,80})",
    re.IGNORECASE,
)
TARGET_DOCKER_IMAGE_PATTERN = re.compile(
    r"\[submit\]\s*docker_image\s*:\s*(?P<image>\S{1,500})",
    re.IGNORECASE,
)
FRAMEWORK_IMAGE_MARKERS = (
    ("vllm-customized", ("vllm-customized", "vllm_customized")),
    ("vllm_fix_tokenizer", ("vllm-fix-tokenizer", "vllm_fix_tokenizer")),
    ("sentence-transformers", ("sentence-transformers", "sentence_transformers")),
    ("sherpa-onnx", ("sherpa-onnx", "sherpa_onnx")),
    ("llamacpp", ("llamacpp", "llama-cpp", "llama.cpp")),
    ("vllm-mlu", ("vllm-mlu", "vllm_mlu")),
    ("vllm-016", ("vllm-016", "vllm_016")),
    ("diffusers", ("diffusers",)),
    ("transformers", ("transformers",)),
    ("sglang", ("sglang",)),
    ("funasr", ("funasr",)),
    ("vllm", ("vllm",)),
)


def fetch_and_classify_failure_log(
    log_url: str,
    *,
    task_context: dict[str, Any] | None = None,
    llm_classifier: LLMAssistedClassifier | None = None,
    timeout_seconds: int = 20,
) -> dict[str, Any]:
    parsed_url = urlparse(log_url)
    if parsed_url.scheme not in {"http", "https"} or not parsed_url.netloc:
        raise ValueError("failure log URL must be HTTP(S)")
    request = urllib.request.Request(
        log_url,
        headers={"User-Agent": "modelhub-submmit-failure-inspector/1"},
    )
    with urllib.request.urlopen(request, timeout=max(1, int(timeout_seconds))) as response:
        archive = response.read(MAX_LOG_ARCHIVE_BYTES + 1)
    if len(archive) > MAX_LOG_ARCHIVE_BYTES:
        raise ValueError("failure log archive exceeds size limit")
    return classify_failure_archive(
        archive,
        task_context=task_context,
        llm_classifier=llm_classifier,
    )


def classify_failure_archive(
    archive: bytes,
    *,
    task_context: dict[str, Any] | None = None,
    llm_classifier: LLMAssistedClassifier | None = None,
) -> dict[str, Any]:
    report: dict[str, Any] = {}
    runtime_log = ""
    with zipfile.ZipFile(io.BytesIO(archive)) as bundle:
        names = set(bundle.namelist())
        if "error_report_raw.json" in names:
            info = bundle.getinfo("error_report_raw.json")
            if info.file_size <= MAX_ERROR_REPORT_BYTES:
                try:
                    parsed = json.loads(bundle.read("error_report_raw.json"))
                    if isinstance(parsed, dict):
                        report = parsed
                except (json.JSONDecodeError, UnicodeDecodeError):
                    report = {}
        if "pod_runtime_log.txt" in names:
            info = bundle.getinfo("pod_runtime_log.txt")
            if info.file_size <= MAX_RUNTIME_LOG_BYTES:
                runtime_log = bundle.read("pod_runtime_log.txt").decode("utf-8", "replace")

    error_lines = _extract_error_lines(runtime_log)
    report_code = str(report.get("code") or "").strip() or None
    suggestion = str(report.get("suggestion") or "")[:500] or None
    classification_inputs = [*error_lines]
    if suggestion:
        classification_inputs.append(suggestion)
    classification = classify_failure_report(report_code, classification_inputs)
    result: dict[str, Any] = {
        "failureCode": report_code,
        "failureSuggestion": suggestion,
        "failureCategory": classification.category,
        "failureScope": classification.scope,
        "failureAction": classification.action,
        "failureDeterministic": classification.deterministic,
        "failureNeedsLlm": classification.needs_llm,
        "failureClassificationReason": classification.reason,
        "failureEvidence": error_lines[-12:],
    }
    observed_memory_gib = _extract_observed_gpu_memory_gib(error_lines)
    if report_code == "PREFLIGHT_OOM" and observed_memory_gib is not None:
        result["failureObservedGpuMemoryGiB"] = observed_memory_gib
    unsupported_architectures, unsupported_model_types = _extract_unsupported_architectures(
        error_lines
    )
    if unsupported_architectures:
        result["failureUnsupportedArchitectures"] = unsupported_architectures
    if unsupported_model_types:
        result["failureUnsupportedModelTypes"] = unsupported_model_types
    detected_framework, framework_source = _extract_submitted_framework(runtime_log)
    if detected_framework:
        result["failureDetectedFramework"] = detected_framework
        result["failureDetectedFrameworkSource"] = framework_source
    if classification.needs_llm and llm_classifier is not None and llm_classifier.enabled:
        llm_decision = llm_classifier.classify_failure(
            task_context=dict(task_context or {}),
            report_code=report_code,
            suggestion=result["failureSuggestion"],
            error_lines=error_lines[-12:],
        )
        result["failureLlmDecision"] = llm_decision
        # Promote only a confident semantic decision. Deterministic rules and
        # platform signatures above never reach this branch, so an LLM cannot
        # override the hard classifiers.
        confidence = float(llm_decision.get("confidence") or 0.0)
        scope = str(llm_decision.get("scope") or "unknown")
        if confidence >= 0.80 and scope != "unknown":
            result.update(
                {
                    "failureCategory": llm_decision.get("category") or result["failureCategory"],
                    "failureScope": scope,
                    "failureAction": llm_decision.get("action") or result["failureAction"],
                    "failureClassificationReason": f"llm:{llm_decision.get('reason') or 'semantic_classification'}",
                    "failureNeedsLlm": False,
                }
            )
    return result


def _extract_submitted_framework(runtime_log: str) -> tuple[str | None, str | None]:
    """Recover only explicit or unambiguous target-framework evidence."""
    explicit = SUBMITTED_FRAMEWORK_PATTERN.search(runtime_log)
    if explicit:
        return explicit.group("framework").strip(), "submit_framework"

    image_match = TARGET_DOCKER_IMAGE_PATTERN.search(runtime_log)
    if image_match is None:
        return None, None
    target_image = image_match.group("image").strip().casefold()
    for framework, markers in FRAMEWORK_IMAGE_MARKERS:
        if any(marker in target_image for marker in markers):
            return framework, "target_docker_image"
    return None, None


def _extract_error_lines(runtime_log: str) -> list[str]:
    selected: list[str] = []
    for raw_line in runtime_log.splitlines():
        line = raw_line.strip()
        if not line or not ERROR_LINE_PATTERN.search(line):
            continue
        line = re.sub(r"\x1b\[[0-9;]*m", "", line)
        line = re.sub(r"\b[0-9a-f]{32,64}\b", "<id>", line, flags=re.IGNORECASE)
        line = line[:1200]
        if line not in selected:
            selected.append(line)
    return selected[-24:]


def _extract_observed_gpu_memory_gib(error_lines: list[str]) -> float | None:
    for line in error_lines:
        match = re.search(
            r"=\s*(?P<available>[0-9]+(?:\.[0-9]+)?)\s*GB\s+available\b",
            line,
            flags=re.IGNORECASE,
        )
        if match:
            value = float(match.group("available"))
            if 0 < value <= 1024:
                return value
    return None


def _extract_unsupported_architectures(
    error_lines: list[str],
) -> tuple[list[str], list[str]]:
    architectures: set[str] = set()
    model_types: set[str] = set()
    for line in error_lines:
        for match in MODEL_TYPE_NOT_RECOGNIZED_PATTERN.finditer(line):
            value = match.group("model_type").strip()
            if value:
                model_types.add(value)
        for match in MODEL_ARCHITECTURES_NOT_SUPPORTED_PATTERN.finditer(line):
            for value in re.findall(r"['\"]([^'\"]+)['\"]", match.group("architectures")):
                value = value.strip()
                if value:
                    architectures.add(value)
    return sorted(architectures, key=str.casefold), sorted(model_types, key=str.casefold)