from __future__ import annotations

import threading
from pathlib import Path

import torch
from midigpt import Score
from midigpt.inference import (
    GenerationRequest,
    InferenceConfig,
    InferenceEngine,
    TrackPrompt,
)

MODEL_REPOSITORY = "Metacreation/MIDI-GPT"
MODEL_NAME = "yellow_small"
MODEL_FILENAME = "yellow_small-final.safetensors"

_engine: InferenceEngine | None = None
_engine_lock = threading.Lock()


def _patch_searchsorted_device_mismatch() -> None:
    if getattr(torch.searchsorted, "_midigpt_device_compat", False):
        return

    original = torch.searchsorted

    def device_safe_searchsorted(sorted_sequence, values, *args, **kwargs):
        if isinstance(values, torch.Tensor) and values.device != sorted_sequence.device:
            values = values.to(sorted_sequence.device)
        return original(sorted_sequence, values, *args, **kwargs)

    device_safe_searchsorted._midigpt_device_compat = True
    torch.searchsorted = device_safe_searchsorted


_patch_searchsorted_device_mismatch()


def _get_engine() -> InferenceEngine:
    global _engine
    if _engine is not None:
        return _engine

    with _engine_lock:
        if _engine is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
            _engine = InferenceEngine.from_pretrained(
                MODEL_REPOSITORY,
                filename=MODEL_FILENAME,
                device=device,
            )
    return _engine


def rewrite_bars(
    input_path: Path,
    output_path: Path,
    track_index: int,
    start_bar: int,
    bar_count: int,
    temperature: float,
    top_p: float,
    seed: int,
    polyphony_limit: int,
) -> dict:
    score = Score.from_midi(str(input_path))
    track_count = len(score.tracks)

    if track_count == 0:
        raise ValueError("The MIDI file does not contain any tracks.")
    if track_index < 0 or track_index >= track_count:
        raise ValueError(
            f"Track index {track_index} is out of range. "
            f"This file has {track_count} tracks (0-{track_count - 1})."
        )

    bar_counts = [len(track.bars) for track in score.tracks]
    if len(set(bar_counts)) != 1:
        raise ValueError(
            "MIDI-GPT requires every track to span the same number of bars."
        )

    total_bars = bar_counts[0]
    if total_bars < 4:
        raise ValueError(
            f"MIDI-GPT requires at least 4 bars. This file has {total_bars}."
        )
    if start_bar < 0 or start_bar + bar_count > total_bars:
        raise ValueError(
            f"Bars {start_bar}-{start_bar + bar_count - 1} are outside "
            f"the available range 0-{total_bars - 1}."
        )

    target_bars = list(range(start_bar, start_bar + bar_count))
    original_notes = [
        [
            (
                note.pitch,
                note.velocity,
                note.onset_ticks,
                note.duration_ticks,
            )
            for note in score.tracks[track_index].bars[index].notes
        ]
        for index in target_bars
    ]
    prompts = []
    for index in range(track_count):
        prompts.append(
            TrackPrompt(
                id=index,
                bars=target_bars if index == track_index else [],
            )
        )

    model_dim = 4 if total_bars < 8 else 8
    request = GenerationRequest(
        tracks=prompts,
        config=InferenceConfig(
            temperature=temperature,
            top_p=top_p,
            seed=seed,
            model_dim=model_dim,
            mask_mode="attention",
            bars_per_step=1,
            polyphony_hard_limit=polyphony_limit,
            novelty_check=False,
        ),
    )

    result = _get_engine().session(score, request).run()
    output_path.parent.mkdir(parents=True, exist_ok=True)
    result.to_midi(str(output_path))

    generated_notes = sum(
        len(result.tracks[track_index].bars[index].notes)
        for index in target_bars
    )
    generated_content = [
        [
            (
                note.pitch,
                note.velocity,
                note.onset_ticks,
                note.duration_ticks,
            )
            for note in result.tracks[track_index].bars[index].notes
        ]
        for index in target_bars
    ]
    return {
        "model": MODEL_NAME,
        "track_count": track_count,
        "total_bars": total_bars,
        "rewritten_track": track_index,
        "rewritten_bars": target_bars,
        "generated_notes": generated_notes,
        "content_changed": generated_content != original_notes,
        "seed": seed,
    }