import streamlit as st
import subprocess
import os
import random
import tempfile
import shutil
import time

from arquivos import salvar_no_dataset, CATEGORIES

st.set_page_config(page_title="TikTok Video Generator", layout="centered")
st.title("🎥 TikTok Video Generator - PRO")

st.markdown("Envie seus vídeos e gere conteúdo com efeitos, zoom, texto, música e filtros!")

# Uploads dos vídeos de cortes (principal)
st.markdown("<div style='background-color:#F0F0FF;padding:10px;border-radius:8px'>", unsafe_allow_html=True)
st.subheader("🎬 Envie os vídeos principais (cortes)")
cortes = st.file_uploader("Vídeos de cortes", type=["mp4"], accept_multiple_files=True)
st.markdown("</div>", unsafe_allow_html=True)

# Upload de múltiplos vídeos ou imagens de fundo
st.markdown("<div style='background-color:#FFF0F0;padding:10px;border-radius:8px;margin-top:15px;'>", unsafe_allow_html=True)
st.subheader("🌄 Envie os vídeos ou imagens de fundo (opcional)")
videos_fundo = st.file_uploader("Vídeos ou Imagens de Fundo", type=["mp4", "mov", "jpg", "jpeg", "png"], accept_multiple_files=True)
st.markdown("</div>", unsafe_allow_html=True)

# Uploads adicionais (múltiplos tutoriais agora)
videos_tutorial = st.file_uploader("📎 Tutoriais (opcional)", type="mp4", accept_multiple_files=True)
musica = st.file_uploader("🎵 Música (opcional, MP3)", type=["mp3"])

# Configurações gerais
num_videos_finais = st.number_input("Quantos vídeos gerar?", 1, 10, 1)
duracao_final = st.number_input("Duração final (s)", 10, 300, 30)
duracao_corte = st.number_input("Duração de cada corte (s)", 1, 30, 5)
zoom = st.slider("Zoom no vídeo principal", 1.0, 2.0, 1.0, 0.1)
blur_strength = st.slider("Blur no fundo", 1, 50, 10)
velocidade_cortes = st.slider("Velocidade dos cortes", 0.5, 2.0, 1.0, 0.1)
velocidade_final = st.slider("Velocidade final", 0.5, 2.0, 1.0, 0.1)
crf_value = st.slider("Qualidade CRF", 18, 30, 23)

# === NOVA OPÇÃO SALVAR NO GERENCIADOR/DATASET ===
st.markdown("### Gerenciamento do vídeo final")
salvar_no_gerenciador_checkbox = st.checkbox("Salvar vídeo no Gerenciador de Arquivos (Dataset)")
if salvar_no_gerenciador_checkbox:
    categoria = st.selectbox("Selecione a categoria:", CATEGORIES)
else:
    categoria = None
# BOTÃO PRINCIPAL
if st.button("Gerar Vídeo(s)"):
    if not cortes:
        st.error("❌ Envie os vídeos de cortes.")
    else:
        with st.spinner("🎬 Processando..."):
            progresso = st.progress(0)
            temp_dir = tempfile.mkdtemp()
            try:
                # Preparar cortes
                cortes_names = []
                for idx, corte in enumerate(cortes):
                    path_in = os.path.join(temp_dir, f"corte_in_{idx}.mp4")
                    path_out = os.path.join(temp_dir, f"corte_padrao_{idx}.mp4")
                    with open(path_in, "wb") as f:
                        f.write(corte.read())
                    subprocess.run([
                        "ffmpeg", "-i", path_in,
                        "-vf", "scale=1280:720:force_original_aspect_ratio=decrease,scale=trunc(iw/2)*2:trunc(ih/2)*2",
                        "-c:v", "libx264", "-preset", "ultrafast", "-crf", "30", path_out
                    ], check=True, stderr=subprocess.PIPE)
                    cortes_names.append(path_out)

                for n in range(num_videos_finais):
                    progresso.progress(5 + n * 5)
                    random.shuffle(cortes_names)
                    tempo_total = 0
                    cortes_prontos = []

                    # Escolher fundo aleatório
                    if videos_fundo:
                        fundo_aleatorio = random.choice(videos_fundo)
                        fundo_ext = os.path.splitext(fundo_aleatorio.name)[-1].lower()
                        fundo_path = os.path.join(temp_dir, f"fundo_{n}{fundo_ext}")

                        with open(fundo_path, "wb") as f:
                            f.write(fundo_aleatorio.read())

                        fundo_convertido = os.path.join(temp_dir, f"fundo_convertido_{n}.mp4")

                        if fundo_ext in [".mp4", ".mov"]:
                            subprocess.run([
                                "ffmpeg", "-i", fundo_path, "-vf",
                                "scale='if(gt(iw/ih,720/1280),max(720,iw),-2)':'if(gt(iw/ih,720/1280),-2,max(1280,ih))',"
                                "scale=720:1280:force_original_aspect_ratio=increase,crop=720:1280",
                                "-t", str(duracao_final), "-preset", "ultrafast", "-crf", "25",
                                fundo_convertido
                            ], check=True, stderr=subprocess.PIPE)
                        else:
                            subprocess.run([
                                "ffmpeg", "-loop", "1", "-i", fundo_path,
                                "-c:v", "libx264", "-t", str(duracao_final),
                                "-pix_fmt", "yuv420p",
                                "-vf", "scale=720:1280:force_original_aspect_ratio=increase,crop=720:1280",
                                "-preset", "ultrafast", "-crf", "25",
                                fundo_convertido
                            ], check=True, stderr=subprocess.PIPE)
                    else:
                        fundo_convertido = os.path.join(temp_dir, f"fundo_preto_{n}.mp4")
                        subprocess.run([
                            "ffmpeg", "-f", "lavfi", "-i", "color=black:s=720x1280:d=600",
                            "-t", str(duracao_final),
                            "-c:v", "libx264", "-preset", "ultrafast", "-crf", "25",
                            fundo_convertido
                        ], check=True, stderr=subprocess.PIPE)
                    # Concatenar cortes
                    lista = os.path.join(temp_dir, f"lista_{n}.txt")
                    with open(lista, "w") as f:
                        for c in cortes_names:
                            f.write(f"file '{c}'\n")

                    video_raw = os.path.join(temp_dir, f"video_raw_{n}.mp4")
                    subprocess.run([
                        "ffmpeg", "-f", "concat", "-safe", "0", "-i", lista,
                        "-vf", "scale=trunc(iw/2)*2:trunc(ih/2)*2",
                        "-c:v", "libx264", "-preset", "ultrafast", "-crf", "30", video_raw
                    ], check=True, stderr=subprocess.PIPE)

                    # Para simplificar: vamos usar o video_raw como vídeo final neste exemplo
                    final_name = f"video_final_{n}_{int(time.time())}.mp4"
                    shutil.copy(video_raw, final_name)

                    # === SALVAR NO DATASET OU DOWNLOAD ===
                    if salvar_no_gerenciador_checkbox and categoria:
                        salvar_no_dataset(final_name, categoria)
                        st.success(f"✅ Vídeo {n+1} salvo no Dataset na categoria '{categoria}'.")
                        st.markdown(
                            f"[🔗 Acessar no Dataset](https://huggingface.co/datasets/pcdoido2/video-gerados/resolve/main/{categoria}/{final_name})"
                        )
                    else:
                        st.video(final_name)
                        with open(final_name, "rb") as f:
                            st.download_button(
                                f"📥 Baixar vídeo {n+1}", f, file_name=final_name
                            )

                progresso.progress(100)
                st.success("✅ Todos os vídeos foram gerados com sucesso!")

            except subprocess.CalledProcessError as e:
                st.error(f"❌ Erro ao gerar vídeo:\n\n{e.stderr.decode(errors='ignore')}")
            finally:
                shutil.rmtree(temp_dir)