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"
    )