import os
import torch
import torchaudio
import psutil
import time
import sys
import numpy as np
import gc
import gradio as gr
from pydub import AudioSegment
from audiocraft.models import MusicGen
from torch.cuda.amp import autocast

# Set PYTORCH_CUDA_ALLOC_CONF to manage memory fragmentation
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:32"

# Setup instructions:
# 1. Create a virtual environment: python3 -m venv venv
# 2. Activate the environment: source venv/bin/activate
# 3. Install dependencies:
#    pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121
#    pip install -r requirements.txt
# 4. Install ffmpeg for MP3 export:
#    sudo apt-get install ffmpeg
# 5. Ensure model weights are in /home/ubuntu/ghostai_music_generator/models/musicgen-medium
# 6. Log in to Hugging Face CLI: huggingface-cli login
# 7. Request access to model if needed: https://huggingface.co/facebook/musicgen-medium

# Check critical dependencies
if np.__version__ != "1.23.5":
    print(f"ERROR: NumPy version {np.__version__} is not compatible. Please install numpy==1.23.5.")
    sys.exit(1)
if not torch.__version__.startswith(("2.1.0", "2.3.1")):
    print(f"WARNING: PyTorch version {torch.__version__} may not be compatible. Expected torch==2.1.0 or 2.3.1.")

# 1) DEVICE SETUP
device = "cuda" if torch.cuda.is_available() else "cpu"
if device != "cuda":
    print("ERROR: CUDA is required for GPU rendering. CPU rendering is disabled to avoid slow performance.")
    sys.exit(1)
print(f"CUDA is available. Using GPU: {torch.cuda.get_device_name(0)}")

# 2) LOAD MUSICGEN INTO VRAM
try:
    print("Loading MusicGen model into VRAM...")
    local_model_path = "/home/ubuntu/ghostai_music_generator/models/musicgen-medium"
    if not os.path.exists(local_model_path):
        print(f"ERROR: Local model path {local_model_path} does not exist. Please ensure the model weights are downloaded.")
        sys.exit(1)
    musicgen_model = MusicGen.get_pretrained(local_model_path, device=device)
except Exception as e:
    print(f"ERROR: Failed to load MusicGen model: {e}")
    print("Please ensure the model weights are in the correct path and dependencies are installed.")
    sys.exit(1)

# 3) RESOURCE MONITORING FUNCTION
def print_resource_usage(stage: str):
    print(f"--- {stage} ---")
    print(f"GPU Memory Allocated: {torch.cuda.memory_allocated() / (1024**3):.2f} GB")
    print(f"GPU Memory Reserved: {torch.cuda.memory_reserved() / (1024**3):.2f} GB")
    print("---------------")

# 4) GENRE PROMPT FUNCTIONS (Updated for distinct song sections with balanced energy)
def set_rock_prompt():
    return "Hard rock with a dynamic intro, expressive verse, and powerful chorus, featuring electric guitars, steady heavy drums, and deep bass"

def set_techno_prompt():
    return "Techno with a pulsing intro, intense build-up, and energetic drop, featuring dark synths, driving bass, and rhythmic fast drums"

def set_jazz_prompt():
    return "Smooth jazz with a warm intro, expressive theme, and lively outro, featuring saxophone, piano, and soft rhythmic drums"

def set_classical_prompt():
    return "Classical orchestral piece with a gentle intro, dramatic development, and grand finale, featuring strings, piano, and subtle rhythmic percussion"

def set_hiphop_prompt():
    return "Hip-hop with a groovy intro, rhythmic verse, and catchy hook, featuring deep bass, tight crisp drums, and funky synths"

# 5) AUDIO PROCESSING FUNCTIONS
def apply_chorus(segment):
    # Enhanced chorus effect for richer sound
    delayed = segment - 4  # Slightly louder delayed copy at -4 dB (was -6 dB)
    delayed = delayed.set_frame_rate(segment.frame_rate)
    return segment.overlay(delayed, position=20)  # Increased delay to 20ms for more noticeable effect

def apply_eq(segment):
    # Apply EQ: Balanced filters for cleaner sound
    segment = segment.low_pass_filter(6000)  # Loosen low-pass filter for brighter highs
    segment = segment.high_pass_filter(100)  # Slightly lower cutoff for low rumble
    return segment

def apply_limiter(segment, max_db=-8.0):
    # Apply a strict limiter to control peaks
    if segment.dBFS > max_db:
        segment = segment - (segment.dBFS - max_db)  # Reduce gain to prevent clipping
    return segment

# 6) GENERATION & I/O FUNCTIONS
def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p: float, temperature: float, total_duration: int, crossfade_duration: int):
    global musicgen_model  # Use global to modify the model instance
    if not instrumental_prompt.strip():
        return None, "⚠️ Please enter a valid instrumental prompt!"
    try:
        # Record start time for total duration calculation
        start_time = time.time()

        # Ensure total duration is within reasonable bounds
        total_duration = min(max(total_duration, 10), 60)  # Between 10 and 60 seconds
        chunk_duration = 15  # Duration of each chunk (e.g., 2 chunks of 15 seconds for a 30-second track)
        num_chunks = max(2, (total_duration + chunk_duration - 1) // chunk_duration)  # Ensure at least 2 chunks
        chunk_duration = total_duration / num_chunks  # Adjust chunk duration to fit total duration

        # Generate slightly longer chunks for overlap
        overlap_duration = min(1.0, crossfade_duration / 1000.0)  # Convert ms to seconds, max 1 second
        generation_duration = chunk_duration + overlap_duration  # Add overlap to each chunk

        # Initialize list to store audio chunks
        audio_chunks = []
        sample_rate = musicgen_model.sample_rate

        # Generate audio in chunks with distinct prompts for song structure
        for i in range(num_chunks):
            # Vary the prompt for each chunk to create clear song sections
            if i == 0:
                chunk_prompt = instrumental_prompt + ", focusing on a dynamic intro and expressive verse"
            else:
                chunk_prompt = instrumental_prompt + ", focusing on a powerful chorus and energetic outro"

            print(f"Generating chunk {i+1}/{num_chunks} on GPU (prompt: {chunk_prompt})...")
            musicgen_model.set_generation_params(
                duration=generation_duration,
                use_sampling=True,
                top_k=top_k,
                top_p=top_p,
                temperature=temperature,
                cfg_coef=cfg_scale
            )

            # Print resource usage before each chunk
            print_resource_usage(f"Before Chunk {i+1} Generation")

            with torch.no_grad():
                with autocast():  # Use mixed precision for lower VRAM usage
                    audio_chunk = musicgen_model.generate([chunk_prompt], progress=True)[0]  # Shape: (2, samples)

            # Normalize to stereo
            audio_chunk = audio_chunk.cpu().to(dtype=torch.float32)
            if audio_chunk.dim() == 1:
                audio_chunk = torch.stack([audio_chunk, audio_chunk], dim=0)  # Mono to stereo
            elif audio_chunk.dim() == 2 and audio_chunk.shape[0] == 1:
                audio_chunk = torch.cat([audio_chunk, audio_chunk], dim=0)
            elif audio_chunk.dim() == 2 and audio_chunk.shape[0] != 2:
                audio_chunk = audio_chunk[:1, :]
                audio_chunk = torch.cat([audio_chunk, audio_chunk], dim=0)
            elif audio_chunk.dim() > 2:
                audio_chunk = audio_chunk.view(2, -1)

            if audio_chunk.shape[0] != 2:
                raise ValueError(f"Expected stereo audio with shape (2, samples), got shape {audio_chunk.shape}")

            # Save chunk temporarily as WAV, then convert to MP3
            temp_wav_path = f"temp_chunk_{i}.wav"
            chunk_path = f"chunk_{i}.mp3"
            torchaudio.save(temp_wav_path, audio_chunk, sample_rate, bits_per_sample=24)
            segment = AudioSegment.from_wav(temp_wav_path)
            segment = apply_limiter(segment, max_db=-8.0)  # Apply limiter early to control peaks
            segment.export(chunk_path, format="mp3", bitrate="320k")
            os.remove(temp_wav_path)  # Clean up temporary WAV file
            audio_chunks.append(chunk_path)

            # Free memory after generating each chunk
            torch.cuda.empty_cache()
            gc.collect()
            time.sleep(0.5)  # Small delay to ensure memory is released
            print_resource_usage(f"After Chunk {i+1} Generation")

        # Combine chunks using pydub with crossfade to smooth transitions
        print("Combining audio chunks...")
        final_segment = AudioSegment.from_mp3(audio_chunks[0])
        for i in range(1, len(audio_chunks)):
            next_segment = AudioSegment.from_mp3(audio_chunks[i])
            # Apply a gentle gain boost to raise quieter sections
            next_segment = next_segment + 2  # Boost by 2 dB to avoid amplitude dips
            next_segment = apply_limiter(next_segment, max_db=-8.0)  # Re-apply limiter after gain boost
            final_segment = final_segment.append(next_segment, crossfade=crossfade_duration)

        # Trim to exact total duration (in milliseconds)
        final_segment = final_segment[:total_duration * 1000]

        # Post-process to enhance audio quality
        print("Post-processing final track...")
        final_segment = apply_eq(final_segment)
        final_segment = apply_chorus(final_segment)

        # Apply final limiter and normalization with large headroom
        final_segment = apply_limiter(final_segment, max_db=-8.0)
        final_segment = final_segment.normalize(headroom=-10.0)  # Large headroom to prevent clipping

        # Export as MP3 only
        mp3_path = "output_cleaned.mp3"
        final_segment.export(
            mp3_path,
            format="mp3",
            bitrate="320k",
            tags={"title": "GhostAI Instrumental", "artist": "GhostAI"}
        )
        print(f"Saved final audio to {mp3_path}")

        # Clean up temporary chunk files
        for chunk_path in audio_chunks:
            os.remove(chunk_path)

        # Print resource usage after generation
        print_resource_usage("After Final Generation")
        print(f"Total Generation Time: {time.time() - start_time:.2f} seconds")

        return mp3_path, "✅ Done!"
    except Exception as e:
        return None, f"❌ Generation failed: {e}"
    finally:
        torch.cuda.empty_cache()
        gc.collect()

def clear_inputs():
    return "", 3.0, 300, 0.95, 1.0, 30, 500

# 7) CUSTOM CSS
css = """
body { 
    background: linear-gradient(135deg, #0A0A0A 0%, #1C2526 100%); 
    color: #E0E0E0; 
    font-family: 'Orbitron', sans-serif; 
    margin: 0; 
    padding: 0; 
}
.header-container { 
    text-align: center; 
    padding: 15px 20px; 
    background: rgba(0, 0, 0, 0.9); 
    border-bottom: 1px solid #00FF9F; 
    box-shadow: 0 0 10px rgba(161, 0, 255, 0.3); 
}
#ghost-logo { 
    font-size: 60px; 
    display: block; 
    margin: 0 auto; 
    animation: glitch-ghost 1.5s infinite; 
    text-shadow: 0 0 10px #A100FF, 0 0 20px #00FF9F; 
}
h1 { 
    color: #A100FF; 
    font-size: 28px; 
    margin: 5px 0; 
    text-shadow: 0 0 5px #A100FF, 0 0 10px #00FF9F; 
    animation: glitch-text 2s infinite; 
}
p { 
    color: #E0E0E0; 
    font-size: 14px; 
    margin: 5px 0; 
}
.input-container { 
    max-width: 1000px; 
    margin: 20px auto; 
    padding: 20px; 
    background: rgba(28, 37, 38, 0.8); 
    border-radius: 10px; 
    box-shadow: 0 0 15px rgba(0, 255, 159, 0.3); 
}
.textbox { 
    background: #1A1A1A; 
    border: 1px solid #A100FF; 
    color: #E0E0E0; 
    border-radius: 5px; 
    padding: 10px; 
    margin-bottom: 20px; 
}
.genre-buttons { 
    display: flex; 
    justify-content: center; 
    gap: 15px; 
    margin-bottom: 20px; 
}
.genre-btn { 
    background: linear-gradient(45deg, #A100FF, #00FF9F); 
    border: none; 
    color: #0A0A0A; 
    font-weight: bold; 
    padding: 10px 20px; 
    border-radius: 5px; 
    transition: transform 0.3s ease, box-shadow 0.3s ease; 
}
.genre-btn:hover { 
    transform: scale(1.05); 
    box-shadow: 0 0 15px #00FF9F; 
}
.settings-container { 
    max-width: 1000px; 
    margin: 20px auto; 
    padding: 20px; 
    background: rgba(28, 37, 38, 0.8); 
    border-radius: 10px; 
    box-shadow: 0 0 15px rgba(0, 255, 159, 0.3); 
}
.action-buttons { 
    display: flex; 
    justify-content: center; 
    gap: 20px; 
    margin-top: 20px; 
}
button { 
    background: linear-gradient(45deg, #A100FF, #00FF9F); 
    border: none; 
    color: #0A0A0A; 
    font-weight: bold; 
    padding: 12px 24px; 
    border-radius: 5px; 
    transition: transform 0.3s ease, box-shadow 0.3s ease; 
}
button:hover { 
    transform: scale(1.05); 
    box-shadow: 0 0 15px #00FF9F; 
}
.output-container { 
    max-width: 1000px; 
    margin: 20px auto; 
    padding: 20px; 
    background: rgba(28, 37, 38, 0.8); 
    border-radius: 10px; 
    box-shadow: 0 0 15px rgba(0, 255, 159, 0.3); 
    text-align: center; 
}
@keyframes glitch-ghost { 
    0% { transform: translate(0, 0); opacity: 1; }
    20% { transform: translate(-5px, 2px); opacity: 0.8; }
    40% { transform: translate(5px, -2px); opacity: 0.6; }
    60% { transform: translate(-3px, 1px); opacity: 0.9; }
    80% { transform: translate(3px, -1px); opacity: 0.7; }
    100% { transform: translate(0, 0); opacity: 1; }
}
@keyframes glitch-text { 
    0% { transform: translate(0, 0); text-shadow: 0 0 10px #A100FF, 0 0 20px #00FF9F; }
    20% { transform: translate(-2px, 1px); text-shadow: 0 0 15px #00FF9F, 0 0 25px #A100FF; }
    40% { transform: translate(2px, -1px); text-shadow: 0 0 10px #A100FF, 0 0 30px #00FF9F; }
    60% { transform: translate(-1px, 2px); text-shadow: 0 0 15px #00FF9F, 0 0 20px #A100FF; }
    80% { transform: translate(1px, -2px); text-shadow: 0 0 10px #A100FF, 0 0 25px #00FF9F; }
    100% { transform: translate(0, 0); text-shadow: 0 0 10px #A100FF, 0 0 20px #00FF9F; }
}
@font-face { 
    font-family: 'Orbitron'; 
    src: url('https://fonts.gstatic.com/s/orbitron/v29/yMJRMIlzdpvBhQQL_Qq7dy0.woff2') format('woff2'); 
}
"""

# 8) BUILD WITH BLOCKS
with gr.Blocks(css=css) as demo:
    gr.Markdown("""
        <div class="header-container">
            <div id="ghost-logo">👻</div>
            <h1>GhostAI Music Generator</h1>
            <p>Summon the Sound of the Unknown</p>
        </div>
    """)
    
    with gr.Column(elem_classes="input-container"):
        instrumental_prompt = gr.Textbox(
            label="Instrumental Prompt",
            placeholder="Click a genre button below or type your own instrumental prompt",
            lines=4,
            elem_classes="textbox"
        )
        with gr.Row(elem_classes="genre-buttons"):
            rock_btn = gr.Button("Rock", elem_classes="genre-btn")
            techno_btn = gr.Button("Techno", elem_classes="genre-btn")
            jazz_btn = gr.Button("Jazz", elem_classes="genre-btn")
            classical_btn = gr.Button("Classical", elem_classes="genre-btn")
            hiphop_btn = gr.Button("Hip-Hop", elem_classes="genre-btn")
    
    with gr.Column(elem_classes="settings-container"):
        cfg_scale = gr.Slider(
            label="Guidance Scale (CFG)",
            minimum=1.0,
            maximum=10.0,
            value=3.0,
            step=0.1,
            info="Higher values make the instrumental more closely follow the prompt, but may reduce diversity."
        )
        top_k = gr.Slider(
            label="Top-K Sampling",
            minimum=10,
            maximum=500,
            value=300,
            step=10,
            info="Limits sampling to the top k most likely tokens. Higher values increase diversity."
        )
        top_p = gr.Slider(
            label="Top-P Sampling (Nucleus Sampling)",
            minimum=0.0,
            maximum=1.0,
            value=0.95,
            step=0.1,
            info="Keeps tokens with cumulative probability above p. Higher values increase diversity."
        )
        temperature = gr.Slider(
            label="Temperature",
            minimum=0.1,
            maximum=2.0,
            value=1.0,
            step=0.1,
            info="Controls randomness. Higher values make output more diverse but less predictable."
        )
        total_duration = gr.Slider(
            label="Total Duration (seconds)",
            minimum=10,
            maximum=60,
            value=30,
            step=1,
            info="Total duration of the track (10 to 60 seconds)."
        )
        crossfade_duration = gr.Slider(
            label="Crossfade Duration (ms)",
            minimum=100,
            maximum=2000,
            value=500,
            step=100,
            info="Crossfade duration between chunks for smoother transitions."
        )
        with gr.Row(elem_classes="action-buttons"):
            gen_btn = gr.Button("Generate Music")
            clr_btn = gr.Button("Clear Inputs")
    
    with gr.Column(elem_classes="output-container"):
        out_audio = gr.Audio(label="Generated Stereo Instrumental Track", type="filepath")
        status = gr.Textbox(label="Status", interactive=False)
    
    rock_btn.click(set_rock_prompt, inputs=None, outputs=[instrumental_prompt])
    techno_btn.click(set_techno_prompt, inputs=None, outputs=[instrumental_prompt])
    jazz_btn.click(set_jazz_prompt, inputs=None, outputs=[instrumental_prompt])
    classical_btn.click(set_classical_prompt, inputs=None, outputs=[instrumental_prompt])
    hiphop_btn.click(set_hiphop_prompt, inputs=None, outputs=[instrumental_prompt])
    gen_btn.click(
        generate_music,
        inputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, crossfade_duration],
        outputs=[out_audio, status]
    )
    clr_btn.click(
        clear_inputs,
        inputs=None,
        outputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, crossfade_duration]
    )

# 9) TURN OFF OPENAPI/DOCS to avoid the gradio-client schema bug
app = demo.launch(
    server_name="0.0.0.0",
    server_port=9999,
    share=False,
    inbrowser=False,
    show_error=True  # keep this so you still see stack traces in console
)
# Access the underlying FastAPI and disable its docs
try:
    fastapi_app = demo._server.app  # Gradio v4.44+ puts FastAPI here
    fastapi_app.docs_url = None
    fastapi_app.redoc_url = None
    fastapi_app.openapi_url = None
except Exception:
    pass