import streamlit as st
import librosa
import soundfile as sf
import pyrubberband as pyrb
from pydub import AudioSegment, effects
import os
import tempfile
from PIL import Image

# ===== CONFIG =====
st.set_page_config(page_title="AutoTune Karaoke Enhancer", page_icon="🎤", layout="centered")

# ===== HEADER =====
st.markdown("""
    <h1 style='text-align:center; color:#FF4B4B;'>🎤 AutoTune Karaoke Enhancer</h1>
    <p style='text-align:center; font-size:18px;'>
        แปลงเสียงร้องคาราโอเกะของคุณให้เพราะขึ้นแบบมืออาชีพ
    </p>
""", unsafe_allow_html=True)

# ===== FUNCTIONS =====
def separate_vocals(input_file):
    os.system(f"spleeter separate -p spleeter:2stems -o output {input_file}")
    base = os.path.splitext(os.path.basename(input_file))[0]
    vocal_path = f"output/('my_vocals.wav')"
    accompaniment_path = f"output/accompaniment.wav"
    return vocal_path, accompaniment_path

def autotune(vocal_file, target_key="C"):
    y, sr = librosa.load('my_vocals.wav')
    pitches, magnitudes = librosa.piptrack(y=y, sr=sr)
    pitch_series = [pitches[magnitudes[:, i].argmax(), i] for i in range(pitches.shape[1]) if pitches[magnitudes[:, i].argmax(), i] > 0]
    semitones = librosa.hz_to_midi(pitch_series)
    shift = - (round(semitones.mean()) - librosa.note_to_midi(target_key))
    y_shifted = pyrb.pitch_shift(y, sr, shift)
    output_file = "autotuned_vocals.wav"
    sf.write(output_file, y_shifted, sr)
    return output_file

def enhance_vocals(vocal_file):
    vocal = AudioSegment.from_file(vocal_file)
    vocal = vocal.low_pass_filter(8000).high_pass_filter(120)
    temp_file = "temp_vocal.wav"
    vocal.export(temp_file, format="wav")
    os.system(f"sox {temp_file} enhanced_vocal.wav reverb 50 50 100")
    return "enhanced_vocal.wav"

def mix_tracks(vocal_file, music_file, output_file="final_mix.wav"):
    vocal = AudioSegment.from_file(vocal_file)
    music = AudioSegment.from_file(music_file)
    mixed = effects.normalize(music.overlay(vocal))
    mixed.export(output_file, format="wav")
    return output_file

# ===== UPLOAD UI =====
uploaded_file = st.file_uploader("📂 เลือกไฟล์เสียง (MP3/WAV)", type=["mp3", "wav"])
target_key = st.selectbox("🎯 เลือกคีย์เป้าหมาย", ["C", "D", "E", "F", "G", "A", "B"])

# ===== PROCESSING =====
if uploaded_file is not None:
    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as temp_file:
        temp_file.write(uploaded_file.read())
        temp_path = temp_file.name
    
    if st.button("🚀 เริ่มประมวลผล"):
        progress = st.progress(0)
        
        st.info("🔍 กำลังแยกเสียงร้อง...")
        vocals, music = separate_vocals(temp_path)
        progress.progress(25)
        
        st.info("🎯 กำลังปรับคีย์...")
        tuned_vocals = autotune(vocals, target_key=target_key)
        progress.progress(50)
        
        st.info("✨ กำลังปรับโทนเสียง...")
        enhanced_vocals = enhance_vocals(tuned_vocals)
        progress.progress(75)
        
        st.info("🎚 กำลังมิกซ์เสียง...")
        final_song = mix_tracks(enhanced_vocals, music)
        progress.progress(100)
        
        st.success("✅ เสร็จแล้ว! ดาวน์โหลดไฟล์ด้านล่าง")
        st.audio(final_song, format="audio/wav")
        with open(final_song, "rb") as f:
            st.download_button("📥 ดาวน์โหลดไฟล์", f, file_name="karaoke_enhanced.wav")