import os
import subprocess
import sys

# `spaces` must be imported before torch so ZeroGPU can patch CUDA init.
import spaces

# ---------------------------------------------------------------------------
# 1. Bootstrap the legato codebase.
#    The repo ships no setup.py/pyproject.toml, so it cannot be pip-installed.
#    Clone it and put the package directory on sys.path instead.
# ---------------------------------------------------------------------------
LEGATO_DIR = "/home/user/app/legato_src"
LEGATO_REPO = "https://github.com/guang-yng/legato.git"

if not os.path.isdir(LEGATO_DIR):
    print("Cloning legato codebase...")
    subprocess.run(["git", "clone", "--depth", "1", LEGATO_REPO, LEGATO_DIR], check=True)

if LEGATO_DIR not in sys.path:
    sys.path.insert(0, LEGATO_DIR)

import gradio as gr
import torch
from transformers import AutoProcessor, GenerationConfig

from legato.models import LegatoModel

# ---------------------------------------------------------------------------
# 2. Load the model once at startup.
#    The vision encoder pulls from the gated meta-llama/Llama-3.2-11B-Vision,
#    so HF_TOKEN must belong to an account that accepted the Llama 3.2 license.
# ---------------------------------------------------------------------------
MODEL_ID = "guangyangmusic/legato-small"
HF_TOKEN = os.environ.get("HF_TOKEN")

print(f"Loading {MODEL_ID}...")
model = LegatoModel.from_pretrained(MODEL_ID, token=HF_TOKEN)
processor = AutoProcessor.from_pretrained(MODEL_ID, token=HF_TOKEN)

# fp16 halves the ~15GB full-precision footprint. Under ZeroGPU the .to("cuda")
# is deferred until the first @spaces.GPU call, which is the supported pattern.
model = model.half().to("cuda")
model.eval()
print("Model loaded.")


# ---------------------------------------------------------------------------
# 3. Inference. The @spaces.GPU decorator is what ZeroGPU scans for at startup.
# ---------------------------------------------------------------------------
@spaces.GPU(duration=120)
def run_legato_omr(image, num_beams, max_length):
    if image is None:
        return "Please upload a sheet music image first!"

    try:
        # convert() guards against RGBA/grayscale/palette PNGs, which are common
        # for score screenshots and would otherwise fail downstream.
        image = image.convert("RGB")

        inputs = processor(images=image, return_tensors="pt")
        inputs = {k: v.to("cuda") for k, v in inputs.items()}

        generation_config = GenerationConfig(
            max_length=int(max_length),
            num_beams=int(num_beams),
            repetition_penalty=1.1,
        )

        with torch.no_grad():
            outputs = model.generate(**inputs, generation_config=generation_config)

        return processor.batch_decode(outputs, skip_special_tokens=True)[0]

    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        return (
            "Out of GPU memory. Try a smaller image crop, or lower the beam count "
            "in Advanced settings."
        )
    except Exception as e:
        return f"Inference failed: {type(e).__name__}: {e}"


# ---------------------------------------------------------------------------
# 4. Interface
# ---------------------------------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("# 🎼 LEGATO End-to-End OMR")
    gr.Markdown(
        "Transcribes typeset sheet music into ABC notation. "
        "Works best on a single clean staff line or a short snippet; "
        "handwritten and low-quality scans are out of scope for this model."
    )

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="1. Upload Sheet Music Snippet")
            submit_btn = gr.Button("Analyze Layout & Transcribe", variant="primary")

            with gr.Accordion("Advanced settings", open=False):
                num_beams = gr.Slider(
                    1, 10, value=4, step=1,
                    label="Beam count",
                    info="The model card uses 10. Higher is more accurate but "
                         "much slower — 10 can exceed the ZeroGPU time limit.",
                )
                max_length = gr.Slider(
                    256, 2048, value=2048, step=256,
                    label="Max tokens",
                )

        with gr.Column():
            output_text = gr.Textbox(
                label="2. LEGATO Output (ABC Notation)", lines=12, show_copy_button=True
            )

    submit_btn.click(
        fn=run_legato_omr,
        inputs=[input_image, num_beams, max_length],
        outputs=output_text,
    )

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)