from __future__ import annotations

import random
import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Tuple

import numpy as np
import torch


def cuda_is_available_quiet() -> bool:
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            message="CUDA initialization: Unexpected error from cudaGetDeviceCount\\(\\).*",
            category=UserWarning,
        )
        return torch.cuda.is_available()


@dataclass(frozen=True)
class ArtifactPaths:
    root: Path
    models: Path
    videos: Path
    logs: Path


def project_root() -> Path:
    return Path(__file__).resolve().parents[1]


def ensure_artifact_paths(root: Optional[Path] = None) -> ArtifactPaths:
    root = root or project_root()
    artifacts = root / "artifacts"
    models = artifacts / "models"
    videos = artifacts / "videos"
    logs = artifacts / "logs"
    for p in (artifacts, models, videos, logs):
        p.mkdir(parents=True, exist_ok=True)
    return ArtifactPaths(root=artifacts, models=models, videos=videos, logs=logs)


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if cuda_is_available_quiet():
        torch.cuda.manual_seed_all(seed)


def resolve_torch_device(requested: str = "auto") -> Tuple[str, str]:
    requested = requested.lower().strip()
    if requested == "cpu":
        return "cpu", "User requested CPU."

    def _probe_cuda(device_name: str = "cuda") -> Tuple[str, str]:
        if not cuda_is_available_quiet():
            return "cpu", "CUDA unavailable, using CPU."
        try:
            x = torch.ones(8, device=device_name)
            _ = (x * 2).cpu().numpy()
            gpu_name = torch.cuda.get_device_name(torch.device(device_name))
            return device_name, f"CUDA is available ({gpu_name})."
        except Exception as exc:  # pragma: no cover - hardware dependent
            return "cpu", f"CUDA check failed ({exc}), fallback to CPU."

    def _probe_mps() -> Tuple[str, str]:
        mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
        if not mps_available:
            return "cpu", "MPS unavailable, using CPU."
        try:
            x = torch.ones(8, device="mps")
            _ = (x * 2).cpu().numpy()
            return "mps", "MPS is available and passed a quick tensor sanity check."
        except Exception as exc:  # pragma: no cover - hardware dependent
            return "cpu", f"MPS check failed ({exc}), fallback to CPU."

    if requested.startswith("cuda"):
        return _probe_cuda(requested)

    if requested == "mps":
        device, msg = _probe_mps()
        if device == "cpu":
            return device, "MPS requested but unavailable, fallback to CPU."
        return device, msg

    if requested == "auto":
        device, msg = _probe_cuda("cuda")
        if device != "cpu":
            return device, msg

        device, msg = _probe_mps()
        if device != "cpu":
            return device, msg

        return "cpu", "Neither CUDA nor MPS is available, using CPU."

    return "cpu", f"Unknown device '{requested}', fallback to CPU."


def collect_torch_runtime_diagnostics() -> str:
    lines: list[str] = []
    lines.append(f"torch={torch.__version__}")
    lines.append(f"torch.version.cuda={torch.version.cuda}")

    try:
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                message="CUDA initialization: Unexpected error from cudaGetDeviceCount\\(\\).*",
                category=UserWarning,
            )
            cuda_available = torch.cuda.is_available()
        lines.append(f"torch.cuda.is_available()={cuda_available}")
    except Exception as exc:  # pragma: no cover - hardware dependent
        cuda_available = False
        lines.append(f"torch.cuda.is_available() raised={exc}")

    try:
        cuda_count = torch.cuda.device_count()
        lines.append(f"torch.cuda.device_count()={cuda_count}")
    except Exception as exc:  # pragma: no cover - hardware dependent
        cuda_count = 0
        lines.append(f"torch.cuda.device_count() raised={exc}")

    if cuda_count > 0:
        try:
            lines.append(f"torch.cuda.get_device_name(0)={torch.cuda.get_device_name(0)}")
        except Exception as exc:  # pragma: no cover - hardware dependent
            lines.append(f"torch.cuda.get_device_name(0) raised={exc}")
        try:
            lines.append(f"torch.cuda.get_device_capability(0)={torch.cuda.get_device_capability(0)}")
        except Exception as exc:  # pragma: no cover - hardware dependent
            lines.append(f"torch.cuda.get_device_capability(0) raised={exc}")

    mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
    lines.append(f"torch.backends.mps.is_available()={mps_available}")

    return " | ".join(lines)


def latest_model_path(models_dir: Path) -> Optional[Path]:
    candidates = [p for p in models_dir.rglob("*.zip") if p.is_file()]
    if not candidates:
        return None
    return max(candidates, key=lambda p: p.stat().st_mtime)


def write_latest_pointer(models_dir: Path, model_path: Path) -> Path:
    pointer = models_dir / "LATEST_MODEL.txt"
    pointer.write_text(str(model_path.resolve()) + "\n", encoding="utf-8")
    return pointer
