"""
Local Streamlit dashboard for Tropiflo run artifacts.

This module serves two purposes:
1. Provides `launch_dashboard()` for the CLI to start Streamlit.
2. Implements the Streamlit app that reads local results from disk.
"""

from __future__ import annotations

import argparse
import difflib
import json
import socket
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any


@dataclass
class TrialRecord:
    run_name: str
    run_path: Path
    trial_path: Path
    sequence: int
    name: str
    hypothesis: str
    kpi: float | None
    code_version_id: str
    timestamp: str | None
    runtime_ms: float | None
    return_code: int | None
    hypothesis_outcome: str | None
    files: list[str]
    stdout: str | None
    stderr: str | None


def _is_port_open(port: int) -> bool:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.settimeout(0.3)
        return sock.connect_ex(("127.0.0.1", port)) == 0


def launch_dashboard(results_root: Path, port: int = 8501) -> subprocess.Popen | None:
    """
    Launch the Streamlit dashboard in a detached subprocess.

    Returns:
        Popen object if launched, None if already running on the port.
    """
    if _is_port_open(port):
        return None

    dashboard_script = Path(__file__).resolve()
    command = [
        sys.executable,
        "-m",
        "streamlit",
        "run",
        str(dashboard_script),
        "--server.headless",
        "true",
        "--server.port",
        str(port),
        "--",
        "--results-root",
        str(results_root),
    ]

    return subprocess.Popen(  # noqa: S603
        command,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        start_new_session=True,
    )


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--results-root", type=str, default="results/runs")
    args, _ = parser.parse_known_args()
    return args


def _load_metadata(path: Path) -> dict[str, Any]:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return {}


def _discover_runs(results_root: Path) -> list[Path]:
    if not results_root.exists() or not results_root.is_dir():
        return []
    run_dirs = [d for d in results_root.iterdir() if d.is_dir()]
    run_dirs.sort(key=lambda p: p.stat().st_mtime, reverse=True)
    return run_dirs


def _load_trials_for_run(run_dir: Path) -> list[TrialRecord]:
    timeline = run_dir / "timeline"
    if not timeline.exists() or not timeline.is_dir():
        return []

    trials: list[TrialRecord] = []
    for item in timeline.iterdir():
        if not item.is_dir():
            continue
        metadata = _load_metadata(item / "metadata.json")
        if not metadata:
            continue

        trials.append(
            TrialRecord(
                run_name=run_dir.name,
                run_path=run_dir,
                trial_path=item,
                sequence=int(metadata.get("sequence", 0) or 0),
                name=str(metadata.get("name", "unknown")),
                hypothesis=str(metadata.get("hypothesis", "") or ""),
                kpi=metadata.get("kpi"),
                code_version_id=str(metadata.get("code_version_id", "")),
                timestamp=metadata.get("timestamp"),
                runtime_ms=metadata.get("runtime_ms"),
                return_code=metadata.get("return_code"),
                hypothesis_outcome=metadata.get("hypothesis_outcome"),
                files=list(metadata.get("files", [])),
                stdout=metadata.get("stdout"),
                stderr=metadata.get("stderr"),
            )
        )

    trials.sort(key=lambda t: t.sequence)
    return trials


def _pick_baseline(trials: list[TrialRecord]) -> TrialRecord | None:
    if not trials:
        return None
    for t in trials:
        if t.name == "baseline":
            return t
    return min(trials, key=lambda t: t.sequence)


def _read_code_file(trial: TrialRecord, rel_path: str) -> str:
    path = trial.trial_path / rel_path
    if not path.exists():
        return ""
    try:
        return path.read_text(encoding="utf-8")
    except Exception:
        return ""


def _render_run_tab(run_dir: Path):
    import streamlit as st
    import pandas as pd

    trials = _load_trials_for_run(run_dir)
    if not trials:
        st.warning(f"No timeline data found for `{run_dir.name}`.")
        return

    baseline = _pick_baseline(trials)
    best = max((t for t in trials if t.kpi is not None), key=lambda t: t.kpi, default=None)

    st.subheader(f"Workflow: {run_dir.name}")
    col1, col2, col3 = st.columns(3)
    col1.metric("Trials", len(trials))
    col2.metric("Best KPI", f"{best.kpi:.6f}" if best and best.kpi is not None else "N/A")
    col3.metric("Baseline KPI", f"{baseline.kpi:.6f}" if baseline and baseline.kpi is not None else "N/A")

    chart_rows = []
    best_rows = []
    running_best = None
    for t in trials:
        if t.kpi is None:
            continue

        chart_rows.append(
            {
                "sequence": t.sequence,
                "kpi": t.kpi,
            }
        )
        if running_best is None or t.kpi >= running_best:
            running_best = t.kpi
            best_rows.append(
                {
                    "sequence": t.sequence,
                    "kpi": t.kpi,
                }
            )

    if chart_rows:
        import altair as alt

        st.markdown("#### KPI over time")
        all_points_df = pd.DataFrame(chart_rows)
        best_points_df = pd.DataFrame(best_rows)

        points = (
            alt.Chart(all_points_df)
            .mark_circle(size=55, opacity=0.55, color="#9CA3AF")
            .encode(
                x=alt.X("sequence:Q", title="Sequence"),
                y=alt.Y("kpi:Q", title="KPI"),
                tooltip=["sequence", "kpi"],
            )
        )

        best_line = (
            alt.Chart(best_points_df)
            .mark_line(color="#2563EB", strokeWidth=3)
            .encode(
                x="sequence:Q",
                y="kpi:Q",
                tooltip=["sequence", "kpi"],
            )
        )

        best_points = (
            alt.Chart(best_points_df)
            .mark_circle(size=90, color="#2563EB")
            .encode(
                x="sequence:Q",
                y="kpi:Q",
                tooltip=["sequence", "kpi"],
            )
        )

        chart = points + best_line + best_points

        if baseline and baseline.kpi is not None:
            baseline_df = pd.DataFrame(
                [
                    {
                        "sequence": baseline.sequence,
                        "kpi": baseline.kpi,
                        "label": "Baseline",
                    }
                ]
            )
            baseline_point = (
                alt.Chart(baseline_df)
                .mark_point(shape="diamond", size=180, color="#DC2626", filled=True)
                .encode(
                    x="sequence:Q",
                    y="kpi:Q",
                    tooltip=["label", "sequence", "kpi"],
                )
            )
            chart = chart + baseline_point

        st.altair_chart(chart.interactive(), use_container_width=True)
    else:
        st.info("No KPI values found for this workflow yet.")

    table_rows = []
    for t in trials:
        table_rows.append(
            {
                "sequence": t.sequence,
                "name": t.name,
                "kpi": t.kpi,
                "runtime_ms": t.runtime_ms,
                "return_code": t.return_code,
                "outcome": t.hypothesis_outcome,
                "hypothesis": t.hypothesis,
                "files": ", ".join(t.files),
            }
        )
    df_table = pd.DataFrame(table_rows).sort_values("sequence")

    st.markdown("#### Hypotheses tried")
    st.dataframe(df_table, use_container_width=True, hide_index=True)

    options = {
        f"{t.sequence:04d} | {t.name} | KPI={t.kpi}": t
        for t in trials
    }
    widget_prefix = f"run_{run_dir.name}"
    selected_label = st.selectbox(
        "Select a trial",
        list(options.keys()),
        index=max(0, len(options) - 1),
        key=f"{widget_prefix}_trial_selectbox",
    )
    selected_trial = options[selected_label]

    st.markdown("#### Trial details")
    detail_col1, detail_col2 = st.columns(2)
    detail_col1.write(f"**Code version ID:** `{selected_trial.code_version_id}`")
    detail_col1.write(f"**Timestamp:** `{selected_trial.timestamp or 'N/A'}`")
    detail_col1.write(f"**Runtime (ms):** `{selected_trial.runtime_ms}`")
    detail_col1.write(f"**Return code:** `{selected_trial.return_code}`")
    detail_col2.write(f"**Outcome:** `{selected_trial.hypothesis_outcome or 'N/A'}`")
    detail_col2.write(f"**Files:** `{', '.join(selected_trial.files) if selected_trial.files else 'N/A'}`")
    detail_col2.write(f"**Hypothesis:** {selected_trial.hypothesis or 'N/A'}")

    files_for_diff = selected_trial.files[:]
    if not files_for_diff:
        discovered = [p.relative_to(selected_trial.trial_path).as_posix() for p in selected_trial.trial_path.rglob("*.py")]
        files_for_diff = sorted(discovered)

    if files_for_diff:
        selected_file = st.selectbox(
            "File for baseline diff",
            files_for_diff,
            key=f"{widget_prefix}_file_selectbox",
        )
        trial_code = _read_code_file(selected_trial, selected_file)
        baseline_code = _read_code_file(baseline, selected_file) if baseline else ""
        if baseline:
            diff = "\n".join(
                difflib.unified_diff(
                    baseline_code.splitlines(),
                    trial_code.splitlines(),
                    fromfile=f"baseline/{selected_file}",
                    tofile=f"selected/{selected_file}",
                    lineterm="",
                )
            )
            st.markdown("#### Diff vs baseline")
            st.code(diff or "No changes for this file.", language="diff")
        else:
            st.info("No baseline found for diff.")

    with st.expander("Stdout"):
        st.code(selected_trial.stdout or "(empty)", language="text")
    with st.expander("Stderr"):
        st.code(selected_trial.stderr or "(empty)", language="text")


def run_streamlit_app():
    import streamlit as st

    args = _parse_args()
    results_root = Path(args.results_root).resolve()

    st.set_page_config(
        page_title="Tropiflo Dashboard",
        page_icon="T",
        layout="wide",
    )
    st.title("Tropiflo Local Experiment Dashboard")
    st.caption(f"Reading local artifacts from `{results_root}`")

    if st.button("Refresh runs"):
        st.rerun()

    run_dirs = _discover_runs(results_root)
    if not run_dirs:
        st.warning("No runs found yet. Start a workflow with `tropiflo run` first.")
        return

    run_names = [r.name for r in run_dirs]
    selected_run_names = st.sidebar.multiselect(
        "Workflows",
        run_names,
        default=run_names[: min(3, len(run_names))],
    )

    if not selected_run_names:
        st.info("Select one or more workflows from the sidebar.")
        return

    selected_dirs = [r for r in run_dirs if r.name in selected_run_names]
    tabs = st.tabs([r.name for r in selected_dirs])
    for tab, run_dir in zip(tabs, selected_dirs):
        with tab:
            _render_run_tab(run_dir)


if __name__ == "__main__":
    run_streamlit_app()
