import streamlit as st
import requests
import os
import pandas as pd
import sqlite3
import altair as alt
from datetime import datetime, timedelta

API_BASE_URL = os.environ.get("API_BASE_URL", "http://api-service:8000")
MODEL_NAME = "Stacked LSTM Autoencoder (L=8, Fine-tuned)"
MODEL_THRESHOLD_DISPLAY = 0.0155

EVAL_ROC_AUC = 0.994
EVAL_FPR = "12.1%"
EVAL_RECALL = "99.7%"

st.set_page_config(
    page_title="NIDS - Real-Time Detection Dashboard",
    page_icon="🛡️",
    layout="wide"
)


def load_data():
    """Read all detections from the FastAPI service into a DataFrame."""
    try:
        resp = requests.get(f"{API_BASE_URL}/detections", timeout=5)
        resp.raise_for_status()
        rows = resp.json()
        df = pd.DataFrame(rows)
        if not df.empty:
            df["timestamp"] = pd.to_datetime(df["timestamp"])
        return df
    except Exception as e:
        st.error(f"Error reading from API: {e}")
        return pd.DataFrame()


def load_active_flow_count():
    """Read the current count of in-progress (not-yet-decided) flows from the API."""
    try:
        resp = requests.get(f"{API_BASE_URL}/active_flows", timeout=5)
        resp.raise_for_status()
        return resp.json().get("active_flows_count", 0)
    except Exception:
        return 0


def render_header():
    """Static header: title only."""
    st.markdown(
        """
        <div style="background-color:#1a1a2e; border:1px solid #2d2d44; border-radius:12px;
                    padding:20px 24px; margin-bottom:20px;">
            <div style="font-size:2em; font-weight:700; color:white;">
                🛡️ Real-Time Malicious Activity Detection
            </div>
        </div>
        """,
        unsafe_allow_html=True
    )


def render_metric_card(label, value, accent):
    """Colored KPI/performance box."""
    st.markdown(
        f'''
        <div style="background-color:{accent}22; border:1px solid {accent}55;
                    border-radius:10px; padding:14px 16px; text-align:center;">
            <div style="color:#ccc; font-size:1.1em; margin-bottom:10px;">{label}</div>
            <div style="color:{accent}; font-size:2.4em; font-weight:700;">{value}</div>
        </div>
        ''',
        unsafe_allow_html=True
    )


def render_flows_table(display_df):
    """Custom HTML table: colored badges, monospace IPs, fixed columns, right-aligned error."""
    rows_html = ""
    for _, row in display_df.iterrows():
        is_malicious = row["decision"] == "MALICIOUS"
        if is_malicious:
            badge = (
                '<span style="background-color:#4a1414; color:#ff6b6b; padding:5px 14px; '
                'border-radius:14px; font-size:1em; white-space:nowrap;">⚠ Malicious</span>'
            )
            error_color = "#ff6b6b"
        else:
            badge = (
                '<span style="background-color:#14301c; color:#4ade80; padding:5px 14px; '
                'border-radius:14px; font-size:1em; white-space:nowrap;">✓ Benign</span>'
            )
            error_color = "#4ade80"

        rows_html += f"""<tr>
<td style="padding:10px 8px; border-bottom:1px solid #222; color:#ddd; font-size:1.05em;">{row['timestamp'].strftime('%Y-%m-%d %H:%M:%S')}</td>
<td style="padding:10px 8px; border-bottom:1px solid #222; font-family:monospace; color:#ddd; font-size:1.05em;">{row['src_ip']}</td>
<td style="padding:10px 8px; border-bottom:1px solid #222; font-family:monospace; color:#ddd; font-size:1.05em;">{row['dst_ip']}</td>
<td style="padding:10px 8px; border-bottom:1px solid #222;">{badge}</td>
<td style="padding:10px 8px; border-bottom:1px solid #222; text-align:right; color:{error_color}; font-family:monospace; font-size:1.1em; font-weight:600;">{row['reconstruction_error']:.3f}</td>
</tr>"""

    table_html = f"""<style>
.flows-table {{ width:100%; border-collapse:collapse; table-layout:fixed; }}
.flows-table th {{ text-align:left; padding:10px; border-bottom:1px solid #333; color:#aaa; font-weight:600; font-size:1.05em; }}
</style>
<table class="flows-table">
<thead>
<tr>
<th style="width:20%;">Time</th>
<th style="width:24%;">Source</th>
<th style="width:24%;">Destination</th>
<th style="width:16%;">Decision</th>
<th style="width:16%; text-align:right;">Error</th>
</tr>
</thead>
<tbody>{rows_html}</tbody>
</table>"""
    st.markdown(
        f'''<div style="max-height:480px; overflow-y:auto; border:1px solid #222; border-radius:8px;">{table_html}</div>''',
        unsafe_allow_html=True
    )


@st.fragment(run_every=2)
def live_section():
    df = load_data()

    if df.empty:
        st.info("Waiting for first data... Run flow_aggregator.py and generate some traffic (curl https://example.com).")
        return

    recent_cutoff = datetime.now() - timedelta(seconds=10)
    recent_malicious = df[(df["decision"] == "MALICIOUS") & (df["timestamp"] >= recent_cutoff)]

    if not recent_malicious.empty:
        last_alert = recent_malicious.iloc[-1]
        st.markdown(
            f"""
            <div style="background-color:#8B0000; padding:18px; border-radius:10px;
                        text-align:center; animation: pulse 1.2s infinite; margin-bottom:20px;">
                <h2 style="color:white; margin:0;">
                    🚨 Malicious activity detected! {last_alert['src_ip']} → {last_alert['dst_ip']}
                    (error: {last_alert['reconstruction_error']:.3f})
                </h2>
            </div>
            <style>
            @keyframes pulse {{
                0% {{ opacity: 1; }}
                50% {{ opacity: 0.6; }}
                100% {{ opacity: 1; }}
            }}
            </style>
            """,
            unsafe_allow_html=True
        )

    total = len(df)
    benign_count = len(df[df["decision"] == "BENIGN"])
    malicious_count = len(df[df["decision"] == "MALICIOUS"])
    detection_rate = (malicious_count / total * 100) if total > 0 else 0
    active_flows_now = load_active_flow_count()

    if "processing_latency_ms" in df.columns and df["processing_latency_ms"].notna().any():
        median_latency = df["processing_latency_ms"].dropna().median()
        latency_display = f"{median_latency:.2f} ms"
    else:
        latency_display = "— ms"

    kpi_cols = st.columns(6)
    with kpi_cols[0]: render_metric_card("Total Decisions", total, "#60a5fa")
    with kpi_cols[1]: render_metric_card("Benign", benign_count, "#4ade80")
    with kpi_cols[2]: render_metric_card("Malicious", malicious_count, "#ff6b6b")
    with kpi_cols[3]: render_metric_card("Detection Rate", f"{detection_rate:.1f}%", "#c084fc")
    with kpi_cols[4]: render_metric_card("Active Flows Now", active_flows_now, "#94a3b8")
    with kpi_cols[5]: render_metric_card("Avg Processing Latency", latency_display, "#2dd4bf")

    st.write("")

    chart_left, chart_right = st.columns(2)

    with chart_left:
        st.markdown("### Model Performance")
        perf_df = pd.DataFrame({
            "Metric": ["ROC-AUC", "Recall", "FPR"],
            "Value": [EVAL_ROC_AUC, 0.997, 0.121]
        })
        perf_chart = alt.Chart(perf_df).mark_bar(size=45, cornerRadiusEnd=6).encode(
            x=alt.X("Metric:N", title=None, sort=None),
            y=alt.Y("Value:Q", title=None, scale=alt.Scale(domain=[0, 1])),
            color=alt.Color(
                "Metric:N",
                scale=alt.Scale(
                    domain=["ROC-AUC", "Recall", "FPR"],
                    range=["#60a5fa", "#4ade80", "#fbbf24"]
                ),
                legend=None
            ),
            tooltip=["Metric:N", alt.Tooltip("Value:Q", format=".3f")]
        ).properties(height=280)
        text_labels = perf_chart.mark_text(
            align="center", baseline="bottom", dy=-5, color="white", fontSize=14
        ).encode(text=alt.Text("Value:Q", format=".1%"))
        st.altair_chart(perf_chart + text_labels, use_container_width=True)

    with chart_right:
        st.markdown("### Decision Breakdown")
        donut_df = pd.DataFrame({
            "Decision": ["Benign", "Malicious"],
            "Count": [benign_count, malicious_count]
        })
        donut_chart = alt.Chart(donut_df).mark_arc(innerRadius=65, outerRadius=115).encode(
            theta=alt.Theta("Count:Q"),
            color=alt.Color(
                "Decision:N",
                scale=alt.Scale(domain=["Benign", "Malicious"], range=["#4ade80", "#ff6b6b"]),
                legend=alt.Legend(title=None, orient="bottom")
            ),
            tooltip=["Decision:N", "Count:Q"]
        ).properties(height=280)
        st.altair_chart(donut_chart, use_container_width=True)

    st.divider()

    if "flow_filter" not in st.session_state:
        st.session_state.flow_filter = "All"

    title_col, btn_col1, btn_col2, btn_col3 = st.columns([4, 1, 1, 1])
    with title_col:
        st.subheader("🗒️ Latest Flows")
    with btn_col1:
        if st.button("All", use_container_width=True):
            st.session_state.flow_filter = "All"
    with btn_col2:
        if st.button("Malicious", use_container_width=True):
            st.session_state.flow_filter = "Malicious"
    with btn_col3:
        if st.button("Benign", use_container_width=True):
            st.session_state.flow_filter = "Benign"

    filtered_df = df.copy()
    if st.session_state.flow_filter == "Malicious":
        filtered_df = filtered_df[filtered_df["decision"] == "MALICIOUS"]
    elif st.session_state.flow_filter == "Benign":
        filtered_df = filtered_df[filtered_df["decision"] == "BENIGN"]

    display_df = filtered_df.tail(20).sort_values("id", ascending=False)[
        ["timestamp", "src_ip", "dst_ip", "dst_port", "stage", "decision", "reconstruction_error"]
    ]

    if display_df.empty:
        st.caption("No flows match this filter yet.")
    else:
        render_flows_table(display_df)

    st.divider()

    st.subheader("📈 Reconstruction Error Over Time")

    threshold_value = df["threshold"].iloc[-1]
    chart_df = df.tail(200).copy()

    base = alt.Chart(chart_df).mark_point(
        filled=True, size=28, opacity=0.75, stroke=None
    ).encode(
        x=alt.X("timestamp:T", title="Time"),
        y=alt.Y("reconstruction_error:Q", title="Reconstruction Error"),
        color=alt.Color(
            "decision:N",
            scale=alt.Scale(domain=["BENIGN", "MALICIOUS"], range=["#2ECC71", "#E74C3C"]),
            legend=alt.Legend(title="Decision")
        ),
        tooltip=["timestamp:T", "src_ip", "dst_ip", "stage", "decision", "reconstruction_error"]
    ).properties(height=350)

    threshold_line = alt.Chart(pd.DataFrame({"y": [threshold_value]})).mark_rule(
        color="orange", strokeDash=[6, 4], size=1.5
    ).encode(y="y:Q")

    st.altair_chart((base + threshold_line).interactive(), use_container_width=True)


render_header()
live_section()
