import gradio as gr
from PIL import Image
import cv2
import numpy as np
import os
import zipfile
import hashlib
import logging
from pathlib import Path
from PIL import ExifTags
import plotly.graph_objects as go

# Configuración de logs
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Directorio de evidencia
evidence_dir = "/home/user/app"
os.makedirs(evidence_dir, exist_ok=True)
logger.info(f"Directorio de evidencia creado en: {evidence_dir}")

def obtener_metadatos(imagen):
    try:
        exif_data = imagen.getexif()
        if not exif_data:
            return {}
        metadata = {}
        for tag_id, value in exif_data.items():
            try:
                tag = ExifTags.TAGS.get(tag_id, tag_id)
                metadata[tag] = value
            except Exception as e:
                logger.debug(f"Error al procesar etiqueta EXIF: {str(e)}")
        return metadata
    except Exception as e:
        logger.error(f"Error al obtener metadatos: {str(e)}")
        return {}

def obtener_coordenadas(exif_data):
    if not exif_data or "GPSInfo" not in exif_data:
        return None
    try:
        gps_info = exif_data["GPSInfo"]
        if not gps_info:
            return None
        def gps_to_degrees(coord):
            d, m, s = coord
            return d + (m / 60.0) + (s / 3600.0)
        lat = gps_info.get(2)
        lon = gps_info.get(4)
        lat_ref = gps_info.get(1)
        lon_ref = gps_info.get(3)
        if lat and lon and lat_ref and lon_ref:
            lat_deg = gps_to_degrees(lat)
            lon_deg = gps_to_degrees(lon)
            if lat_ref == "S":
                lat_deg = -lat_deg
            if lon_ref == "W":
                lon_deg = -lon_deg
            return lat_deg, lon_deg
    except Exception as e:
        logger.error(f"Error al procesar coordenadas GPS: {str(e)}")
    return None

def calcular_hash(imagen):
    return hashlib.sha3_256(imagen.tobytes()).hexdigest()

def analizar_manipulacion(imagen, metadatos):
    manipulada = False
    razones = []
    if imagen.mode == "P":
        razones.append("La imagen tiene marcas de agua o es una imagen indexada.")
        manipulada = True
    if "Software" in metadatos:
        razones.append(f"La imagen fue editada con: {metadatos['Software']}")
        manipulada = True
    return manipulada, razones

def calcular_porcentaje_ela(mask):
    if mask is None or mask.size == 0:
        return 0.0
    total_pixeles = mask.size
    pixeles_detectados = np.count_nonzero(mask)
    porcentaje = (pixeles_detectados / total_pixeles) * 100
    return porcentaje

def estimar_probabilidad_manipulacion(porcentaje_ela):
    if porcentaje_ela < 0.5:
        return "Muy baja (< 0.5%) - Imagen probablemente auténtica."
    elif porcentaje_ela < 2.0:
        return "Baja (0.5% - 2%) - Posible compresión o ruido, pero no manipulación evidente."
    elif porcentaje_ela < 5.0:
        return "Moderada (2% - 5%) - Posible edición o retoque localizado."
    elif porcentaje_ela < 15.0:
        return "Alta (5% - 15%) - Alta probabilidad de manipulación o zonas borradas."
    else:
        return "Muy alta (> 15%) - Manipulación extensa o generación por IA detectada."


def realizar_ela(imagen):
    try:
        img_np = np.array(imagen.convert("RGB"))
        img_cv = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)

        # Fuerza la compresión para detectar artefactos
        quality = 75  
        temp_path = "/tmp/temp_image.jpg"
        cv2.imwrite(temp_path, img_cv, [cv2.IMWRITE_JPEG_QUALITY, quality])
        img_comprimida = cv2.imread(temp_path)
        if img_comprimida is None:
            raise ValueError("No se pudo leer la imagen comprimida.")

        # Calcula diferencias
        diferencia = cv2.absdiff(img_cv.astype(np.float32), img_comprimida.astype(np.float32))
        scaled_diff = np.clip(diferencia * 15, 0, 255).astype(np.uint8)  # Escalado fuerte

        # Convertir a escala de grises y aplicar umbral más bajo
        gray_diff = cv2.cvtColor(scaled_diff, cv2.COLOR_BGR2GRAY)
        _, mask = cv2.threshold(gray_diff, 5, 255, cv2.THRESH_BINARY)

        # Resalta con color donde hay diferencias
        ela_color = cv2.applyColorMap(gray_diff, cv2.COLORMAP_TURBO)

        # Combinar con la original en escala de grises
        img_gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
        img_gray = cv2.cvtColor(img_gray, cv2.COLOR_GRAY2BGR)
        result = np.where(mask[..., None] > 0, ela_color, img_gray)

        os.remove(temp_path)
        return result, mask

    except Exception as e:
        logger.error(f"Error en ELA: {str(e)}")
        error_img = np.zeros((300, 600, 3), dtype=np.uint8) + 30
        cv2.putText(error_img, "ERROR AL PROCESAR ELA", (50, 150), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
        return error_img, None




def crear_graficos_estadisticas(porcentaje_ela, dimensiones, metadatos):
    """Crear gráficos estadísticos para el análisis ELA"""
    
    # Gráfico 1: Porcentaje de manipulación (Barra)
    fig_bar = go.Figure()
    fig_bar.add_trace(go.Bar(
        x=['Áreas Detectadas'],
        y=[porcentaje_ela],
        marker_color=['#4CAF50' if porcentaje_ela > 2 else '#36A2EB'],
        text=[f'{porcentaje_ela:.3f}%'],
        textposition='auto',
    ))
    fig_bar.update_layout(
        title=dict(text="<b>Porcentaje de Áreas Manipuladas</b>", x=0.5),
        yaxis_title="Porcentaje (%)",
        height=250,
        showlegend=False,
        margin=dict(l=20, r=20, t=40, b=20)
    )
    
    # Gráfico 2: Probabilidad de manipulación (Pie)
    niveles = ['Muy baja', 'Baja', 'Moderada', 'Alta', 'Muy alta']
    color_map = ['#00CC96', '#36A2EB', '#FFCE56', '#FF9F40', '#FF6384']
    
    prob_index = 0
    if porcentaje_ela >= 15:
        prob_index = 4
    elif porcentaje_ela >= 5:
        prob_index = 3
    elif porcentaje_ela >= 2:
        prob_index = 2
    elif porcentaje_ela >= 0.5:
        prob_index = 1
        
    fig_pie = go.Figure()
    fig_pie.add_trace(go.Pie(
        labels=[niveles[prob_index]],
        values=[100],
        hole=.6,
        marker_colors=[color_map[prob_index]],
        textinfo='label'
    ))
    fig_pie.update_layout(
        title=dict(text="<b>Nivel de Probabilidad</b>", x=0.5),
        height=250,
        showlegend=False,
        margin=dict(l=20, r=20, t=40, b=20)
    )
    
    # Gráfico 3: Distribución de píxeles (Box plot)
    fig_box = go.Figure()
    fig_box.add_trace(go.Box(
        y=[porcentaje_ela],
        name="Distribución ELA",
        marker_color='#FF6B35',
        boxpoints='all'
    ))
    fig_box.update_layout(
        title=dict(text="<b>Distribución de Anomalías</b>", x=0.5),
        yaxis_title="Porcentaje (%)",
        height=250,
        showlegend=False,
        margin=dict(l=20, r=20, t=40, b=20)
    )
    
    # Gráfico 4: Tendencias temporales (Scatter)
    fig_scatter = go.Figure()
    fig_scatter.add_trace(go.Scatter(
        x=['Compresión', 'Ruido', 'Manipulación'],
        y=[max(0.1, porcentaje_ela-1), max(0.1, porcentaje_ela-0.5), max(0.1, porcentaje_ela)],
        mode='lines+markers',
        line=dict(color='#FF6B35', width=3),
        marker=dict(size=10)
    ))
    fig_scatter.update_layout(
        title=dict(text="<b>Análisis de Componentes</b>", x=0.5),
        yaxis_title="Intensidad",
        height=250,
        margin=dict(l=20, r=20, t=40, b=20)
    )
    
    return fig_bar, fig_pie, fig_box, fig_scatter

def procesar_imagen(archivo_imagen):
    if not archivo_imagen:
        return [None, "❌ **ERROR: Por favor, cargue una imagen antes de analizar.**", None, ""] + [None] * 4

    if not os.path.exists(archivo_imagen):
        return [None, "❌ **ERROR: El archivo de imagen no existe o es inválido.**", None, ""] + [None] * 4

    try:
        img = Image.open(archivo_imagen)
        logger.info(f"Imagen cargada: {archivo_imagen}")

        nombre_original = os.path.basename(archivo_imagen)
        nombre = os.path.splitext(nombre_original)[0]

        original_path = os.path.join(evidence_dir, f"{nombre}_original.jpg")
        ela_path = os.path.join(evidence_dir, f"{nombre}_ela.jpg")
        text_path = os.path.join(evidence_dir, f"{nombre}_analisis.txt")
        zip_path = os.path.join(evidence_dir, f"{nombre}_errorELA.zip")

        img.save(original_path, "JPEG")
        ela_result, mask = realizar_ela(img)

        if mask is None:
            raise Exception("No se pudo generar el análisis ELA.")

        cv2.imwrite(ela_path, ela_result)

        porcentaje_ela = calcular_porcentaje_ela(mask)
        probabilidad = estimar_probabilidad_manipulacion(porcentaje_ela)

        # Crear gráficos estadísticos
        fig_bar, fig_pie, fig_box, fig_scatter = crear_graficos_estadisticas(
            porcentaje_ela, img.size, obtener_metadatos(img)
        )

        file_size_bytes = os.path.getsize(archivo_imagen)
        if file_size_bytes < 1024 * 1024:
            file_size = f"{file_size_bytes / 1024:.2f} KB"
        else:
            file_size = f"{file_size_bytes / (1024 * 1024):.2f} MB"

        info_basica = f"Nombre del archivo: {nombre_original}\r\n"
        info_basica += f"Tamaño del archivo: {file_size}\r\n"
        info_basica += f"Dimensiones: {img.size[0]} x {img.size[1]} píxeles\r\n"
        info_basica += f"Formato: {img.format}\r\n"
        info_basica += f"Modo: {img.mode}\r\n\r\n"

        metadatos = obtener_metadatos(img)
        info_metadatos = "ANÁLISIS FORENSE DE LOS METADATOS:\r\n\r\n"

        google_maps_url = None

        if metadatos:
            for tag, value in metadatos.items():
                if tag == "DateTime":
                    info_metadatos += f"- Fecha y hora de captura: {value}\r\n"
                elif tag == "Make":
                    info_metadatos += f"- Fabricante de cámara: {value}\r\n"
                elif tag == "Model":
                    info_metadatos += f"- Modelo de cámara: {value}\r\n"
                elif tag == "Software":
                    info_metadatos += f"- Software de edición: {value}\r\n"
                elif tag == "GPSInfo":
                    coords = obtener_coordenadas(metadatos)
                    if coords:
                        lat, lon = coords
                        info_metadatos += f"- Coordenadas GPS: {lat:.6f}, {lon:.6f}\r\n"
                        google_maps_url = f"https://www.google.com/maps?q={lat},{lon}"
                        info_metadatos += f"- Enlace a Google Maps: {google_maps_url}\r\n"
                    else:
                        info_metadatos += "- **Coordenadas GPS:** No se encontraron coordenadas válidas\r\n"
                else:
                    info_metadatos += f"- **{tag}:** {value}\r\n"
        else:
            info_metadatos += "- **Metadatos EXIF:** No se encontraron metadatos EXIF\r\n"

        sha3_hash = calcular_hash(img)
        info_metadatos += f"\r\n SHA3-256: {sha3_hash}\r\n\r\n"

        manipulada, razones = analizar_manipulacion(img, metadatos)
        info_manipulacion = "ANÁLISIS DE MANIPULACIÓN:\r\n\r\n"
        info_manipulacion += f"- Porcentaje de áreas detectadas: {porcentaje_ela:.3f}%\r\n"
        info_manipulacion += f"- Estimación forense: {probabilidad}\r\n\r\n"

        analysis_text = info_basica + info_metadatos + info_manipulacion

        with open(text_path, "w", encoding="utf-8", newline="\r\n") as f:
            f.write(analysis_text)

        with zipfile.ZipFile(zip_path, "w") as zipf:
            zipf.write(original_path, os.path.basename(original_path))
            zipf.write(ela_path, os.path.basename(ela_path))
            zipf.write(text_path, os.path.basename(text_path))

        logger.info(f"Análisis completado. Archivo ZIP: {zip_path}")

        ela_rgb = cv2.cvtColor(ela_result, cv2.COLOR_BGR2RGB)

        return [zip_path, analysis_text, ela_rgb, google_maps_url or "", fig_bar, fig_pie, fig_box, fig_scatter]

    except Exception as e:
        logger.error(f"Error en procesamiento: {str(e)}")
        mensaje_error = f"❌ **ERROR CRÍTICO:** {str(e)}"
        error_img = np.zeros((300, 600, 3), dtype=np.uint8)
        cv2.putText(error_img, "ERROR INTERNO", (80, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
        return [None, mensaje_error, error_img, ""] + [None] * 4

# 🎨 Tema moderno
theme = gr.themes.Soft(primary_hue="blue", secondary_hue="slate", neutral_hue="stone")

# 🖥️ Interfaz Gradio
with gr.Blocks(title="Análisis Forense de Imágenes con ELA", theme=theme, css="""
    .orange-download-btn {
        background: linear-gradient(45deg, #FF6B35, #FF8E53) !important;
        color: white !important;
        border: none !important;
        font-weight: bold !important;
    }
    .orange-download-btn:hover {
        background: linear-gradient(45deg, #E55A2B, #FF7B42) !important;
        transform: scale(1.02) !important;
        box-shadow: 0 4px 12px rgba(255, 107, 53, 0.3) !important;
    }
    .equal-height {
        height: 400px !important;
    }
""") as demo:
    gr.Markdown("""
    # 📸 Análisis Forense de Imágenes (Error Level Analysis - ELA)
    **Programa de computación forense para analizar imágenes en busca de evidencia de manipulación o edición.**
    """)
    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():
            input_image = gr.Image(
                label="Subir imagen (JPG/PNG)",
                type="filepath",
                height=400,
                sources=["upload"],
                elem_classes=["equal-height"]
            )
            process_btn = gr.Button("Analizar imagen", variant="primary")

            download_zip = gr.DownloadButton(
                label="⬇️ Descargar resultados (ZIP)",
                variant="secondary",
                visible=False,
                interactive=True,
                elem_classes=["orange-download-btn"]
            )

            # Estadísticas en 2x2 grid
            with gr.Row(visible=False) as stats_row:
                with gr.Column():
                    stats_bar = gr.Plot(label="Porcentaje de Manipulación", show_label=True)
                    stats_pie = gr.Plot(label="Nivel de Probabilidad", show_label=True)
                with gr.Column():
                    stats_box = gr.Plot(label="Distribución de Anomalías", show_label=True)
                    stats_scatter = gr.Plot(label="Análisis de Componentes", show_label=True)

        with gr.Column():
            with gr.Accordion(open=True):
                ela_image = gr.Image(
                    label="ÁREAS TURQUESA = manipulaciones o borrados",
                    type="numpy",
                    height=400,
                    show_label=True,
                    elem_classes=["equal-height"]
                )

            with gr.Accordion("Informe Detallado", open=True):
                analysis_text = gr.Textbox(label="📝 Resultados del análisis forense", lines=15, max_lines=25)
                google_maps_btn = gr.Button("📍 Ver ubicación en Google Maps", visible=False)
                google_maps_url_state = gr.State("")

    def reset_on_upload():
        return (
            gr.update(value=None),
            gr.update(value=None),
            gr.update(visible=False),
            gr.update(visible=False),
            "",
            gr.update(visible=False),
            None, None, None, None
        )

    process_btn.click(
        fn=procesar_imagen,
        inputs=input_image,
        outputs=[download_zip, analysis_text, ela_image, google_maps_url_state, stats_bar, stats_pie, stats_box, stats_scatter],
        api_name="analyze_image"
    ).then(
        fn=lambda url: gr.update(visible=bool(url)),
        inputs=google_maps_url_state,
        outputs=google_maps_btn
    ).then(
        fn=lambda: gr.update(visible=True),
        inputs=None,
        outputs=download_zip
    ).then(
        fn=lambda: gr.update(visible=True),
        inputs=None,
        outputs=stats_row
    )

    google_maps_btn.click(
        fn=lambda url: url,
        inputs=google_maps_url_state,
        outputs=None,
        js="(url) => { window.open(url, '_blank'); }"
    )

    input_image.upload(
        fn=reset_on_upload,
        inputs=None,
        outputs=[analysis_text, ela_image, download_zip, google_maps_btn, google_maps_url_state, stats_row, stats_bar, stats_pie, stats_box, stats_scatter]
    )

# ▶️ Ejecución
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        inbrowser=True,
        favicon_path="https://www.forensedigital.gt/wp-content/uploads/2019/07/cropped-40fb84a6-c75a-4c38-bfa0-c0a9777430cd-1.jpg"
    )