import gradio as gr
import torch
import numpy as np
from transformers import AutoProcessor, MusicgenForConditionalGeneration

# 모델 로드 (첫 실행 시 다운로드 및 로딩 시간이 소요됩니다.)
print("모델을 로드하는 중입니다... (facebook/musicgen-melody)")
device = "cuda" if torch.cuda.is_available() else "cpu"
model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-melody").to(device)
processor = AutoProcessor.from_pretrained("facebook/musicgen-melody")
print("모델 로딩 완료!")

def generate_music(text_prompt, audio_melody):
    """
    텍스트 프롬프트와 선택적인 멜로디 오디오를 받아 음악을 생성하는 함수.
    """
    if not text_prompt:
        return None, "Please enter a text description."

    try:
        # 입력 전처리
        # audio_melody가 None이면 텍스트 기반으로만 생성합니다.
        inputs = processor(
            text=[text_prompt],
            audio=audio_melody,
            sampling_rate=processor.sampling_rate,
            padding=True,
            return_tensors="pt",
        )

        # 모델 생성 (GPU 사용 가능 시 GPU로 이동)
        inputs = {k: v.to(device) for k, v in inputs.items()}
        audio_values = model.generate(**inputs, max_new_tokens=256)

        # 결과를 numpy 배열로 변환 (Gradio 오디오 출력을 위해)
        sampling_rate = model.config.audio_encoder.sampling_rate
        audio_data = audio_values[0].cpu().numpy()
        
        return (sampling_rate, audio_data), "생성이 완료되었습니다!"
    
    except Exception as e:
        return None, f"오류가 발생했습니다: {str(e)}"

# Gradio 6 애플리케이션 정의
with gr.Blocks() as demo:
    # 헤더 섹션
    gr.Markdown(
        """
        # 🎵 AI 음악 생성기 (MusicGen Melody)
        텍스트 설명을 입력하거나 자신만의 멜로디를 업로드하여 AI가 음악을 만들어드립니다.
        """
    )
    
    # 필수 요구사항: Built with anycoder 링크
    gr.HTML(
        '<div style="text-align: center; margin-bottom: 20px;">'
        '<a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #f97316; font-weight: bold; text-decoration: none;">Built with anycoder</a>'
        '</div>'
    )

    with gr.Row():
        with gr.Column(scale=1):
            # 입력 섹션
            text_input = gr.Textbox(
                label="📝 음악 설명 (Prompt)",
                placeholder="예: A relaxing jazz song with piano and saxophone, 80s pop song with heavy synth",
                lines=3,
                max_lines=5
            )
            
            melody_input = gr.Audio(
                label="🎼 참고 멜로디 (선택사항)",
                sources=["upload", "microphone"],
                type="numpy",
                info="특정 멜로디를 기반으로 생성하고 싶다면 오디오를 업로드하세요."
            )
            
            generate_btn = gr.Button("🎶 음악 생성하기", variant="primary", size="lg")

        with gr.Column(scale=1):
            # 출력 섹션
            status_output = gr.Textbox(label="상태", interactive=False)
            audio_output = gr.Audio(label="🎧 생성된 음악", type="numpy")

    # 예시 섹션
    gr.Examples(
        examples=[
            ["A happy pop song with synth sounds", None],
            ["Sad violin melody in a rainy day", None],
            ["Epic orchestral trailer music with drums", None],
            ["Lo-fi hip hop beats for studying", None],
        ],
        inputs=[text_input, melody_input],
    )

    # 이벤트 리스너 연결
    generate_btn.click(
        fn=generate_music,
        inputs=[text_input, melody_input],
        outputs=[audio_output, status_output],
        api_visibility="public"
    )

# Gradio 6 런치 메서드 (테마 및 기타 설정 여기서 적용)
demo.launch(
    theme=gr.themes.Soft(
        primary_hue="orange",
        secondary_hue="red",
        neutral_hue="slate",
        text_size="lg",
        spacing_size="lg",
        radius_size="md"
    ),
    footer_links=[
        {"label": "Hugging Face Model", "url": "https://huggingface.co/facebook/musicgen-melody"},
        {"label": "Gradio", "url": "https://gradio.app"}
    ]
)