import time
import threading
import random
import gradio as gr
import matplotlib.pyplot as plt
from obspy.clients.seedlink.easyseedlink import create_client
from obspy import Stream

# Shared buffers
stream = Stream()
buffer = []
alert_threshold = 5000
use_dummy = True
selected_station = {"network": "IU", "station": "ANMO", "channel": "00BHZ"}  # default

# ObsPy callback
def handle_trace(trace):
    global buffer, use_dummy
    use_dummy = False
    stream.append(trace)
    stream.merge()
    tr = stream[-1]
    buffer.extend(tr.data.tolist())

def dummy_data():
    global buffer
    while True:
        if use_dummy:
            data = [int(random.gauss(0, 500)) for _ in range(200)]
            if random.random() < 0.01:
                data = [x * 20 for x in data]
            buffer.extend(data)
        time.sleep(1)

def start_obspy():
    global selected_station
    try:
        client = create_client("rtserve.iris.washington.edu", on_data=handle_trace)
        client.select_stream(
            selected_station["network"],
            selected_station["station"],
            selected_station["channel"]
        )
        threading.Thread(target=client.run, daemon=True).start()
        print(f"✅ ObsPy started for {selected_station}")
    except Exception as e:
        print("⚠️ ObsPy failed, using dummy mode:", e)

def stream_waveform():
    global buffer
    fig, ax = plt.subplots()
    if buffer:
        batch = buffer[-200:]
        buffer.clear()
        alert = max(abs(x) for x in batch) > alert_threshold
        ax.plot(batch, color="red" if alert else "blue")
        ax.set_ylim(-10000, 10000)
        return fig, ("⚠️ ALERT DETECTED!" if alert else "✅ Normal")
    else:
        ax.text(0.5, 0.5, "Waiting for data...", ha="center")
        ax.set_axis_off()
        return fig, "Waiting for data..."

def change_station(station_choice):
    global selected_station, buffer, stream, use_dummy
    buffer = []
    stream = Stream()
    use_dummy = True
    net, sta, cha = station_choice.split(".")
    selected_station = {"network": net, "station": sta, "channel": cha}
    threading.Thread(target=start_obspy, daemon=True).start()
    return f"📡 Switched to station: {station_choice}"

# Start background threads
threading.Thread(target=dummy_data, daemon=True).start()
threading.Thread(target=start_obspy, daemon=True).start()

STATIONS = [
    "IU.ANMO.00BHZ",
    "IU.TUC.00BHZ",
    "IU.COLA.00BHZ",
    "IU.KONO.00BHZ",
    "IU.KEV.00BHZ",
    "IU.DWPF.00BHZ",
]

with gr.Blocks() as demo:
    gr.Markdown("## 🌍 Live Seismic Waveform (IRIS Stations)")

    station_dropdown = gr.Dropdown(choices=STATIONS, value="IU.ANMO.00BHZ", label="Select Station")
    switch_status = gr.Label(value="📡 Connected to IU.ANMO.00BHZ")

    chart = gr.Plot()
    alert_box = gr.Label(label="Status")

    station_dropdown.change(fn=change_station, inputs=station_dropdown, outputs=switch_status)

    # ✅ Timer replaces every=
    timer = gr.Timer(1.0)
    timer.tick(stream_waveform, None, [chart, alert_box])

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)