import gradio as gr
import numpy as np
import plotly.graph_objects as go
import pandas as pd   # ← Explicit early import fixes Plotly's pandas validator crash

# ====================== DRONE SIMULATOR CLASS ======================
class DroneSimulator:
    def __init__(self):
        self.reset()

    def reset(self):
        self.x = self.y = self.z = 0.0
        self.vx = self.vy = self.vz = 0.0
        self.roll = self.pitch = self.yaw = 0.0
        self.battery = 100.0
        self.armed = False
        self.mode = "DISARMED"
        self.waypoints = []
        self.current_wp = 0
        self.sim_time = 0.0
        self.history = {'t': [], 'x': [], 'y': [], 'z': [], 'battery': []}
        self.logs = ["[0.0s] Ground station online – High-reliability autonomy demo ready."]
        self.link_quality = 95.0
        self.gps_ok = True
        self.home = (0.0, 0.0, 0.0)
        self.failsafe_active = False

    def add_log(self, msg):
        ts = f"[{self.sim_time:.1f}s]"
        self.logs.append(f"{ts} {msg}")
        if len(self.logs) > 50: self.logs = self.logs[-50:]

    def arm(self):
        if not self.armed:
            self.armed = True
            self.mode = "STABILIZE"
            self.add_log("DRONE ARMED – Stabilize mode active")
            return True
        return False

    def disarm(self):
        self.armed = False
        self.mode = "DISARMED"
        self.vx = self.vy = self.vz = 0.0
        self.add_log("DRONE DISARMED")

    def load_waypoints(self, text):
        if not text or not text.strip():
            self.waypoints = []
            self.add_log("Waypoints cleared")
            return True
        try:
            lines = [line.strip() for line in text.strip().splitlines() if line.strip()]
            self.waypoints = [tuple(float(p.strip()) for p in line.split(',')) for line in lines]
            self.current_wp = 0
            self.add_log(f"✅ Loaded {len(self.waypoints)} waypoints")
            return True
        except:
            self.add_log("❌ Invalid format – use x,y,z one per line")
            return False

    def update(self, dt=0.5):
        if not self.armed:
            self.z = max(0.0, self.z - dt)
            self.sim_time += dt
            self.record_history()
            return

        self.sim_time += dt
        self.battery = max(0.0, self.battery - (0.35 + 0.08 * (abs(self.vx) + abs(self.vy) + abs(self.vz))) * dt)
        self.link_quality = max(10.0, self.link_quality - np.random.normal(0, 0.08))

        self.check_failsafes()

        target_z = self.z
        tvx = tvy = 0.0

        if self.mode == "TAKEOFF":
            target_z = 15.0
            if abs(self.z - 15.0) < 1.5:
                self.mode = "AUTO" if self.waypoints else "STABILIZE"
                self.add_log("Takeoff complete → autonomous mode")
        elif self.mode == "LAND":
            target_z = 0.0
            if self.z < 0.5:
                self.disarm()
                self.add_log("Safe landing confirmed")
        elif self.mode == "RTL":
            tx, ty, tz = self.home
            dx = tx - self.x; dy = ty - self.y
            dist = (dx**2 + dy**2)**0.5
            if dist < 3.0 and abs(self.z - tz) < 2.0:
                self.mode = "LAND"
            else:
                speed = 7.0
                if dist > 0:
                    tvx = (dx / dist) * speed
                    tvy = (dy / dist) * speed
                target_z = tz
        elif self.mode == "AUTO" and self.waypoints and self.current_wp < len(self.waypoints):
            tx, ty, tz = self.waypoints[self.current_wp]
            dx = tx - self.x; dy = ty - self.y; dz = tz - self.z
            dist = (dx**2 + dy**2 + dz**2)**0.5
            if dist < 3.0:
                self.current_wp += 1
                self.add_log(f"Waypoint {self.current_wp-1} reached")
                if self.current_wp >= len(self.waypoints):
                    self.mode = "RTL"
                    self.add_log("Mission complete – RTL initiated")
            else:
                speed = 8.0
                if dist > 0:
                    tvx = (dx / dist) * speed
                    tvy = (dy / dist) * speed
                target_z = tz

        self.vx += (tvx - self.vx) * 1.2 * dt
        self.vy += (tvy - self.vy) * 1.2 * dt
        self.vz += ((target_z - self.z) * 2.0 - self.vz) * 0.9 * dt

        self.x += self.vx * dt
        self.y += self.vy * dt
        self.z = max(0.0, self.z + self.vz * dt)

        self.roll = np.clip(self.vy * 3, -30, 30)
        self.pitch = np.clip(-self.vx * 3, -30, 30)

        self.record_history()

        if int(self.sim_time) % 2 == 0:
            self.add_log(f"TX Telemetry → Ground (alt={self.z:.1f}m, batt={self.battery:.0f}%)")

    def check_failsafes(self):
        if self.battery < 25 and self.mode not in ["RTL", "LAND"]:
            self.mode = "RTL"; self.failsafe_active = True
            self.add_log("*** CRITICAL FAILSAFE: Low battery → RTL ACTIVATED ***")
        if self.link_quality < 25 and self.mode not in ["RTL", "LAND"]:
            self.mode = "LAND"
            self.add_log("*** CRITICAL FAILSAFE: Link loss → Emergency land ***")
        if not self.gps_ok and self.mode == "AUTO":
            self.mode = "RTL"
            self.add_log("*** FAILSAFE: GPS failure → RTL ***")

    def record_history(self):
        self.history['t'].append(self.sim_time)
        self.history['x'].append(self.x)
        self.history['y'].append(self.y)
        self.history['z'].append(self.z)
        self.history['battery'].append(self.battery)
        if len(self.history['t']) > 300:
            for k in self.history: self.history[k] = self.history[k][-250:]

    def trigger_failure(self, ftype):
        if ftype == "low_battery": self.battery = 18.0
        elif ftype == "gps_loss": self.gps_ok = False
        elif ftype == "comm_loss": self.link_quality = 12.0
        elif ftype == "imu":
            self.roll = np.random.normal(0, 18)
            self.pitch = np.random.normal(0, 18)
        self.add_log(f"SIMULATED FAILURE: {ftype.replace('_', ' ').upper()}")
        self.check_failsafes()

# ====================== PLOTS ======================
def gauge(value, title, maxv=100, suffix="", color="#00ffcc"):
    fig = go.Figure(go.Indicator(
        mode="gauge+number",
        value=value,
        title={"text": title, "font": {"size": 18, "color": "#f8f8f8"}},
        gauge={"axis": {"range": [0, maxv], "tickcolor": "#f8f8f8"},
               "bar": {"color": color},
               "steps": [{"range": [0, maxv*0.3], "color": "#ff3333"},
                         {"range": [maxv*0.3, maxv*0.7], "color": "#ffaa00"},
                         {"range": [maxv*0.7, maxv], "color": "#00ff88"}]},
        number={"font": {"color": "#f8f8f8", "size": 28}}
    ))
    fig.update_layout(height=210, margin=dict(l=10,r=10,t=50,b=10), paper_bgcolor="#0d0d0d")
    return fig

def path_plot(d):
    fig = go.Figure()
    if d.history['x']:
        fig.add_trace(go.Scatter(x=d.history['x'], y=d.history['y'], mode='lines',
                                 line=dict(color='#00ffcc', width=3), name='Path'))
    fig.add_trace(go.Scatter(x=[d.x], y=[d.y], mode='markers+text',
                             marker=dict(size=20, color='#ff00aa', symbol='star'),
                             text=["🚁"], textposition="top center", name='Live'))

    pad = 8 if d.history['x'] else 25
    if d.history['x']:
        minx = min(d.history['x']) - pad
        maxx = max(d.history['x']) + pad
        miny = min(d.history['y']) - pad
        maxy = max(d.history['y']) + pad
    else:
        minx = maxx = miny = maxy = 0
        minx -= pad; maxx += pad; miny -= pad; maxy += pad

    fig.update_layout(
        title="Top-Down Flight Map (X-Y)",
        xaxis_title="East (m)", yaxis_title="North (m)",
        height=460, plot_bgcolor="#111111", paper_bgcolor="#0d0d0d",
        font_color="#f8f8f8",
        xaxis=dict(range=[minx, maxx], gridcolor="#333333", zerolinecolor="#555555"),
        yaxis=dict(range=[miny, maxy], gridcolor="#333333", zerolinecolor="#555555"),
    )
    return fig

def alt_plot(d):
    fig = go.Figure(go.Scatter(x=d.history['t'], y=d.history['z'], mode='lines', line=dict(color='#ffaa00', width=3)))
    fig.update_layout(title="Altitude History", xaxis_title="Time (s)", yaxis_title="Altitude (m)",
                      height=260, plot_bgcolor="#111111", paper_bgcolor="#0d0d0d", font_color="#f8f8f8")
    return fig

def batt_plot(d):
    fig = go.Figure(go.Scatter(x=d.history['t'], y=d.history['battery'], mode='lines', line=dict(color='#00ccff', width=3)))
    fig.update_layout(title="Battery History", xaxis_title="Time (s)", yaxis_title="Battery %",
                      height=260, plot_bgcolor="#111111", paper_bgcolor="#0d0d0d", font_color="#f8f8f8")
    return fig

# ====================== CUSTOM CSS (2px radius + premium dark) ======================
custom_css = """
    .gradio-container { background-color: #0d0d0d !important; color: #f8f8f8 !important; }
    .gr-button, .gr-textbox textarea, .gr-plot, .gr-textbox, .gr-tab { 
        border-radius: 2px !important; 
        border: 1px solid #333333 !important; 
        transition: all 0.2s;
    }
    .gr-button.primary { background-color: #ff5e00 !important; color: #f8f8f8 !important; }
    .gr-button.secondary { background-color: #333333 !important; color: #f8f8f8 !important; }
    .gr-button:hover { filter: brightness(1.15); }
    .gr-markdown h1, .gr-markdown h2, .gr-markdown p { color: #f8f8f8 !important; }
"""

# ====================== DARK THEME ======================
dark_theme = gr.themes.Default().set(
    body_background_fill="#0d0d0d",
    body_text_color="#f8f8f8",
    panel_background_fill="#1a1a1a",
    block_background_fill="#1a1a1a",
    block_border_color="#333333",
    button_primary_background_fill="#ff5e00",
    button_primary_background_fill_hover="#ff7a33",
    button_secondary_background_fill="#333333",
    button_secondary_background_fill_hover="#444444",
    input_background_fill="#222222",
    input_border_color="#444444"
)

# ====================== GRADIO UI ======================
with gr.Blocks(title="AeroGuard – High-Reliability Autonomous Drone Simulator") as demo:
    gr.Markdown("# 🚁 AeroGuard – High-Reliability Autonomous Drone Simulator\n"
                "**Real-time telemetry • Autonomous waypoint navigation • Layered fail-safes • Ground-station comms**")

    sim_state = gr.State(DroneSimulator())

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Ground Station Controls")
            with gr.Row():
                arm_btn = gr.Button("ARM", variant="primary", size="large")
                disarm_btn = gr.Button("DISARM", size="large")
            with gr.Row():
                takeoff_btn = gr.Button("TAKEOFF", variant="primary")
                land_btn = gr.Button("LAND")
                rtl_btn = gr.Button("RTL")
            emergency_btn = gr.Button("🚨 EMERGENCY STOP", variant="stop", size="large")

            gr.Markdown("### Waypoints (x,y,z one per line)")
            demo_value = "0,0,15\n30,0,20\n30,30,15\n0,30,10"
            wp_box = gr.Textbox(lines=6, value=demo_value, placeholder="0,0,15\n30,0,20...")

            with gr.Row():
                load_btn = gr.Button("Load Mission", variant="secondary")
                demo_btn = gr.Button("Load Demo Mission", variant="secondary")

        with gr.Column(scale=2):
            gr.Markdown("### Live Telemetry")
            with gr.Row():
                alt_g = gr.Plot(label="Altitude")
                batt_g = gr.Plot(label="Battery")
                link_g = gr.Plot(label="Comms Link")
            with gr.Row():
                status = gr.Textbox(label="Status", interactive=False)
                position = gr.Textbox(label="Position (m)", interactive=False)
                attitude = gr.Textbox(label="Attitude (°)", interactive=False)
            path = gr.Plot(label="Flight Map")

    with gr.Tabs():
        with gr.Tab("Plots"):
            with gr.Row():
                alt_h = gr.Plot()
                batt_h = gr.Plot()
        with gr.Tab("Ground Station Log"):
            log_box = gr.Textbox(lines=18, autoscroll=True, interactive=False)
        with gr.Tab("Fail-Safe Tester"):
            gr.Markdown("**Inject failures to see autonomous response**")
            with gr.Row():
                low_batt_btn = gr.Button("Low Battery")
                gps_btn = gr.Button("GPS Loss")
                comms_btn = gr.Button("Comms Loss")
                imu_btn = gr.Button("IMU Drift")

    with gr.Row():
        advance_btn = gr.Button("Advance 1 s", variant="primary")
        auto_btn = gr.Button("Auto Simulate 10 s", variant="primary")
        reset_btn = gr.Button("Reset", variant="secondary")

    def ui_outputs():
        return [alt_g, batt_g, link_g, status, position, attitude, path, alt_h, batt_h, log_box]

    def update_ui(s):
        return (
            gauge(s.z, "Altitude", 50, " m", "#ffaa00"),
            gauge(s.battery, "Battery", 100, " %", "#00ccff"),
            gauge(s.link_quality, "Link", 100, " %", "#00ffcc"),
            f"{s.mode} | Armed: {s.armed} | Failsafe: {s.failsafe_active}",
            f"{s.x:.1f}, {s.y:.1f}, {s.z:.1f}",
            f"Roll {s.roll:.1f}° / Pitch {s.pitch:.1f}° / Yaw {s.yaw:.1f}°",
            path_plot(s), alt_plot(s), batt_plot(s),
            "\n".join(s.logs)
        )

    # ====================== EVENTS ======================
    arm_btn.click(lambda s: (s if s.arm() else s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())
    disarm_btn.click(lambda s: (s if s.disarm() else s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())
    takeoff_btn.click(lambda s: (setattr(s, 'mode', 'TAKEOFF') or s.update() or s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())
    land_btn.click(lambda s: (setattr(s, 'mode', 'LAND') or s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())
    rtl_btn.click(lambda s: (setattr(s, 'mode', 'RTL') or s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())
    emergency_btn.click(lambda s: (s if s.disarm() else s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())

    load_btn.click(lambda s, txt: (s if s.load_waypoints(txt) else s, *update_ui(s)), inputs=[sim_state, wp_box], outputs=[sim_state] + ui_outputs())
    demo_btn.click(lambda s: (s if s.load_waypoints(demo_value) else s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())

    low_batt_btn.click(lambda s: (s.trigger_failure("low_battery") or s.update() or s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())
    gps_btn.click(lambda s: (s.trigger_failure("gps_loss") or s.update() or s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())
    comms_btn.click(lambda s: (s.trigger_failure("comm_loss") or s.update() or s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())
    imu_btn.click(lambda s: (s.trigger_failure("imu") or s.update() or s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())

    advance_btn.click(lambda s: (s.update() or s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())
    auto_btn.click(lambda s: ([s.update() for _ in range(20)] or s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())
    reset_btn.click(lambda s: (s.reset() or s, *update_ui(s)), inputs=sim_state, outputs=[sim_state] + ui_outputs())

    demo.load(lambda s: update_ui(s), inputs=sim_state, outputs=ui_outputs())

# ====================== LAUNCH (Gradio 6+ compatible) ======================
if __name__ == "__main__":
    demo.launch(
        theme=dark_theme,
        css=custom_css,
        server_name="0.0.0.0",
        server_port=7860
    )