import gradio as gr
import cv2
import os
import random
import zipfile
from PIL import Image
from datetime import datetime
import io
import hashlib

def procesar_video(video):
    try:
        if isinstance(video, dict):
            original_name = video.get("name", "video")
            video_path = video.get("file", video.get("data"))
        else:
            original_name = os.path.basename(video)
            video_path = video

        allowed_extensions = ('.mp4', '.avi', '.mov', '.mkv')
        if not original_name.lower().endswith(allowed_extensions):
            raise gr.Error("Solo se permiten archivos de video (mp4, avi, mov, mkv)")

        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        temp_dir = f"temp_{datetime.now().strftime('%Y%m%d%H%M%S')}"
        os.makedirs(temp_dir, exist_ok=True)

        # Extracción de fotogramas
        cap = cv2.VideoCapture(video_path)
        frame_count = 0
        frame_paths = []
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(frame_rgb)
            img_path = os.path.join(temp_dir, f"frame_{frame_count:04d}.jpg")
            img.save(img_path)
            frame_paths.append(img_path)
            frame_count += 1
        cap.release()

        if frame_count == 0:
            raise gr.Error("No se pudieron extraer fotogramas del video")

        # Crear collage
        selected_frames = random.sample(frame_paths, min(4, len(frame_paths)))
        images = [Image.open(img) for img in selected_frames]
        margin = 10
        collage = Image.new('RGB', 
            (images[0].width * 2 + margin, images[0].height * 2 + margin), 
            (255,255,255))
        positions = [
            (0, 0),
            (images[0].width + margin, 0),
            (0, images[0].height + margin),
            (images[0].width + margin, images[0].height + margin)
        ]
        for i, img in enumerate(images):
            collage.paste(img, positions[i])
        collage_path = os.path.join(temp_dir, "collage.jpg")
        collage.save(collage_path)

        # Generación del ZIP
        base_name = os.path.splitext(original_name)[0]
        zip_filename = f"{base_name}.zip"
        final_zip_path = os.path.join(temp_dir, zip_filename)
        
        with zipfile.ZipFile(final_zip_path, mode="w") as zipf:
            # Añadir todos los frames
            for img_path in frame_paths:
                zipf.write(img_path, os.path.basename(img_path))
            
            # Añadir TXT de cadena de custodia
            chain_content = (
                f"Nombre del archivo: {original_name}\n"
                f"Fecha de carga y extracción: {timestamp}\n"
                f"Número de fotogramas: {frame_count}\n"
                f"Hash MD5 del video: {hashlib.md5(open(video_path, 'rb').read()).hexdigest()}\n"
            )
            zipf.writestr("cadena_custodia.txt", chain_content)

        return collage_path, final_zip_path, temp_dir

    except Exception as e:
        raise gr.Error(f"Error al procesar el video: {str(e)}")

def limpiar_cache(temp_dir):
    if temp_dir and os.path.exists(temp_dir):
        for file in os.listdir(temp_dir):
            os.remove(os.path.join(temp_dir, file))
        os.rmdir(temp_dir)

with gr.Blocks(title="Extracción de Fotogramas Forenses") as demo:
    gr.Markdown("# 📷 Herramienta de Extracción de Fotogramas Forenses")
    gr.Markdown("**Carga un video para extraer TODOS los fotogramas y generar un collage de muestra.**")
    gr.Markdown("Desarrollado por José R. Leonett para el Grupo de Peritos Forenses Digitales de Guatemala - [www.forensedigital.gt](https://www.forensedigital.gt)")
    
    with gr.Row():
        with gr.Column():
            video_input = gr.Video(label="Subir Video", interactive=True)
            procesar_btn = gr.Button("Procesar Fotogramas", interactive=False)
        with gr.Column():
            gallery_output = gr.Image(label="Collage de Muestra")
            download_btn = gr.Button("DESCARGAR FOTOGRAMAS", interactive=False)
            download_file = gr.File(label="Archivo ZIP generado", visible=True)  # Cambiado a visible
    
    temp_dir_state = gr.State(None)
    zip_path_state = gr.State(None)
    
    def on_video_change(video):
        if video:
            return gr.update(interactive=True), gr.update(interactive=False)
        return gr.update(interactive=False), gr.update(interactive=False)
    
    video_input.change(
        fn=on_video_change,
        inputs=video_input,
        outputs=[procesar_btn, download_btn],
        queue=False
    )
    
    def procesar_y_mostrar(video):
        if temp_dir_state.value:
            limpiar_cache(temp_dir_state.value)
        collage_path, zip_path, temp_dir = procesar_video(video)
        return collage_path, zip_path, temp_dir, zip_path, gr.update(interactive=True), gr.update(interactive=False)
    
    procesar_btn.click(
        fn=procesar_y_mostrar,
        inputs=video_input,
        outputs=[gallery_output, download_file, temp_dir_state, zip_path_state, download_btn, procesar_btn],
    )
    
    def trigger_download(zip_path):
        try:
            with open(zip_path, "rb") as f:
                file_bytes = f.read()
            return os.path.basename(zip_path), file_bytes
        except Exception as e:
            raise gr.Error(f"Error al generar descarga: {str(e)}")
    
    download_btn.click(
        fn=trigger_download,
        inputs=zip_path_state,
        outputs=download_file,
    )

if __name__ == "__main__":
    demo.launch()