import os
import sys
import json
import pickle
import random
import tempfile
from pathlib import Path

import torch
import gradio as gr
from pyharp import ModelCard, build_endpoint, load_midi, save_midi

sys.path.insert(0, str(Path(__file__).parent / "improvnet"))
sys.path.insert(0, str(Path(__file__).parent))

os.environ["TOKENIZERS_PARALLELISM"] = "true"

ARTIFACT_FOLDER = Path(__file__).parent / "artifacts"
model = None
tokenizer = None
decode_tokenizer = None
configs = None

def download_artifacts():
    if not (ARTIFACT_FOLDER / "style_transfer" / "fine_tuned_model").exists():
        print("Downloading artifacts from Google Drive...", flush=True)
        import gdown
        import zipfile
        zip_path = Path(__file__).parent / "artifacts.zip"
        gdown.download(
            "https://drive.google.com/uc?id=11H3y2sFUFldf6nS5pSpk8B-bIDHtFH4K",
            str(zip_path),
            quiet=False
        )
        print("Extracting artifacts...", flush=True)
        with zipfile.ZipFile(zip_path, "r") as z:
            z.extractall(str(Path(__file__).parent))
        zip_path.unlink()
        print("Artifacts ready.", flush=True)

download_artifacts()

def get_model():
    global model, tokenizer, decode_tokenizer, configs
    if model is None:
        import yaml
        from transformers import EncoderDecoderModel

        config_path = Path(__file__).parent / "configs" / "config_style_transfer.yaml"
        with open(config_path, "r") as f:
            configs = yaml.safe_load(f)

        tokenizer_path = ARTIFACT_FOLDER / "style_transfer" / "vocab_corrupted.pkl"
        with open(tokenizer_path, "rb") as f:
            tokenizer = pickle.load(f)
        decode_tokenizer = {v: k for k, v in tokenizer.items()}

        print("Loading ImprovNet model...", flush=True)
        model = EncoderDecoderModel.from_pretrained(
            str(ARTIFACT_FOLDER / "style_transfer" / "fine_tuned_model")
        )
        model.eval()
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model.to(device)
        print("Model loaded.", flush=True)

    return model, tokenizer, decode_tokenizer, configs


model_card = ModelCard(
    name="ImprovNet",
    description="Generate expressive musical improvisations from piano MIDI. Supports Classical→Jazz and Classical→Classical style transfer.",
    author="Keshav Bhandari, Sungkyun Chang, Tongyu Lu, Fareza R. Enus, Louis B. Bradshaw, Dorien Herremans, Simon Colton",
    tags=["midi", "improvisation", "style-transfer", "jazz", "classical", "piano"],
)


def process_fn(input_midi_path: str, convert_to: str, num_passes: int, corruption_rate: float, t_segment_start: int) -> str:
    print(f"Processing: convert_to={convert_to}, passes={num_passes}, rate={corruption_rate}", flush=True)

    fusion_model, tok, decode_tok, cfg = get_model()

    from generation import generate

    passes = {}
    corruption_types = ["skyline", "skyline", "pitch_velocity_mask", "incorrect_transposition",
                        "permute_pitches", "note_modification", "onset_duration_mask", "fragmentation",
                        "whole_mask", "random"]
    for i in range(num_passes):
        passes[f"pass_{i+1}"] = {
            "corruption_rate": corruption_rate,
            "corruption_type": corruption_types[i % len(corruption_types)]
        }

    cfg["generation"]["passes"] = passes
    cfg["generation"]["convert_to"] = convert_to
    cfg["generation"]["t_segment_start"] = t_segment_start
    cfg["generation"]["novel_peaks_pct"] = 0.0
    cfg["generation"]["write_intermediate_passes"] = False
    cfg["generation"]["context_before"] = 5
    cfg["generation"]["context_after"] = 5
    cfg["generation"]["temperature"] = 1.0
    cfg["generation"]["end_original"] = True
    cfg["generation"]["t_segment_stop"] = -1

    output_dir = Path(tempfile.mkdtemp())

    generate(
        midi_file_path=input_midi_path,
        audio_file_path=None,
        fusion_model=fusion_model,
        configs=cfg,
        novel_peaks_pct=0.0,
        t_segment_start=t_segment_start,
        convert_to=convert_to,
        context_before=5,
        context_after=5,
        corruption_passes=passes,
        tokenizer=tok,
        decode_tokenizer=decode_tok,
        output_folder=str(output_dir),
        save_original=False,
        quiet=False,
        write_intermediate_passes=False,
        temperature=1.0,
        end_original=True,
        t_segment_stop=-1,
    )

    # Find output MIDI
    output_files = list(output_dir.glob("*.mid")) + list(output_dir.glob("*.midi"))
    if not output_files:
        raise ValueError("No output MIDI generated.")

    output_path = tempfile.mktemp(suffix=".mid")
    import shutil
    shutil.copy(output_files[0], output_path)
    print("Done.", flush=True)
    return output_path


with gr.Blocks() as demo:
    input_components = [
        gr.File(
            type="filepath",
            label="Input Piano MIDI",
            file_types=[".mid", ".midi"],
        ).harp_required(True),
        gr.Dropdown(
            choices=["jazz", "classical"],
            value="jazz",
            label="Convert To",
        ),
        gr.Slider(minimum=1, maximum=10, step=1, value=5, label="Number of Passes"),
        gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Corruption Rate"),
        gr.Slider(minimum=0, maximum=10, step=1, value=2, label="Start Segment (5s each)"),
    ]
    output_components = [
        gr.File(
            type="filepath",
            label="Output MIDI",
            file_types=[".mid", ".midi"],
        ).set_info("Improvised MIDI output."),
    ]
    app = build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

print("Launching Gradio...", flush=True)
demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True, pwa=True)