import streamlit as st
import streamlit.components.v1 as components
import numpy as np
from PIL import Image
import base64
from io import BytesIO
from music_engine.llm_bridge import LLMAcousticBridge

# Вспомогательная функция для крохотного предпросмотра
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 render_live_tab(matcher, image_processor):
    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")

    # ==========================================
    # CSS ИНЪЕКЦИИ
    # ==========================================
    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); } }
    </style>
    """, unsafe_allow_html=True)

    # ==========================================
    # ЭКРАН 1: ЗАГРУЗКА
    # ==========================================
    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:
                # 1. КНОПКА СРАЗУ ПОСЛЕ ЗАГРУЗКИ (Не нужно скроллить вниз)
                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()
                
                # 2. МИНИАТЮРЫ ПОД КНОПКОЙ
                st.markdown("<br>", unsafe_allow_html=True)
                st.caption("Выбранные кадры:")
                # Генерируем компактный блок миниатюр
                st.markdown(get_thumbnail_html(uploaded_files), unsafe_allow_html=True)

    # ==========================================
    # ЭКРАН 2: АНАЛИЗ (СПИННЕР)
    # ==========================================
    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></div>', unsafe_allow_html=True)
        status_text = st.empty()
        
        images = []
        all_objects = []
        all_v, all_a = [], []
        
        for i, file in enumerate(files):
            status_text.markdown(f"<h3 style='text-align: center; font-weight: 400;'>Анализ кадра {i + 1} из {len(files)}...</h3>", unsafe_allow_html=True)
            
            img = Image.open(file)
            images.append(img)
            
            embedding = image_processor.extract_embedding(img)
            v, a = matcher.predict_va(embedding)
            all_v.append(v)
            all_a.append(a)
            
            caption = image_processor.describe_scene(img)
            all_objects.append(caption)
            
        target_v, target_a = np.mean(all_v), np.mean(all_a)
        
        status_text.markdown("<h3 style='text-align: center; font-weight: 400;'>Трансляция семантики в аудиопрофиль...</h3>", unsafe_allow_html=True)
        llm = LLMAcousticBridge()
        llm_profile = llm.get_acoustic_profile(target_v, target_a, list(set(all_objects)))
        
        status_text.markdown("<h3 style='text-align: center; font-weight: 400;'>Поиск идеальных композиций...</h3>", unsafe_allow_html=True)
        playlist = matcher.find_nearest_tracks(target_v, target_a, llm_profile=llm_profile, top_k=15)
        
        st.session_state.result_data = {
            "target_v": target_v,
            "target_a": target_a,
            "llm_profile": llm_profile,
            "playlist": playlist,
            "semantics": list(set(all_objects))
        }
        st.session_state.live_state = "result"
        st.rerun()
        
    # ==========================================
    # ЭКРАН 3: РЕЗУЛЬТАТЫ
    # ==========================================
    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("Рекомендованный плейлист")
        
        for _, row in data["playlist"].iterrows():
            with st.container(border=True):
                if viewport == "desktop":
                    c1, c2 = st.columns([1, 3])
                    with c1:
                        st.write(f"**Track:** {int(row['song_id'])}")
                        st.caption(f"Score: {row['final_score']:.2f}")
                    with c2:
                        audio_path = matcher.get_audio_path(row['song_id'])
                        if audio_path:
                            st.audio(str(audio_path))
                        else:
                            st.warning("Файл не найден")
                else:
                    st.write(f"**Track:** {int(row['song_id'])} (Score: {row['final_score']:.2f})")
                    audio_path = matcher.get_audio_path(row['song_id'])
                    if audio_path:
                        st.audio(str(audio_path))
                    else:
                        st.warning("Файл не найден")

        st.markdown("<br>", unsafe_allow_html=True)
        
        with st.expander("Технические параметры анализа"):
            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.write("**Акустические таргеты (LLM):**")
            if data["llm_profile"]:
                cols_per_row = 2 if viewport == "mobile" else 3
                llm_items = list(data["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]
                            cols[j].metric(k, f"{v:.2f}")
                            
            st.markdown("---")
            st.write("**Обнаруженная семантика:**")
            st.write(", ".join([str(c).capitalize() for c in data["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()