import streamlit as st
import tempfile
import os
import subprocess
import imageio_ffmpeg
import time
import gc
import shutil
import yt_dlp
from PIL import Image

# --- 1. CẤU HÌNH & CSS ---
ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()
st.set_page_config(page_title="Subtitles Pro Fixed", layout="wide", page_icon="🎬")

st.markdown("""
<style>
    .main .block-container { padding-top: 1rem; }
    div[data-testid="stImage"], div[data-testid="stVideo"] {
        display: flex; justify-content: center; align-items: center; width: 100%;
    }
    .stSpinner { margin-bottom: 0px; }
    .group-label { font-weight: bold; color: #333; margin-bottom: 5px; display: block; text-decoration: underline;}
</style>
""", unsafe_allow_html=True)

# --- 2. SESSION STATE ---
if 'preview_img_path' not in st.session_state: st.session_state.preview_img_path = None
if 'final_video_path' not in st.session_state: st.session_state.final_video_path = None
if 'sub_configs' not in st.session_state: st.session_state.sub_configs = {}
if 'input_video_path' not in st.session_state: st.session_state.input_video_path = None
if 'custom_font_name' not in st.session_state: st.session_state.custom_font_name = None

# --- 3. HÀM HỖ TRỢ ---

def hex_to_ass_color(hex_color, opacity_percent=100):
    """
    Chuyển đổi màu Hex sang chuẩn ASS (&HAABBGGRR).
    opacity_percent: 
        - 100: Màu đậm đặc (Solid) -> Alpha 00
        - 0: Trong suốt (Invisible) -> Alpha FF
    """
    # Tính giá trị Alpha (đảo ngược: 100% hiện = 00 alpha)
    alpha_val = 255 - int(opacity_percent * 255 / 100)
    alpha_hex = f"{alpha_val:02X}"
    
    hex_color = hex_color.lstrip('#')
    if len(hex_color) != 6: return "&H00FFFFFF"
    
    # Thứ tự ASS là Blue-Green-Red
    r, g, b = hex_color[0:2], hex_color[2:4], hex_color[4:6]
    return f"&H{alpha_hex}{b}{g}{r}"

def get_linux_path(path): return path.replace("\\", "/")

def install_custom_font(uploaded_font):
    font_dir = os.path.join(os.getcwd(), "fonts")
    os.makedirs(font_dir, exist_ok=True)
    font_path = os.path.join(font_dir, uploaded_font.name)
    with open(font_path, "wb") as f:
        f.write(uploaded_font.getvalue())
    font_name = os.path.splitext(uploaded_font.name)[0]
    return font_path, font_dir, font_name

def download_video_from_url(url):
    temp_dir = tempfile.gettempdir()
    ydl_opts = {
        'format': 'bestvideo[height<=720][ext=mp4]+bestaudio[ext=m4a]/best[height<=720][ext=mp4]',
        'outtmpl': os.path.join(temp_dir, 'dl_video_%(id)s.%(ext)s'),
        'quiet': True, 'no_warnings': True, 'socket_timeout': 30,
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=True)
        filename = ydl.prepare_filename(info)
        return filename

def create_static_video(image_file, audio_file, resolution_mode):
    temp_dir = tempfile.gettempdir()
    img_path = os.path.join(temp_dir, "temp_bg.jpg")
    audio_path = os.path.join(temp_dir, "temp_audio.mp3")
    out_path = os.path.join(temp_dir, f"static_vid_{int(time.time())}.mp4")
    
    with open(img_path, "wb") as f: f.write(image_file.getvalue())
    with open(audio_path, "wb") as f: f.write(audio_file.getvalue())

    if resolution_mode == "1080p (Full HD)": vf = "scale=-2:1080" 
    elif resolution_mode == "720p (HD)": vf = "scale=-2:720"
    elif resolution_mode == "Vuông 1:1": vf = "scale=1080:1080:force_original_aspect_ratio=decrease,pad=1080:1080:(ow-iw)/2:(oh-ih)/2"
    elif resolution_mode == "Dọc 9:16": vf = "scale=1080:1920:force_original_aspect_ratio=decrease,pad=1080:1920:(ow-iw)/2:(oh-ih)/2"
    else: vf = "scale=trunc(iw/2)*2:trunc(ih/2)*2"

    cmd = [
        ffmpeg_exe, "-y", "-loop", "1", "-i", img_path, "-i", audio_path,
        "-c:v", "libx264", "-tune", "stillimage", "-c:a", "aac", "-b:a", "192k",
        "-vf", vf, "-pix_fmt", "yuv420p", "-shortest", "-preset", "ultrafast", out_path
    ]
    subprocess.run(cmd, check=True)
    return out_path

def generate_preview_multi(video_path, sub_configs_list, font_dir=None):
    out_img = os.path.join(tempfile.gettempdir(), f"prev_{int(time.time())}.jpg")
    filter_chain = []
    
    for conf in sub_configs_list:
        p = get_linux_path(conf['path'])
        s = conf['style']
        if font_dir:
            filter_chain.append(f"subtitles='{p}':fontsdir='{get_linux_path(font_dir)}':force_style='{s}'")
        else:
            filter_chain.append(f"subtitles='{p}':force_style='{s}'")
    
    full_filter = ",".join(filter_chain)
    cmd = [
        ffmpeg_exe, "-y", "-ss", "00:00:05", "-i", video_path,
        "-vf", full_filter, "-vframes", "1", "-q:v", "5", out_img
    ]
    subprocess.run(cmd, capture_output=True)
    return out_img

# --- 4. GIAO DIỆN CHÍNH ---

st.title("🎬 Subtitles Pro (Fixed Colors)")

col_sidebar, col_main = st.columns([1, 2.5])

with col_sidebar:
    st.header("1. Nguồn Video")
    source_type = st.radio("Chọn nguồn:", ["Upload từ máy", "Link URL", "Tạo từ Ảnh + Nhạc"])
    
    input_ready = False
    
    if source_type == "Upload từ máy":
        v_file = st.file_uploader("Chọn file", type=["mp4", "avi", "mov", "mkv"])
        if v_file:
            t_dir = os.path.join(tempfile.gettempdir(), "uploads")
            os.makedirs(t_dir, exist_ok=True)
            v_path = os.path.join(t_dir, v_file.name)
            with open(v_path, "wb") as f: f.write(v_file.getvalue())
            st.session_state.input_video_path = v_path
            input_ready = True

    elif source_type == "Link URL":
        url = st.text_input("Dán link:")
        if url and st.button("📥 Tải Video"):
            with st.spinner("Đang tải..."):
                try:
                    v_path = download_video_from_url(url)
                    st.session_state.input_video_path = v_path
                    input_ready = True
                    st.success("OK!")
                except Exception as e: st.error(f"Lỗi: {e}")
        elif st.session_state.input_video_path and "dl_video" in st.session_state.input_video_path:
            input_ready = True

    elif source_type == "Tạo từ Ảnh + Nhạc":
        bg_img = st.file_uploader("1. Ảnh", type=["jpg", "png"])
        audio_file = st.file_uploader("2. Nhạc", type=["mp3", "wav"])
        res_choice = st.selectbox("Độ phân giải", ["Giữ nguyên", "1080p (Full HD)", "720p (HD)", "Dọc 9:16", "Vuông 1:1"])
        if bg_img and audio_file and st.button("🎞️ Tạo Video"):
            with st.spinner("Đang tạo..."):
                v_path = create_static_video(bg_img, audio_file, res_choice)
                st.session_state.input_video_path = v_path
                input_ready = True
                st.success("OK!")
        elif st.session_state.input_video_path and "static_vid" in st.session_state.input_video_path:
            input_ready = True

    st.divider()
    st.header("2. Subtitle & Font")
    subtitle_files = st.file_uploader("Upload Sub", type=["srt", "ass"], accept_multiple_files=True)
    custom_font_file = st.file_uploader("Upload Font (.ttf)", type=["ttf", "otf"])
    
    font_dir_local = None
    if custom_font_file:
        f_path, f_dir, f_name = install_custom_font(custom_font_file)
        st.session_state.custom_font_name = f_name
        font_dir_local = f_dir
        st.success(f"Font: {f_name}")

if input_ready and subtitle_files and st.session_state.input_video_path:
    video_path = st.session_state.input_video_path
    t_dir = os.path.join(tempfile.gettempdir(), "subs")
    os.makedirs(t_dir, exist_ok=True)
    saved_subs = []
    for f in subtitle_files:
        path = os.path.join(t_dir, f.name)
        with open(path, "wb") as file: file.write(f.getvalue())
        saved_subs.append({"name": f.name, "path": path})

    with col_main:
        c_settings, c_preview = st.columns([1.3, 1.5], gap="medium")
        
        with c_settings:
            st.subheader("3. Tùy chỉnh Sub")
            tabs = st.tabs([s['name'] for s in saved_subs])
            current_configs = []

            with st.form(key='main_form'):
                for i, tab in enumerate(tabs):
                    with tab:
                        font_opts = ["Arial (Mặc định)", "Times New Roman", "Noto Sans CJK SC (Trung)", "Noto Sans CJK KR (Hàn)", "DejaVu Sans (Nga)", "Monospace"]
                        if st.session_state.custom_font_name: font_opts.insert(0, f"CUSTOM: {st.session_state.custom_font_name}")

                        c1, c2 = st.columns(2)
                        with c1: font_choice = st.selectbox(f"Font", font_opts, key=f"fn_{i}")
                        with c2: f_size = st.number_input(f"Cỡ chữ", 1, 200, 7, key=f"fs_{i}")

                        col_text, col_style = st.columns(2)
                        with col_text:
                            st.markdown('<span class="group-label">Văn bản</span>', unsafe_allow_html=True)
                            text_color = st.color_picker(f"Màu chữ", "#FFFFFF", key=f"tc_{i}")
                            is_bold = st.checkbox("In Đậm", False, key=f"bd_{i}")
                            is_italic = st.checkbox("In Nghiêng", False, key=f"it_{i}")

                        with col_style:
                            st.markdown('<span class="group-label">Nền & Viền</span>', unsafe_allow_html=True)
                            # Màu nền: Mặc định Đen
                            bg_color = st.color_picker(f"Màu Nền (Hộp)", "#000000", key=f"bgc_{i}")
                            # Opacity: Mặc định 0 (Tắt)
                            bg_opacity = st.number_input(f"Độ đậm nền %", 0, 100, 0, key=f"bgo_{i}", help="0=Không nền, 100=Nền đặc")
                            
                            outline_color = st.color_picker(f"Màu Viền", "#000000", key=f"oc_{i}")
                            outline_width = st.number_input(f"Độ dày viền", 0.0, 20.0, 1.0, step=0.1, key=f"ow_{i}")

                        st.write("---")
                        st.markdown('<span class="group-label">Vị trí</span>', unsafe_allow_html=True)
                        cp1, cp2 = st.columns(2)
                        with cp1: v_pos = st.selectbox(f"Vị trí Dọc", ["Dưới cùng", "Trên cùng", "Giữa"], 0, key=f"vp_{i}")
                        with cp2: align = st.selectbox(f"Căn lề Ngang", ["Giữa", "Trái", "Phải"], key=f"al_{i}")
                        
                        cm1, cm2 = st.columns(2)
                        with cm1: marg_v = st.number_input(f"Cách lề Dọc", 0, 1000, 10, key=f"mv_{i}")
                        with cm2: marg_h = st.number_input(f"Cách lề Ngang", 0, 1000, 10, key=f"mh_{i}")

                        # --- LOGIC XỬ LÝ STYLE (ĐÃ SỬA LỖI MÀU) ---
                        final_font = "sans-serif"
                        if "CUSTOM:" in font_choice: final_font = st.session_state.custom_font_name
                        elif "Times" in font_choice: final_font = "serif"
                        elif "Trung" in font_choice: final_font = "Noto Sans CJK SC"
                        elif "Hàn" in font_choice: final_font = "Noto Sans CJK KR"
                        elif "Nga" in font_choice: final_font = "DejaVu Sans"

                        a_code = 2
                        if v_pos == "Dưới cùng": a_code = 1 if align=="Trái" else (3 if align=="Phải" else 2)
                        elif v_pos == "Trên cùng": a_code = 5 if align=="Trái" else (7 if align=="Phải" else 6)
                        else: a_code = 9 if align=="Trái" else (11 if align=="Phải" else 10)

                        # LOGIC QUAN TRỌNG:
                        # Nếu Opacity > 0 -> Dùng BorderStyle=3 (Hộp).
                        # Lúc này BackColour sẽ là màu của Hộp Nền.
                        if bg_opacity > 0:
                            border_style = 3
                            back_colour_ass = hex_to_ass_color(bg_color, bg_opacity)
                        else:
                            # Nếu Opacity = 0 -> Dùng BorderStyle=1 (Chữ thường).
                            border_style = 1
                            # BackColour lúc này là màu bóng đổ (Shadow), ta để trong suốt để ẩn bóng đi
                            back_colour_ass = "&HFF000000"

                        # OutlineColour luôn là màu viền chữ, độ đậm luôn là 100% (Solid)
                        outline_colour_ass = hex_to_ass_color(outline_color, 100)

                        ass_style = (
                            f"FontName={final_font},FontSize={f_size},"
                            f"Bold={'-1' if is_bold else '0'},Italic={'-1' if is_italic else '0'},"
                            f"PrimaryColour={hex_to_ass_color(text_color, 100)},"
                            f"OutlineColour={outline_colour_ass},"
                            f"BackColour={back_colour_ass}," 
                            f"Outline={outline_width},BorderStyle={border_style}," 
                            f"Alignment={a_code},MarginV={marg_v},MarginL={marg_h},MarginR={marg_h},"
                            f"Shadow=0" # Tắt bóng đổ mặc định để tránh nhiễu
                        )
                        current_configs.append({"path": saved_subs[i]['path'], "style": ass_style})

                st.write("")
                submit_btn = st.form_submit_button("📸 CẬP NHẬT PREVIEW", type="primary")

        if submit_btn:
            st.session_state.sub_configs = current_configs
            st.toast("Đang tạo preview...", icon="⏳")
            f_dir = font_dir_local if st.session_state.custom_font_name else None
            new_img = generate_preview_multi(video_path, current_configs, f_dir)
            st.session_state.preview_img_path = new_img
            st.toast("Đã xong!", icon="✅")

        with c_preview:
            st.subheader("4. Xem trước")
            with st.container(height=650, border=True):
                if st.session_state.preview_img_path and os.path.exists(st.session_state.preview_img_path):
                    st.image(st.session_state.preview_img_path, use_container_width=True)
                else:
                    st.info("👈 Bấm Cập nhật để xem")
                    
    st.divider()
    if st.button("🚀 XUẤT VIDEO (FULL)", type="secondary", use_container_width=True):
        if st.session_state.sub_configs:
            gc.collect()
            out_path = os.path.join(tempfile.gettempdir(), f"final_{int(time.time())}.mp4")
            
            with st.status("Đang Render...", expanded=True) as status:
                filter_chain = []
                f_dir = font_dir_local if st.session_state.custom_font_name else None
                for conf in st.session_state.sub_configs:
                    p = get_linux_path(conf['path'])
                    s = conf['style']
                    if f_dir:
                        filter_chain.append(f"subtitles='{p}':fontsdir='{get_linux_path(f_dir)}':force_style='{s}'")
                    else:
                        filter_chain.append(f"subtitles='{p}':force_style='{s}'")
                full_filter = ",".join(filter_chain)

                cmd = [
                    ffmpeg_exe, "-y", "-i", video_path,
                    "-vf", full_filter, "-c:a", "copy", "-c:v", "libx264", 
                    "-preset", "ultrafast", "-crf", "28", out_path
                ]
                
                res = subprocess.run(cmd, capture_output=True)
                if res.returncode == 0:
                    status.update(label="Xong!", state="complete", expanded=False)
                    st.session_state.final_video_path = out_path
                else:
                    status.update(label="Lỗi!", state="error")
                    st.error(f"Render lỗi: {res.stderr}")
        else:
            st.error("Chưa cấu hình sub!")

    if st.session_state.final_video_path and os.path.exists(st.session_state.final_video_path):
        c1, c2 = st.columns([2, 1])
        with c1: st.video(st.session_state.final_video_path)
        with c2:
            with open(st.session_state.final_video_path, "rb") as f:
                st.download_button("⬇️ Tải về", f, "video_final.mp4", "video/mp4", type="primary")
else:
    st.info("👈 Chọn nguồn Video để bắt đầu.")