from __future__ import annotations

import streamlit as st
import streamlit.components.v1 as components

from multi_swarm.dashboard.aquarium import (
    STYLE_COLORS,
    build_aquarium_html,
    build_fish_dataset,
    build_lineage_index,
)
from multi_swarm.dashboard.data import (
    evaluations_df,
    genomes_df,
    get_repo,
    list_runs_df,
)

st.title("Aquarium 2D")
st.caption(
    "Pesci colorati per stile cognitivo, dimensione proporzionale a fitness. "
    "Click su un pesce per dettaglio + discendenza."
)

db_path = st.session_state.get("db_path", "./runs.db")
repo = get_repo(db_path)

runs = list_runs_df(repo)
if runs.empty:
    st.info("Nessuna run nel database.")
    st.stop()

selected_run = st.selectbox("Run", runs["id"].tolist())

# Fetch ALL genomes of the run (no gen filter): needed to build the lineage
# index across generations. The active set is filtered afterwards.
all_genomes = genomes_df(repo, selected_run)
all_evals = evaluations_df(repo, selected_run)

if all_genomes.empty:
    st.warning("Nessun genoma per questa run.")
    st.stop()

available_gens = sorted(all_genomes["generation_idx"].unique().tolist())
selected_gen = st.selectbox(
    "Generazione",
    available_gens,
    index=len(available_gens) - 1,  # default ultima
)

active_genomes = all_genomes[all_genomes["generation_idx"] == selected_gen]
active_evals = (
    all_evals[all_evals["genome_id"].isin(active_genomes["id"])]
    if not all_evals.empty
    else all_evals
)
if not active_evals.empty:
    active_merged = active_genomes.merge(
        active_evals,
        left_on="id",
        right_on="genome_id",
        how="left",
        suffixes=("", "_eval"),
    )
else:
    active_merged = active_genomes.copy()
    active_merged["genome_id"] = active_merged["id"]

lineage = build_lineage_index(all_genomes, all_evals)
fish = build_fish_dataset(active_merged, lineage, max_lineage_levels=5)

if not fish:
    st.warning("Nessun agente attivo in questa generazione.")
    st.stop()

st.caption(f"{len(fish)} agenti in generazione {selected_gen}")

html_str = build_aquarium_html(fish, canvas_w=1000, canvas_h=600)
components.html(html_str, height=620, scrolling=False)

with st.expander("Legenda colori"):
    legend_md = "\n".join(
        f"- <span style='color:{color};font-weight:bold;'>&#9679;</span> "
        f"`{style}`"
        for style, color in STYLE_COLORS.items()
    )
    st.markdown(legend_md, unsafe_allow_html=True)