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,
)


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
    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_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