import os
import io
import gradio as gr
from PIL import Image
from huggingface_hub import InferenceClient

# 1. Access secure system keys
hf_token = os.environ.get("HF_TOKEN")
MODEL_ID = "guangyangmusic/legato-small"

print("Hugging Face API Pipeline active. Forwarding inference tasks...")

# 2. Build the explicit serverless inference router
def run_legato_omr(image):
    if image is None:
        return "Please upload a sheet music image first!"
    
    try:
        # Convert incoming canvas/image assets to generic bytes array
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        image_bytes = buffered.getvalue()
        
        # Instantiate localized pipeline execution clients
        client = InferenceClient(model=MODEL_ID, token=hf_token, timeout=30)
        
        print("Streaming asset matrices to Hugging Face cluster...")
        response = client.post(data=image_bytes, model=MODEL_ID)
        result_text = response.decode("utf-8")
        
        if not result_text or "error" in result_text.lower():
            return f"Cluster Warmup Alert: The remote model is currently spinning up its parameters. Please resubmit this image in 10-15 seconds.\n\nDetails: {result_text}"
            
        return result_text
            
    except Exception as e:
        return (
            f"API Communication Exception: {str(e)}\n\n"
            "💡 Troubleshooting Tips:\n"
            "1. If a 'Timeout' occurred, the remote server is spinning up. Try again in a minute.\n"
            "2. Ensure your Space Secret Key is exactly named HF_TOKEN and contains a valid 'Read' token."
        )

# 3. Formulate the visual block interface
with gr.Blocks() as demo:
    gr.Markdown("# 🎼 LEGATO End-to-End OMR Engine (API Cluster Layer)")
    gr.Markdown("Proxied via serverless API hooks to bypass local 16GB CPU limits. Conversions take 2-5 seconds per line snippet.")
    
    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.Column():
            output_text = gr.Textbox(label="2. LEGATO Output (ABC Notation)", lines=12)
            
    submit_btn.click(fn=run_legato_omr, inputs=input_image, outputs=output_text)

# 4. FIXED: Removed 'ssr' parameter to prevent version mismatch crash
if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0", 
        server_port=7860,
        ssr=False
    )