import os
import base64
import inspect
from queue import Queue
import threading

import torch
from PIL import Image
from transformers import AutoProcessor, MiniCPMV4_6ForConditionalGeneration
import gradio as gr

try:
    import spaces
except ImportError:
    # Allows the app to run locally without the Hugging Face ZeroGPU package.
    class _SpacesFallback:
        @staticmethod
        def GPU(*args, **kwargs):
            def decorator(fn):
                return fn
            return decorator

    spaces = _SpacesFallback()


ORIGINAL_MODEL_ID = "openbmb/MiniCPM-V-4.6"
FINETUNED_MODEL_ID = "jon-fernandes/noteworthy"

GPT_MODEL_ID = "gpt-5.5"
GPT_FALLBACK_MODEL_ID = "gpt-5.4-mini"
GPT_MAX_COMPLETION_TOKENS = 4096
GPT_REASONING_EFFORT = "none"

NOTES_PROMPT = "Transcribe the musical notes in this image. Return only the transcription."

# For ZeroGPU:
# large = default ZeroGPU size.
# xlarge = more VRAM, but uses more quota and may queue longer.
# If you get CUDA out-of-memory errors, add this Space variable:
# ZERO_GPU_SIZE=xlarge
ZERO_GPU_SIZE = os.environ.get("ZERO_GPU_SIZE", "large")

CAMERA_CAPTURE_JS = """
function () {
  const attachTapCapture = () => {
    const root = document.getElementById("sheet-music-input");
    if (!root || root.dataset.tapCaptureReady === "1") {
      return Boolean(root);
    }

    root.dataset.tapCaptureReady = "1";
    root.addEventListener("click", (event) => {
      if (event.target.closest("button, input, select, textarea, a")) {
        return;
      }

      if (!root.querySelector("video")) {
        return;
      }

      const buttons = Array.from(root.querySelectorAll("button"));
      const captureButton = buttons.find((button) => {
        const text = [
          button.textContent,
          button.getAttribute("aria-label"),
          button.getAttribute("title"),
        ].filter(Boolean).join(" ").toLowerCase();
        return text.includes("capture") || text.includes("photo") || text.includes("snapshot");
      });

      if (captureButton) {
        captureButton.click();
      }
    });

    return true;
  };

  if (!attachTapCapture()) {
    const timer = setInterval(() => {
      if (attachTapCapture()) {
        clearInterval(timer);
      }
    }, 300);
    setTimeout(() => {
      clearInterval(timer);
    }, 5000);
  }
}
"""

CUSTOM_CSS = """
#sheet-music-input video,
#sheet-music-input canvas,
#sheet-music-input img {
  cursor: pointer;
}
"""


def env_flag(name: str, default: bool = False) -> bool:
    value = os.environ.get(name)

    if value is None:
        return default

    return value.strip().lower() in {"1", "true", "yes", "on"}


# Important for ZeroGPU:
# Keep this False by default. GPU is allocated when @spaces.GPU functions run,
# so startup warmup can cause unnecessary errors or quota use.
ENABLE_MODEL_WARMUP = env_flag("NOTEWORTHY_WARMUP", False)


def supports_keyword(callable_obj, keyword):
    try:
        signature = inspect.signature(callable_obj)
    except (TypeError, ValueError):
        return False

    return keyword in signature.parameters


print("Loading processor...")
processor = AutoProcessor.from_pretrained(
    ORIGINAL_MODEL_ID,
    trust_remote_code=True,
)

MODEL_LOAD_ERRORS = {}


def load_local_model(label, model_id):
    print(f"Loading {label} model...")

    try:
        model = MiniCPMV4_6ForConditionalGeneration.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
            attn_implementation="sdpa",
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )

        # Important for ZeroGPU:
        # Hugging Face recommends placing the model on CUDA at module level.
        # ZeroGPU uses CUDA emulation during startup and real CUDA inside @spaces.GPU.
        model = model.to("cuda").eval()

        print(f"{label} model loaded.")
        return model

    except Exception as e:
        message = f"{type(e).__name__}: {e}"
        MODEL_LOAD_ERRORS[label] = message
        print(f"Failed to load {label} model: {message}")
        return None


original_model = load_local_model("original", ORIGINAL_MODEL_ID)
finetuned_model = load_local_model("fine-tuned", FINETUNED_MODEL_ID)

print("Models loaded.")


def _get_model_device(model):
    try:
        return next(model.parameters()).device
    except StopIteration:
        return torch.device("cuda")


def _move_model_inputs(inputs, device):
    moved = {}

    for key, value in inputs.items():
        if isinstance(value, torch.Tensor):
            if torch.is_floating_point(value):
                value = value.to(dtype=torch.bfloat16)

            moved[key] = value.to(device)
        else:
            moved[key] = value

    return moved


def _build_model_inputs(image: Image.Image):
    input_variants = [
        (
            [
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "image": image},
                        {"type": "text", "text": NOTES_PROMPT},
                    ],
                }
            ],
            {},
        ),
        (
            [
                {
                    "role": "user",
                    "content": [
                        {"type": "image"},
                        {"type": "text", "text": NOTES_PROMPT},
                    ],
                }
            ],
            {"images": [image]},
        ),
    ]

    errors = []

    for messages, extra_processor_kwargs in input_variants:
        try:
            inputs = processor.apply_chat_template(
                messages,
                add_generation_prompt=True,
                tokenize=True,
                return_dict=True,
                return_tensors="pt",
                enable_thinking=False,
                processor_kwargs={
                    **extra_processor_kwargs,
                    "downsample_mode": "4x",
                    "max_slice_nums": 9,
                    "use_image_id": True,
                },
            )

            if hasattr(inputs, "items"):
                return dict(inputs)

            errors.append(f"Unexpected input type: {type(inputs).__name__}")

        except TypeError as e:
            errors.append(str(e))

            try:
                inputs = processor.apply_chat_template(
                    messages,
                    add_generation_prompt=True,
                    tokenize=True,
                    return_dict=True,
                    return_tensors="pt",
                    processor_kwargs={
                        **extra_processor_kwargs,
                        "downsample_mode": "4x",
                        "max_slice_nums": 9,
                        "use_image_id": True,
                    },
                )

                if hasattr(inputs, "items"):
                    return dict(inputs)

                errors.append(f"Unexpected input type: {type(inputs).__name__}")

            except Exception as fallback_error:
                errors.append(str(fallback_error))

        except Exception as e:
            errors.append(str(e))

    raise RuntimeError("; ".join(errors[-4:]))


def generate_model_text(model, image: Image.Image, max_new_tokens: int):
    if model is None:
        raise RuntimeError("Model failed to load.")

    device = _get_model_device(model)
    inputs = _move_model_inputs(_build_model_inputs(image), device)

    with torch.inference_mode():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            num_beams=1,
            downsample_mode="4x",
        )

    input_ids = inputs.get("input_ids")

    if (
        isinstance(input_ids, torch.Tensor)
        and isinstance(generated_ids, torch.Tensor)
        and generated_ids.shape[-1] > input_ids.shape[-1]
    ):
        generated_ids = generated_ids[:, input_ids.shape[-1]:]

    return processor.tokenizer.batch_decode(
        generated_ids,
        skip_special_tokens=True,
    )[0].strip()


def stream_model(model, image: Image.Image, label: str):
    if model is None:
        yield f"[Error: {label} model failed to load: {MODEL_LOAD_ERRORS.get(label, 'unknown error')}]"
        return

    try:
        yield generate_model_text(
            model,
            image,
            max_new_tokens=1024,
        )

    except Exception as e:
        yield f"[Error: {type(e).__name__}: {e}]"


def stream_model_text(model, image: Image.Image, label: str):
    text = ""

    for chunk in stream_model(model, image, label):
        text += chunk
        yield text


def postprocess_finetuned(text: str) -> str:
    text = text.replace("note-", "")
    text = text.replace("barline", "|")
    text = text.replace("whole", "semibreve")
    text = text.replace("half", "minim")
    text = text.replace("quarter", "crotchet")
    text = text.replace("eighth", "quaver")
    text = text.replace("sixteenth", "semiquaver")
    text = text.replace("thirtysecond", "demisemiquaver")
    return text


def warmup_models():
    if not ENABLE_MODEL_WARMUP:
        print("Model warmup disabled.")
        return

    warmup_path = "examples/000100005-1_1_1.png"

    if not os.path.exists(warmup_path):
        print("Skipping model warmup; example image is missing.")
        return

    print("Warming up local models...")
    image = Image.open(warmup_path).convert("RGB")

    for name, model in (
        ("fine-tuned", finetuned_model),
        ("original", original_model),
    ):
        if model is None:
            print(f"  Skipping {name} warmup; model failed to load.")
            continue

        print(f"  Warming {name} model...")

        try:
            generate_model_text(
                model,
                image,
                max_new_tokens=8,
            )
        except Exception as e:
            print(f"  Warmup failed for {name} model: {type(e).__name__}: {e}")

    print("Model warmup complete.")


warmup_models()


def _stream_openai_transcription(client, model_id, mime, image_data):
    response_stream = client.chat.completions.create(
        model=model_id,
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:{mime};base64,{image_data}",
                            "detail": "low",
                        },
                    },
                    {
                        "type": "text",
                        "text": NOTES_PROMPT,
                    },
                ],
            }
        ],
        max_completion_tokens=GPT_MAX_COMPLETION_TOKENS,
        reasoning_effort=GPT_REASONING_EFFORT,
        stream=True,
    )

    content = ""
    refusal = ""
    finish_reason = None

    for chunk in response_stream:
        if not chunk.choices:
            continue

        choice = chunk.choices[0]
        finish_reason = choice.finish_reason or finish_reason
        delta = choice.delta

        delta_content = getattr(delta, "content", None)

        if delta_content:
            content += delta_content
            yield content, finish_reason

        delta_refusal = getattr(delta, "refusal", None)

        if delta_refusal:
            refusal += delta_refusal
            yield f"[Refusal: {refusal}]", finish_reason

    if not content and not refusal:
        yield "", finish_reason

    elif finish_reason == "length":
        yield (
            f"{content}\n\n[Stopped because the completion token limit was reached.]",
            finish_reason,
        )


def stream_gpt_text(image_path):
    yield f"Calling {GPT_MODEL_ID}..."

    api_key = os.environ.get("OPENAI_API_KEY")

    if not api_key:
        yield "[Error: OPENAI_API_KEY is not set.]"
        return

    try:
        from openai import OpenAI

        with open(image_path, "rb") as f:
            image_data = base64.b64encode(f.read()).decode("utf-8")

        ext = os.path.splitext(image_path)[1].lstrip(".").lower()

        if ext == "jpg":
            ext = "jpeg"

        mime = f"image/{ext}" if ext in ("png", "jpeg", "gif", "webp") else "image/jpeg"

        client = OpenAI(
            api_key=api_key,
            timeout=45,
        )

        last_text = ""
        last_finish_reason = None

        for text, finish_reason in _stream_openai_transcription(
            client,
            GPT_MODEL_ID,
            mime,
            image_data,
        ):
            last_text = text
            last_finish_reason = finish_reason

            if text:
                yield text

        if last_text:
            return

        yield f"{GPT_MODEL_ID} returned no text; trying {GPT_FALLBACK_MODEL_ID}..."

        for text, finish_reason in _stream_openai_transcription(
            client,
            GPT_FALLBACK_MODEL_ID,
            mime,
            image_data,
        ):
            last_text = text
            last_finish_reason = finish_reason

            if text:
                yield text

        if not last_text:
            yield f"[Empty response. finish_reason={last_finish_reason}]"

    except Exception as e:
        yield f"[Error: {e}]"


def _run_stream(index, stream, updates):
    try:
        for text in stream:
            updates.put((index, text))

    except Exception as e:
        updates.put((index, f"[Error: {e}]"))

    finally:
        updates.put((index, None))


@spaces.GPU(duration=180, size=ZERO_GPU_SIZE)
def predict_all(image_path):
    if image_path is None:
        message = "Please upload an image."
        yield message, message, message
        return

    try:
        image = Image.open(image_path).convert("RGB")
    except Exception as e:
        message = f"[Error opening image: {type(e).__name__}: {e}]"
        yield message, message, message
        return

    updates = Queue()
    outputs = ["", "", ""]

    yield outputs[0], outputs[1], outputs[2]

    streams = [
        stream_model_text(finetuned_model, image.copy(), "fine-tuned"),
        stream_model_text(original_model, image.copy(), "original"),
        stream_gpt_text(image_path),
    ]

    threads = [
        threading.Thread(
            target=_run_stream,
            args=(index, stream, updates),
            daemon=True,
        )
        for index, stream in enumerate(streams)
    ]

    for thread in threads:
        thread.start()

    running = len(threads)

    while running:
        index, text = updates.get()

        if text is None:
            running -= 1
            continue

        outputs[index] = postprocess_finetuned(text) if index == 0 else text

        yield outputs[0], outputs[1], outputs[2]

    for thread in threads:
        thread.join()

    if torch.cuda.is_available():
        torch.cuda.empty_cache()


blocks_kwargs = {
    "title": "Noteworthy — Sheet Music Transcription",
    "theme": gr.themes.Soft(),
}

if supports_keyword(gr.Blocks, "css"):
    blocks_kwargs["css"] = CUSTOM_CSS

if supports_keyword(gr.Blocks, "js"):
    blocks_kwargs["js"] = CAMERA_CAPTURE_JS


with gr.Blocks(**blocks_kwargs) as demo:
    gr.Markdown(
        """
        # Noteworthy

        Sheet Music Transcription

        Take a photo or upload sheet music, then click **Transcribe Music** to compare models.
        """
    )

    image_input = gr.Image(
        type="filepath",
        label="Sheet Music Image",
        show_label=False,
        sources=["upload", "webcam", "clipboard"],
        webcam_options=gr.WebcamOptions(
            mirror=False,
            constraints={"facingMode": "environment"},
        ),
        placeholder="Upload sheet music, then click Transcribe Music.",
        elem_id="sheet-music-input",
    )

    gr.Examples(
        examples=[
            ["examples/000100005-1_1_1.png"],
            ["examples/000100014-1_1_1.png"],
            ["examples/000100059-1_1_1.png"],
        ],
        inputs=image_input,
    )

    notes_btn = gr.Button(
        "Transcribe Music",
        variant="primary",
        size="lg",
    )

    with gr.Row():
        finetuned_output = gr.Textbox(
            label="Noteworthy Fine-tuned",
            lines=20,
        )

        original_output = gr.Textbox(
            label="MiniCPM-V-4.6 Original",
            lines=20,
        )

        gpt_output = gr.Textbox(
            label=f"{GPT_MODEL_ID.upper()} no reasoning",
            lines=20,
        )

    notes_btn.click(
        fn=predict_all,
        inputs=[image_input],
        outputs=[
            finetuned_output,
            original_output,
            gpt_output,
        ],
        api_name="transcribe_music",
    )


demo.queue(max_size=20)


launch_kwargs = {
    "server_name": os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
    "server_port": int(os.environ.get("GRADIO_SERVER_PORT", "7860")),
    "share": env_flag("GRADIO_SHARE"),
}

if supports_keyword(demo.launch, "mcp_server"):
    launch_kwargs["mcp_server"] = True


demo.launch(**launch_kwargs)