import gradio as gr
import obspy
from obspy import Stream, Trace
from obspy.clients.seedlink.easyseedlink import create_client
from seisbench.models import PhaseNet
from groq import Groq
import threading
import numpy as np
import collections
import pandas as pd
import os
import time
from datetime import datetime
import requests
import logging

logging.getLogger("seisbench").setLevel(logging.ERROR)

GROQ_API_KEY = os.environ.get("api_groq")
TELEGRAM_BOT_TOKEN = os.environ.get("TELEGRAM_BOT_TOKEN")
TELEGRAM_CHAT_ID = "-5515143628"

SERVIDOR_CHILE_CSN = "eew.csn.uchile.cl" 
SERVIDOR_ASIA_IRIS = "rtserve.iris.washington.edu"

ESTACIONES_CHILE = [
    # --- NORTE GRANDE (Arica, Tarapacá, Antofagasta) ---
    {"net": "C1", "sta": "GO01", "loc": "Iquique (Tarapacá - Costa) *Latencia ~6h"},  # 
    {"net": "CX", "sta": "PB01", "loc": "Pica (Tarapacá - Interior)"},               # 
    {"net": "CX", "sta": "PB02", "loc": "Pozo Almonte (Tarapacá)"},                  # 
    {"net": "C1", "sta": "LMEL", "loc": "María Elena (Antofagasta - Interior)"},      # 
    {"net": "CX", "sta": "PATCX","loc": "Pampa Alta (Antofagasta) *Latencia ~3h"},    # 
    {"net": "CX", "sta": "PB10", "loc": "Antofagasta (Norte Interior)"},              # 

    # --- NORTE CHICO & CENTRAL (Atacama, Coquimbo, Valparaíso, RM) ---
    {"net": "C1", "sta": "GO02", "loc": "Copiapó (Atacama)"},                         # 
    {"net": "C1", "sta": "GO03", "loc": "Vallenar / Huasco (Atacama)"},               # 
    {"net": "IU", "sta": "LCO",  "loc": "Las Campanas (Observatorio - Coquimbo)"},    # 
    {"net": "C1", "sta": "GO06", "loc": "La Serena / Coquimbo (Costa)"},              # 
    {"net": "C1", "sta": "MT01", "loc": "Santiago / Farellones (Metropolitana)"},     # 
    {"net": "C1", "sta": "MT02", "loc": "Santiago / San José de Maipo (RM)"},         # 
    {"net": "C1", "sta": "ROC1", "loc": "Rancagua / El Teniente (O'Higgins)"},        # 

    # --- ZONA SUR & AUSTRAL (Biobío, La Araucanía, Los Lagos, Magallanes) ---
    {"net": "C1", "sta": "CO01", "loc": "Concepción / Talcahuano (Biobío)"},          # 
    {"net": "C1", "sta": "CO03", "loc": "Chillán / Ñuble (Interior)"},                # 
    {"net": "C1", "sta": "TA01", "loc": "Temuco / Araucanía"},                        # 
    {"net": "C1", "sta": "TA02", "loc": "Valdivia / Los Ríos"},                       # 
    {"net": "C1", "sta": "TA03", "loc": "Puerto Montt / Los Lagos"},                  # 
    {"net": "G",  "sta": "COYC", "loc": "Coyhaique (Aysén)"}                          # 
]

REDES_ASIA = [
    {"net": "IU", "sta": "MAJO", "pais": "Japón", "loc": "Matsushiro"},
    {"net": "II", "sta": "ERM",  "pais": "Japón", "loc": "Erimo (Sanriku)"},
    {"net": "IU", "sta": "PAB",  "pais": "Filipinas", "loc": "San Pablo (Luzón)"},
    {"net": "II", "sta": "DAV",  "pais": "Filipinas", "loc": "Dávao (Mindanao)"},
    {"net": "GE", "sta": "PSI",  "pais": "Indonesia", "loc": "Prapat (Sumatra)"},
    {"net": "GE", "sta": "JCJI", "pais": "Indonesia", "loc": "Jeti (Java)"}
]

MUESTRAS_VENTANA = 3000
buffers_globales = {}

print("[SISTEMA] Cargando e inicializando PhaseNet...")
model = PhaseNet.from_pretrained("original")
model.eval()
print("[SISTEMA] PhaseNet configurado y listo para producción.")

def analizar_paquete_global(trace):
    net = trace.stats.network
    estacion = trace.stats.station
    channel = trace.stats.channel
    componente = channel[-1] if channel else None

    if componente not in ['Z', 'N', 'E']:
        return

    if estacion not in buffers_globales:
        match_cl = next((e for e in ESTACIONES_CHILE if e["sta"] == estacion), None)
        if match_cl:
            pais = "Chile"
            loc = match_cl["loc"]
        else:
            match_as = next((e for e in REDES_ASIA if e["sta"] == estacion), None)
            pais = match_as["pais"] if match_as else "Internacional"
            loc = match_as["loc"] if match_as else "Ubicación Remota"

        buffers_globales[estacion] = {
            'Z': collections.deque(maxlen=MUESTRAS_VENTANA),
            'N': collections.deque(maxlen=MUESTRAS_VENTANA),
            'E': collections.deque(maxlen=MUESTRAS_VENTANA),
            'net': net,
            'channel': channel[:-1],
            'sampling_rate': trace.stats.sampling_rate,
            'pais': pais,
            'loc': loc
        }

    buffers_globales[estacion][componente].extend(trace.data)

    if (len(buffers_globales[estacion]['Z']) == MUESTRAS_VENTANA and 
        len(buffers_globales[estacion]['N']) == MUESTRAS_VENTANA and 
        len(buffers_globales[estacion]['E']) == MUESTRAS_VENTANA):
        
        evaluar_ruptura(estacion)

def evaluar_ruptura(estacion):
    meta = buffers_globales[estacion]
    sr = meta['sampling_rate']
    
    st = Stream()
    amplitudes_maximas = []
    
    for comp in ['Z', 'N', 'E']:
        data_array = np.array(meta[comp], dtype=np.float32)
        
        # --- FILTRO 1: ANTI-GLITCH ELECTRÓNICO ---
        # Si hay un pico de energía plano absurdo (típico error de sensor muerto), descartamos.
        amp_pico = np.max(np.abs(data_array))
        if amp_pico > 1e7:  # Umbral de tolerancia eléctrica estándar para cuentas digitales masivas
            return
            
        data_array = data_array - np.mean(data_array)
        amplitudes_maximas.append(amp_pico)
        
        header = {
            'network': meta['net'],
            'station': estacion,
            'channel': f"{meta['channel']}{comp}",
            'sampling_rate': sr,
            'starttime': obspy.UTCDateTime()
        }
        st.append(Trace(data=data_array, header=header))

    try:
        # --- FILTRO 2: BANDPASS FILTER (1.0 Hz - 15.0 Hz) ---
        # Remueve ruidos ambientales mecánicos y frecuencias parásitas antes de PhaseNet
        st.detrend("linear")
        st.taper(max_percentage=0.05, type="cosine")
        st.filter("bandpass", fmin=1.0, fmax=15.0, corners=4, zerophase=True)
        
        # Normalización estricta por desviación estándar post-filtrado
        for trace in st:
            std = np.std(trace.data)
            if std > 0: 
                trace.data /= std

        # Pasamos la traza limpia por el modelo de IA sismológica
        annotations = model.annotate(st)
        prob_P = np.max(annotations[0].data)
        prob_S = np.max(annotations[1].data)
        
        # --- MODIFICACIÓN: Ajuste de estrictez ---
        # Bajamos P a 0.70 y S a 0.60. Mantiene precisión exigiendo ambas, pero es menos estricto con la dispersión real de la señal.
        if prob_P > 0.70 and prob_S > 0.60:
            idx_P = np.argmax(annotations[0].data)
            idx_S = np.argmax(annotations[1].data)
            
            if idx_S > idx_P:
                delta_t = (idx_S - idx_P) / float(sr)
                
                # --- VALIDACIÓN FÍSICA SÍSMICA ---
                # Un tiempo P-S menor a 0.5s o mayor a 120s en redes regionales suele ser ruido o anomalía.
                if delta_t < 0.5 or delta_t > 120.0:
                    return
                    
                pais = meta['pais']
                ubicacion = meta['loc']
                
                distancia_km = delta_t * 8.2
                amplitud_maxima_total = np.max(amplitudes_maximas)
                
                # Cálculo calibrado de magnitud local
                if amplitud_maxima_total > 0 and distancia_km > 0:
                    magnitud_estimada = np.log10(amplitud_maxima_total) + 1.6 * np.log10(distancia_km) - 0.15
                    magnitud_estimada = round(max(1.0, min(9.5, magnitud_estimada)), 1)
                else:
                    magnitud_estimada = 0.0
                
                disparar_alerta_ia_global(estacion, ubicacion, pais, delta_t, prob_P, prob_S, magnitud_estimada)
                
                for c in ['Z', 'N', 'E']: meta[c].clear()
    except:
        pass

import asyncio
import aiohttp
import time

def enviar_a_telegram_background(url, payload):
    """
    Punto de entrada compatible con hilos que inicializa y ejecuta 
    el bucle asíncrono dedicado para saltarse el estrangulamiento de CPU.
    """
    try:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        loop.run_until_complete(despachar_hacia_gateway_async(url, payload))
        loop.close()
    except Exception as e:
        print(f"⚠️ [SISTEMA] Error crítico al inicializar bucle asíncrono: {e}")

async def despachar_hacia_gateway_async(url, payload):
    """Despacha la alerta usando sockets asíncronos no bloqueantes"""
    intentos_maximos = 3
    timeout_estricto = aiohttp.ClientTimeout(total=25)
    
    await asyncio.sleep(0.5)

    async with aiohttp.ClientSession(trust_env=False, timeout=timeout_estricto) as session:
        for intento in range(1, intentos_maximos + 1):
            try:
                print(f"🔄 [ASYNC GATEWAY] Intentando envío (Intento {intento}/{intentos_maximos})...", flush=True)
                
                headers = {
                    "Connection": "close",
                    "Content-Type": "application/json"
                }
                
                async with session.post(url, json=payload, headers=headers) as response:
                    status = response.status
                    text_response = await response.text()
                    
                    print(f"📥 [ASYNC GATEWAY] Status Code recibido: {status}")
                    
                    if status == 200:
                        print(f"✅ [TELEGRAM] ¡Alerta sísmica enviada con éxito mediante Cloudflare Workers!")
                        return True
                    elif status == 429:
                        print("⏳ [ASYNC GATEWAY] Rate-limit activo en la API. Esperando reintento...")
                        await asyncio.sleep(4)
                    else:
                        print(f"❌ [ASYNC GATEWAY] Error devuelto por endpoint: {text_response}")
                        return False
                        
            except asyncio.TimeoutError:
                print(f"⚠️ [ASYNC GATEWAY] Timeout en intento {intento}. CPU saturada por SeisBench, reintentando...")
                await asyncio.sleep(2)
            except Exception as e:
                print(f"⚠️ [ASYNC GATEWAY] Error de conexión en socket asíncrono: {e}")
                await asyncio.sleep(1)
                
    print("❌ [TELEGRAM] Envío cancelado. No se pudo liberar el socket tras 3 intentos.")
    return False
    
def disparar_alerta_ia_global(estacion, ubicacion, pais, delta_t, p_prob, s_prob, magnitud):
    distancia_epicentro_km = round(delta_t * 8.2, 1)
    
    str_magnitud = f"{magnitud} Ml" if magnitud > 0 else "En cálculo preliminar"
    reporte_final = f"⚠️ <b>[ALERTA SÍSMICA AUTOMÁTICA]</b><br><br>• <b>País:</b> {pais}<br>• <b>Estación:</b> {estacion} ({ubicacion})<br>• <b>Tiempo P-S:</b> {delta_t}s<br>• <b>Magnitud Estimada:</b> {str_magnitud}<br>• <b>Epicentro estimado:</b> a ~{distancia_epicentro_km} km del sensor."

    if GROQ_API_KEY:
        client_groq = Groq(api_key=GROQ_API_KEY)
        
        prompt = f"""[DETECCIÓN PRELIMINAR AUTOMÁTICA INSTRUMENTAL - MONITOREO SEISNET]
        El sistema automático ha filtrado y verificado una posible señal sísmica en la corteza:
        - Región / País: {pais}
        - Ubicación técnica: {ubicacion} (Sensor: {estacion})
        - Tiempo transcurrido entre Onda P y Onda S: {delta_t} segundos.
        - Radio estimado al epicentro: {distancia_epicentro_km} km.
        - Certeza de Fase: Onda P ({round(p_prob*100,1)}%), Onda S ({round(s_prob*100,1)}%).
        - Magnitud Local Estimada: {str_magnitud}.
        
        Instrucciones de formato de salida OBLIGATORIAS:
        1. Encabeza OBLIGATORIAMENTE el mensaje con el título: "⚠️ DETECCIÓN PRELIMINAR INSTRUMENTAL AUTOMÁTICA (EN PROCESO DE VALIDACIÓN)" en mayúsculas destacadas con emojis formales.
        2. Redacta el informe en ESPAÑOL de manera clara, objetiva y breve para lectura móvil.
        3. NO uses asteriscos para negrita. Si quieres usar negrita utiliza la etiqueta HTML <b>texto</b> y para cursiva <i>texto</i>.
        4. Explica explícitamente que es un cálculo de software automático en tiempo real procesado por el equipo de SeisNet y que NO reemplaza bajo ningún motivo el reporte oficial ni la opinión experta del Centro Sismológico Nacional de Chile (CSN).
        5. Llama a mantener la calma e infórmale a las zonas cercanas el tiempo estimado de viaje que tienen las ondas secundarias basado en el radio calculado.
        6. Muestra la Magnitud Local Estimada de forma clara en su sección."""

        try:
            completion = client_groq.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[{"role": "user", "content": prompt}]
            )
            reporte_final = completion.choices[0].message.content
            reporte_final = reporte_final.replace("\n", "<br>")
        except Exception as e:
            print(f"Error en Groq API: {e}. Se enviará reporte técnico base.")

    print(f"\n🚨 [ALERTA GENERADA] - Despachando a cola de Telegram... 🚨\n")

    if TELEGRAM_BOT_TOKEN and TELEGRAM_CHAT_ID:
        url_gateway = f"https://seisnet.bleondubos.workers.dev/bot{TELEGRAM_BOT_TOKEN}/sendMessage"
        
        texto_telegram = reporte_final.replace("<br>", "\n").replace("<br/>", "\n")
        texto_telegram = texto_telegram.replace("<tbd>", "TBD").replace("<TBD>", "TBD")
        texto_telegram = texto_telegram.replace("</tbd>", "")
        
        payload = {
            "chat_id": str(TELEGRAM_CHAT_ID).strip(), 
            "text": texto_telegram,
        }
        
        threading.Thread(target=enviar_a_telegram_background, args=(url_gateway, payload), daemon=True).start()
    else:
        print(f"⚠️ [TELEGRAM] Envío omitido: Faltan credenciales (Token: {'OK' if TELEGRAM_BOT_TOKEN else 'FALTA'}, ChatID: {'OK' if TELEGRAM_CHAT_ID else 'FALTA'})")
        
def conectar_chile_csn():
    while True:
        try:
            print(f"[CHILE CSN] Abriendo canal explícito con {SERVIDOR_CHILE_CSN}...")
            client = create_client(SERVIDOR_CHILE_CSN, on_data=analizar_paquete_global)
            
            for est in ESTACIONES_CHILE:
                try:
                    client.select_stream(est["net"], est["sta"], "?H?")
                except:
                    continue
            client.run()
        except Exception as e:
            print(f"[RECONEXIÓN CHILE] Reintentando en 15 segundos... Info: {e}")
            time.sleep(15)

def conectar_asia_iris():
    while True:
        try:
            print(f"[ASIA IRIS] Abriendo canal con {SERVIDOR_ASIA_IRIS}...")
            client = create_client(SERVIDOR_ASIA_IRIS, on_data=analizar_paquete_global)
            
            for est in REDES_ASIA:
                try:
                    client.select_stream(est["net"], est["sta"], "BH?")
                except:
                    continue
            client.run()
        except Exception as e:
            print(f"[RECONEXIÓN ASIA] Reintentando en 15 segundos... Info: {e}")
            time.sleep(15)

threading.Thread(target=conectar_chile_csn, daemon=True).start()
threading.Thread(target=conectar_asia_iris, daemon=True).start()

with gr.Blocks() as demo:
    gr.Markdown("# 🌍 Consola de Vigilancia Sísmica Total - Cobertura País Chile")
    gr.Markdown("Escucha activa multipunto en el servidor del CSN sin bloqueos de comodines.")
    
    with gr.Row():
        gr.DataFrame(pd.DataFrame(ESTACIONES_CHILE), label="Malla de Sensores Chilenos en Monitoreo Directo")

demo.launch()