import gradio as gr
import subprocess
import os 
import shutil
import tempfile
import torch
import logging
import numpy as np
from concurrent.futures import ThreadPoolExecutor
from functools import lru_cache

# 로깅 설정
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('yue_generation.log'),
        logging.StreamHandler()
    ]
)

# GPU 설정 최적화
def optimize_gpu_settings():
    if torch.cuda.is_available():
        # L40S에 최적화된 설정
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.deterministic = False
        torch.backends.cudnn.enabled = True
        
        # GPU 메모리 설정
        torch.cuda.empty_cache()
        torch.cuda.set_device(0)
        
        logging.info(f"Using GPU: {torch.cuda.get_device_name(0)}")
        logging.info(f"Available GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
    else:
        logging.warning("GPU not available!")

# flash-attn 설치 함수 개선
def install_flash_attn():
    try:
        logging.info("Installing flash-attn...")
        subprocess.run(
            ["pip", "install", "flash-attn", "--no-build-isolation"], 
            check=True,
            capture_output=True
        )
        logging.info("flash-attn installed successfully!")
    except subprocess.CalledProcessError as e:
        logging.error(f"Failed to install flash-attn: {e}")
        raise

# 초기화 함수
def initialize_system():
    optimize_gpu_settings()
    install_flash_attn()
    
    from huggingface_hub import snapshot_download
    
    # xcodec_mini_infer 폴더 생성
    folder_path = './inference/xcodec_mini_infer'
    os.makedirs(folder_path, exist_ok=True)
    logging.info(f"Created folder at: {folder_path}")

    # 모델 다운로드
    snapshot_download(
        repo_id="m-a-p/xcodec_mini_infer",
        local_dir="./inference/xcodec_mini_infer",
        resume_download=True
    )

    # inference 디렉토리로 이동
    try:
        os.chdir("./inference")
        logging.info(f"Working directory changed to: {os.getcwd()}")
    except FileNotFoundError as e:
        logging.error(f"Directory error: {e}")
        raise

# 캐시를 활용한 파일 관리
@lru_cache(maxsize=100)
def get_cached_file_path(content_hash, prefix):
    return create_temp_file(content_hash, prefix)

def empty_output_folder(output_dir):
    try:
        shutil.rmtree(output_dir)
        os.makedirs(output_dir)
        logging.info(f"Output folder cleaned: {output_dir}")
    except Exception as e:
        logging.error(f"Error cleaning output folder: {e}")
        raise

def create_temp_file(content, prefix, suffix=".txt"):
    temp_file = tempfile.NamedTemporaryFile(delete=False, mode="w", prefix=prefix, suffix=suffix)
    content = content.strip() + "\n\n"
    content = content.replace("\r\n", "\n").replace("\r", "\n")
    temp_file.write(content)
    temp_file.close()
    logging.debug(f"Temporary file created: {temp_file.name}")
    return temp_file.name

def get_last_mp3_file(output_dir):
    mp3_files = [f for f in os.listdir(output_dir) if f.endswith('.mp3')]
    if not mp3_files:
        logging.warning("No MP3 files found")
        return None
    
    mp3_files_with_path = [os.path.join(output_dir, f) for f in mp3_files]
    mp3_files_with_path.sort(key=os.path.getmtime, reverse=True)
    return mp3_files_with_path[0]

# L40S에 최적화된 추론 함수
def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
    try:
        # 임시 파일 생성
        genre_txt_path = create_temp_file(genre_txt_content, prefix="genre_")
        lyrics_txt_path = create_temp_file(lyrics_txt_content, prefix="lyrics_")
        
        output_dir = "./output"
        os.makedirs(output_dir, exist_ok=True)
        empty_output_folder(output_dir)

        # L40S에 최적화된 명령어
        command = [
            "python", "infer.py",
            "--stage1_model", "m-a-p/YuE-s1-7B-anneal-en-cot",
            "--stage2_model", "m-a-p/YuE-s2-1B-general",
            "--genre_txt", genre_txt_path,
            "--lyrics_txt", lyrics_txt_path,
            "--run_n_segments", str(num_segments),
            "--stage2_batch_size", "8",  # L40S에 맞게 증가
            "--output_dir", output_dir,
            "--cuda_idx", "0",
            "--max_new_tokens", str(max_new_tokens),
            "--disable_offload_model",
            "--use_flash_attention_2",  # Flash Attention 2 활성화
            "--bf16"  # BF16 정밀도 사용
        ]

        # CUDA 환경 변수 설정
        env = os.environ.copy()
        env.update({
            "CUDA_VISIBLE_DEVICES": "0",
            "CUDA_HOME": "/usr/local/cuda",
            "PATH": f"/usr/local/cuda/bin:{env.get('PATH', '')}",
            "LD_LIBRARY_PATH": f"/usr/local/cuda/lib64:{env.get('LD_LIBRARY_PATH', '')}",
            "PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:512"
        })

        # 명령 실행
        process = subprocess.run(command, env=env, check=True, capture_output=True)
        logging.info("Inference completed successfully")

        # 결과 처리
        last_mp3 = get_last_mp3_file(output_dir)
        if last_mp3:
            logging.info(f"Generated audio file: {last_mp3}")
            return last_mp3
        else:
            logging.warning("No output audio file generated")
            return None

    except Exception as e:
        logging.error(f"Inference error: {e}")
        raise
    finally:
        # 임시 파일 정리
        for file in [genre_txt_path, lyrics_txt_path]:
            try:
                os.remove(file)
                logging.debug(f"Removed temporary file: {file}")
            except Exception as e:
                logging.warning(f"Failed to remove temporary file {file}: {e}")

# Gradio 인터페이스
with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("# YuE: Open Music Foundation Models for Full-Song Generation (L40S Optimized)")
        gr.HTML("""
        <div style="display:flex;column-gap:4px;">
            <a href="https://github.com/multimodal-art-projection/YuE">
                <img src='https://img.shields.io/badge/GitHub-Repo-blue'>
            </a> 
            <a href="https://map-yue.github.io">
                <img src='https://img.shields.io/badge/Project-Page-green'>
            </a>
        </div>
        """)
        
        with gr.Row():
            with gr.Column():
                genre_txt = gr.Textbox(
                    label="Genre",
                    placeholder="Enter music genre and style descriptions..."
                )
                lyrics_txt = gr.Textbox(
                    label="Lyrics",
                    placeholder="Enter song lyrics...",
                    lines=10
                )
                
            with gr.Column():
                num_segments = gr.Number(
                    label="Number of Song Segments",
                    value=2,
                    minimum=1,
                    maximum=4,
                    step=1,
                    interactive=True
                )
                max_new_tokens = gr.Slider(
                    label="Max New Tokens",
                    minimum=500,
                    maximum=32000,  # L40S의 큰 메모리를 활용
                    step=500,
                    value=4000,
                    interactive=True
                )
                submit_btn = gr.Button("Generate Music", variant="primary")
                music_out = gr.Audio(label="Generated Audio")

        gr.Examples(
            examples=[
                [
                    "female blues airy vocal bright vocal piano sad romantic guitar jazz",
                    """[verse]
In the quiet of the evening, shadows start to fall
Whispers of the night wind echo through the hall
Lost within the silence, I hear your gentle voice
Guiding me back homeward, making my heart rejoice

[chorus]
Don't let this moment fade, hold me close tonight
With you here beside me, everything's alright
Can't imagine life alone, don't want to let you go
Stay with me forever, let our love just flow
                    """
                ],
                [
                    "rap piano street tough piercing vocal hip-hop synthesizer clear vocal male",
                    """[verse]
Woke up in the morning, sun is shining bright
Chasing all my dreams, gotta get my mind right
City lights are fading, but my vision's clear
Got my team beside me, no room for fear
Walking through the streets, beats inside my head
Every step I take, closer to the bread
People passing by, they don't understand
Building up my future with my own two hands

[chorus]
This is my life, and I'm aiming for the top
Never gonna quit, no, I'm never gonna stop
Through the highs and lows, I'mma keep it real
Living out my dreams with this mic and a deal
                    """
                ]
            ],
            inputs=[genre_txt, lyrics_txt]
        )

    # 시스템 초기화
    initialize_system()
    
    # 이벤트 핸들러
    submit_btn.click(
        fn=infer,
        inputs=[genre_txt, lyrics_txt, num_segments, max_new_tokens],
        outputs=[music_out]
    )

# 서버 설정으로 실행
demo.queue(concurrency_count=2).launch(
    server_name="0.0.0.0",
    server_port=7860,
    share=True,
    enable_queue=True,
    show_api=True,
    show_error=True
)