from __future__ import annotations

import hashlib
import os
import shutil
import sys
import tarfile
import tempfile
import threading
import urllib.request
import zipfile
from pathlib import Path
from types import SimpleNamespace

os.environ.setdefault("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "1")

ROOT = Path(__file__).resolve().parent
CACHE_ROOT = Path(
    os.environ.get("GAME_CACHE_DIR", Path.home() / ".cache" / "game")
)

SOURCE_REVISION = "4ad815c90dfe2442730f3fdc866fd23e737cbc97"
SOURCE_NAME = f"GAME-{SOURCE_REVISION}"
SOURCE_DIR = CACHE_ROOT / SOURCE_NAME
SOURCE_ARCHIVE = CACHE_ROOT / f"{SOURCE_NAME}.tar.gz"
SOURCE_URL = (
    f"https://codeload.github.com/openvpi/GAME/tar.gz/{SOURCE_REVISION}"
)
SOURCE_SHA256 = (
    "b1c1584d2326d6920228695a3b401f6483e6ef50c7d1695b984b63da6ba86f3b"
)

MODEL_NAME = "GAME-1.0-small"
MODEL_FILES = ("model.pt", "config.yaml", "lang_map.json")
MODEL_DIR = CACHE_ROOT / MODEL_NAME
MODEL_ARCHIVE = CACHE_ROOT / f"{MODEL_NAME}.zip"
MODEL_URL = (
    "https://github.com/openvpi/GAME/releases/"
    "download/v1.0.0/GAME-1.0-small.zip"
)
MODEL_SHA256 = (
    "3d3e1ac0a83234b2a163a3d43043455d15670765eaa25ef6285c399da1ccc576"
)

_source_lock = threading.Lock()
_model_archive_lock = threading.Lock()
_runtime_lock = threading.Lock()
_inference_model_lock = threading.Lock()

_runtime: SimpleNamespace | None = None
_inference_model = None
_language_map: dict[str, int] | None = None


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as file:
        for chunk in iter(lambda: file.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _download(
    url: str,
    destination: Path,
    expected_sha256: str,
    label: str,
) -> None:
    CACHE_ROOT.mkdir(parents=True, exist_ok=True)

    if destination.exists():
        if _sha256(destination) == expected_sha256:
            return
        destination.unlink()

    partial = Path(f"{destination}.part")
    partial.unlink(missing_ok=True)
    request = urllib.request.Request(
        url,
        headers={"User-Agent": "TEAMuP-GAME-Space/1.0"},
    )

    print(f"Downloading {label}...", flush=True)
    try:
        with (
            urllib.request.urlopen(request, timeout=60) as response,
            partial.open("wb") as file,
        ):
            shutil.copyfileobj(response, file)

        actual_sha256 = _sha256(partial)
        if actual_sha256 != expected_sha256:
            raise RuntimeError(
                f"{label} failed SHA-256 validation: "
                f"expected {expected_sha256}, received {actual_sha256}"
            )
        partial.replace(destination)
    except Exception:
        partial.unlink(missing_ok=True)
        raise


def _source_complete(path: Path) -> bool:
    required = (
        path / "inference" / "api.py",
        path / "inference" / "callbacks.py",
        path / "inference" / "data.py",
        path / "inference" / "slicer2.py",
        path / "lib" / "config" / "schema.py",
    )
    return all(file.is_file() for file in required)


def _get_source_dir() -> Path:
    local_source = ROOT / "GAME"
    if _source_complete(local_source):
        return local_source
    if _source_complete(SOURCE_DIR):
        return SOURCE_DIR

    with _source_lock:
        if _source_complete(SOURCE_DIR):
            return SOURCE_DIR

        _download(
            SOURCE_URL,
            SOURCE_ARCHIVE,
            SOURCE_SHA256,
            f"GAME source revision {SOURCE_REVISION}",
        )
        with tempfile.TemporaryDirectory(
            prefix="game-source-",
            dir=CACHE_ROOT,
        ) as temporary_dir:
            temporary_path = Path(temporary_dir)
            with tarfile.open(SOURCE_ARCHIVE, "r:gz") as archive:
                archive.extractall(temporary_path, filter="data")

            extracted = temporary_path / SOURCE_NAME
            if not _source_complete(extracted):
                raise RuntimeError(
                    "The GAME source archive is missing inference files."
                )
            if SOURCE_DIR.exists():
                shutil.rmtree(SOURCE_DIR)
            shutil.move(str(extracted), str(SOURCE_DIR))

    return SOURCE_DIR


def _model_complete(path: Path) -> bool:
    return all((path / filename).is_file() for filename in MODEL_FILES)


def _get_model_dir() -> Path:
    local_model = ROOT / "models" / MODEL_NAME
    if _model_complete(local_model):
        return local_model
    if _model_complete(MODEL_DIR):
        return MODEL_DIR

    with _model_archive_lock:
        if _model_complete(MODEL_DIR):
            return MODEL_DIR

        _download(
            MODEL_URL,
            MODEL_ARCHIVE,
            MODEL_SHA256,
            f"{MODEL_NAME} checkpoint",
        )
        with tempfile.TemporaryDirectory(
            prefix="game-model-",
            dir=CACHE_ROOT,
        ) as temporary_dir:
            temporary_path = Path(temporary_dir)
            with zipfile.ZipFile(MODEL_ARCHIVE) as archive:
                archive.extractall(temporary_path)

            extracted = temporary_path / MODEL_NAME
            if not _model_complete(extracted):
                raise RuntimeError(
                    "The GAME checkpoint archive is incomplete."
                )
            if MODEL_DIR.exists():
                shutil.rmtree(MODEL_DIR)
            shutil.move(str(extracted), str(MODEL_DIR))

    return MODEL_DIR


def _get_runtime() -> SimpleNamespace:
    global _runtime
    if _runtime is not None:
        return _runtime

    with _runtime_lock:
        if _runtime is not None:
            return _runtime

        source_path = str(_get_source_dir())
        if source_path not in sys.path:
            sys.path.insert(0, source_path)

        from inference.api import infer_model, load_inference_model
        from inference.callbacks import (
            SaveCombinedMidiFileCallback,
            SaveCombinedTextFileCallback,
        )
        from inference.data import SlicedAudioFileIterableDataset
        from inference.slicer2 import Slicer
        from lib.config.schema import ValidationConfig

        _runtime = SimpleNamespace(
            infer_model=infer_model,
            load_inference_model=load_inference_model,
            MidiCallback=SaveCombinedMidiFileCallback,
            TextCallback=SaveCombinedTextFileCallback,
            Dataset=SlicedAudioFileIterableDataset,
            Slicer=Slicer,
            ValidationConfig=ValidationConfig,
        )

    return _runtime


def _get_model(runtime: SimpleNamespace):
    global _inference_model, _language_map
    if _inference_model is not None:
        return _inference_model, _language_map

    with _inference_model_lock:
        if _inference_model is None:
            _inference_model, _language_map = runtime.load_inference_model(
                _get_model_dir() / "model.pt"
            )

    return _inference_model, _language_map


def _language_id(
    language_code: str,
    language_map: dict[str, int] | None,
) -> int:
    if not language_code:
        return 0
    if language_map is None or language_code not in language_map:
        supported = ", ".join(language_map or ())
        raise ValueError(
            f"Language '{language_code}' is not supported. "
            f"Supported languages: {supported}"
        )
    return language_map[language_code]


def transcribe(
    audio_path: Path,
    output_dir: Path,
    language_code: str,
    steps: int,
) -> None:
    runtime = _get_runtime()
    model, language_map = _get_model(runtime)
    sample_rate = model.inference_config.features.audio_sample_rate

    dataset = runtime.Dataset(
        filemap={audio_path.stem: audio_path},
        samplerate=sample_rate,
        slicer=runtime.Slicer(
            sr=sample_rate,
            threshold=-40.0,
            min_length=1000,
            min_interval=200,
            max_sil_kept=100,
        ),
        language=_language_id(language_code, language_map),
    )
    callbacks = [
        runtime.MidiCallback(output_dir=output_dir, tempo=120),
        runtime.TextCallback(
            output_dir=output_dir,
            file_format="csv",
            pitch_format="name",
            round_pitch=False,
        ),
    ]
    config = runtime.ValidationConfig(
        d3pm_sample_t0=0.0,
        d3pm_sample_steps=steps,
        d3pm_sample_ts=None,
        boundary_decoding_threshold=0.2,
        boundary_decoding_radius=round(0.02 / model.timestep),
        note_presence_threshold=0.2,
    )
    runtime.infer_model(
        model=model,
        dataset=dataset,
        config=config,
        callbacks=callbacks,
        batch_size=1,
        num_workers=0,
        precision="32-true",
    )