import os
import requests
import streamlit as st
import streamlit.components.v1 as components
from PIL import Image
import base64
from io import BytesIO

st.set_page_config(page_title="EmoM Playlist Generator", layout="wide", initial_sidebar_state="collapsed")

API_URL = os.getenv("BACKEND_API_URL", "http://emom_inference:8000") + "/analyze"
DEAM_AUDIO_DIR = "/app/dataset/DEAM/DEAM_audio/MEMD_audio"

def get_thumbnail_html(images, max_display=12):
    html_images = ""
    for file in images[:max_display]:
        img = Image.open(file)
        img.thumbnail((100, 100))
        if img.mode != "RGB": 
            img = img.convert("RGB")
        buffered = BytesIO()
        img.save(buffered, format="JPEG")
        b64_str = base64.b64encode(buffered.getvalue()).decode()
        html_images += f'<img src="data:image/jpeg;base64,{b64_str}" style="width: 60px; height: 60px; object-fit: cover; border-radius: 8px; margin-right: 8px; margin-bottom: 8px; border: 1px solid rgba(255, 255, 255, 0.2);">'
    
    if len(images) > max_display:
        html_images += f'<span style="display: inline-block; width: 60px; height: 60px; line-height: 60px; text-align: center; background: rgba(150, 150, 150, 0.2); border-radius: 8px; vertical-align: top; font-size: 14px;">+{len(images) - max_display}</span>'
    return f'<div style="display: flex; flex-wrap: wrap;">{html_images}</div>'

def main():
    if "live_state" not in st.session_state: 
        st.session_state.live_state = "upload"
    if "result_data" not in st.session_state: 
        st.session_state.result_data = None
        
    viewport = st.query_params.get("viewport", "desktop")

    st.markdown("""
    <style>
    [data-testid="stFileUploadDropzone"] { min-height: 250px !important; display: flex; align-items: center; justify-content: center; border-radius: 16px; background-color: rgba(255, 75, 75, 0.03); }
    .spinner-container { display: flex; flex-direction: column; align-items: center; justify-content: center; min-height: 40vh; margin-top: 10vh; }
    .big-spinner { width: 120px; height: 120px; border: 10px solid rgba(255, 75, 75, 0.1); border-top: 10px solid #ff4b4b; border-radius: 50%; animation: spin 1s linear infinite; margin-bottom: 2rem; }
    @keyframes spin { 0% { transform: rotate(0deg); } 100% { transform: rotate(360deg); } }
    #MainMenu {visibility: hidden;} footer {visibility: hidden;}
    </style>
    """, unsafe_allow_html=True)

    if st.session_state.live_state == "upload":
        upload_placeholder = st.empty()
        with upload_placeholder.container():
            st.write("Загрузите изображения для визуально-семантического анализа.")
            if viewport == "mobile": 
                st.markdown("<br>", unsafe_allow_html=True)
            
            uploaded_files = st.file_uploader(
                "Загрузка файлов", 
                type=['png', 'jpg', 'jpeg'], 
                accept_multiple_files=True,
                label_visibility="collapsed" if viewport == "mobile" else "visible"
            )
            
            if uploaded_files:
                st.markdown("<br>", unsafe_allow_html=True)
                if st.button("Выполнить анализ", type="primary", use_container_width=True):
                    st.session_state.uploaded_images = uploaded_files
                    st.session_state.live_state = "processing"
                    upload_placeholder.empty()
                    st.rerun()
                
                st.markdown("<br>", unsafe_allow_html=True)
                st.caption("Выбранные файлы:")
                st.markdown(get_thumbnail_html(uploaded_files), unsafe_allow_html=True)

    elif st.session_state.live_state == "processing":
        components.html("<script>window.parent.scrollTo(0, 0);</script>", height=0, width=0)
        files = st.session_state.get("uploaded_images", [])
        st.markdown('<div class="spinner-container"><div class="big-spinner"></div><h3 style="text-align: center; font-weight: 400;">Обработка данных...</h3></div>', unsafe_allow_html=True)
        
        try:
            upload_data = [('files', (f.name, f.getvalue(), f.type)) for f in files]
            response = requests.post(API_URL, files=upload_data, timeout=300) 
            
            if response.status_code == 200:
                st.session_state.result_data = response.json()
                st.session_state.live_state = "result"
                st.rerun()
            else:
                st.error(f"Ошибка сервера: {response.status_code}")
                if st.button("Назад"):
                    st.session_state.live_state = "upload"
                    st.rerun()
        except Exception as e:
            st.error(f"Ошибка соединения: {str(e)}")
            if st.button("Назад"):
                st.session_state.live_state = "upload"
                st.rerun()
                
    elif st.session_state.live_state == "result":
        components.html("<script>window.parent.scrollTo(0, 0);</script>", height=0, width=0)
        data = st.session_state.result_data
        
        st.header(f"Сгенерированный плейлист (обработано файлов: {data['images_processed']})")
        
        for row in data.get("tracks", []):
            with st.container(border=True):
                song_id = int(row['song_id'])
                score = row['final_score']
                
                audio_path = f"{DEAM_AUDIO_DIR}/{song_id}.mp3"
                if not os.path.exists(audio_path): 
                    audio_path = audio_path.replace('.mp3', '.wav')
                
                if viewport == "desktop":
                    c1, c2 = st.columns([1, 3])
                    with c1:
                        st.write(f"**Track ID:** {song_id}")
                        st.caption(f"Score: {score:.4f}")
                    with c2:
                        if os.path.exists(audio_path): 
                            st.audio(audio_path)
                        else: 
                            st.caption("Аудиофайл не найден")
                else:
                    st.write(f"**Track ID:** {song_id} (Score: {score:.4f})")
                    if os.path.exists(audio_path): 
                        st.audio(audio_path)
                    else: 
                        st.caption("Аудиофайл не найден")

        st.markdown("<br>", unsafe_allow_html=True)
        
        with st.expander("Отладочная информация (Метрики)"):
            st.subheader("Координаты V/A")
            c_v, c_a = st.columns(2)
            c_v.metric("Valence", f"{data['target_v']:.2f}")
            c_a.metric("Arousal", f"{data['target_a']:.2f}")
            
            st.markdown("---")
            st.subheader("Акустические признаки (LLM)")
            
            feature_titles = {
                "energy": "RMS Energy",
                "flux": "Spectral Flux",
                "centroid": "Spectral Centroid",
                "pitch": "F0 (Pitch)",
                "hnr": "HNR",
                "zcr": "ZCR"
            }
            
            # Развернутые описания
            feature_helps = {
                "energy": "Среднеквадратичная амплитуда (громкость). Бывает высокой в плотных, интенсивных композициях, отражает общую акустическую энергию сцены.",
                "flux": "Спектральный поток. Измеряет резкость изменений в спектре. Высок при четком, агрессивном ритме и частой смене нот.",
                "centroid": "Спектральный центроид («яркость» звука). Высокие значения указывают на преобладание высоких частот (звонкие инструменты, открытые пространства).",
                "pitch": "Основная частота звука. Высокий pitch характерен для позитивных, легких или, напротив, напряженных мелодий.",
                "hnr": "Отношение гармоник к шуму. Высокий HNR — чистая мелодия и вокал. Низкий HNR — присутствие дисторшна, шумов или перкуссии.",
                "zcr": "Частота пересечения нуля. Отражает шумовую составляющую сигнала. Высок в треках с выраженными ударными (hi-hats) или атмосферным шумом."
            }
            
            llm_profile = data.get("llm_profile")
            if llm_profile and isinstance(llm_profile, dict) and len(llm_profile) > 0:
                cols_per_row = 2 if viewport == "mobile" else 3
                llm_items = list(llm_profile.items())
                
                for i in range(0, len(llm_items), cols_per_row):
                    cols = st.columns(cols_per_row)
                    for j in range(cols_per_row):
                        if i + j < len(llm_items):
                            k, v = llm_items[i + j]
                            label = feature_titles.get(k, k)
                            tooltip = feature_helps.get(k, "")
                            cols[j].metric(label, f"{v:.2f}", help=tooltip)
            else:
                st.caption("Акустический профиль недоступен. Применен fallback-алгоритм.")
                            
            st.markdown("---")
            st.write("**Извлеченные теги (BLIP-2):**")
            st.write(", ".join([str(c).capitalize() for c in data.get("semantics", [])]))

        st.markdown("<br>", unsafe_allow_html=True)
        
        if st.button("Новый запрос", use_container_width=True):
            st.session_state.live_state = "upload"
            st.session_state.result_data = None
            st.session_state.pop("uploaded_images", None)
            st.rerun()

if __name__ == "__main__":
    main()