"""System Config tab — parameter editor with presets and JSON import/export."""

import json

import numpy as np
import streamlit as st

from src.composed.myco_surface import MycoSystemConfig
from src.primitives.redemption_curve import PAMMParams
from dashboard.presets import PRESETS, DEFAULT_PRESET


def _apply_preset(name: str) -> None:
    """Bulk-set session_state values from a preset."""
    preset = PRESETS[name]
    for key, value in preset.items():
        st.session_state[f"cfg_{key}"] = value
    st.session_state["cfg_active_preset"] = name


def _build_config() -> MycoSystemConfig:
    """Build a MycoSystemConfig from current session_state."""
    n = st.session_state.get("cfg_n_reserve_assets", 3)
    lambdas = st.session_state.get("cfg_surface_lambdas", [1.0] * n)
    weights = st.session_state.get("cfg_initial_target_weights", [1.0 / n] * n)
    vault_names = st.session_state.get("cfg_vault_names", ["USDC", "ETH", "DAI"][:n])

    return MycoSystemConfig(
        n_reserve_assets=n,
        surface_lambdas=np.array(lambdas[:n]),
        vault_names=vault_names[:n],
        initial_target_weights=np.array(weights[:n]),
        pamm_params=PAMMParams(
            alpha_bar=st.session_state.get("cfg_pamm_alpha_bar", 10.0),
            xu_bar=st.session_state.get("cfg_pamm_xu_bar", 0.8),
            theta_bar=st.session_state.get("cfg_pamm_theta_bar", 0.5),
        ),
        flow_memory=st.session_state.get("cfg_flow_memory", 0.999),
        flow_threshold=st.session_state.get("cfg_flow_threshold", 0.1),
        static_fee=st.session_state.get("cfg_static_fee", 0.003),
        surge_fee_rate=st.session_state.get("cfg_surge_fee_rate", 0.05),
        imbalance_threshold=st.session_state.get("cfg_imbalance_threshold", 0.2),
        max_labor_mint_fraction=st.session_state.get("cfg_max_labor_mint_fraction", 0.1),
        max_subscription_mint_fraction=st.session_state.get("cfg_max_subscription_mint_fraction", 0.1),
        max_staking_bonus_fraction=st.session_state.get("cfg_max_staking_bonus_fraction", 0.05),
    )


def config_to_dict(config: MycoSystemConfig) -> dict:
    """Serialize a MycoSystemConfig to a JSON-safe dict."""
    return {
        "n_reserve_assets": config.n_reserve_assets,
        "vault_names": config.vault_names,
        "surface_lambdas": config.surface_lambdas.tolist() if config.surface_lambdas is not None else None,
        "initial_target_weights": config.initial_target_weights.tolist() if config.initial_target_weights is not None else None,
        "pamm_alpha_bar": config.pamm_params.alpha_bar,
        "pamm_xu_bar": config.pamm_params.xu_bar,
        "pamm_theta_bar": config.pamm_params.theta_bar,
        "static_fee": config.static_fee,
        "surge_fee_rate": config.surge_fee_rate,
        "imbalance_threshold": config.imbalance_threshold,
        "flow_memory": config.flow_memory,
        "flow_threshold": config.flow_threshold,
        "max_labor_mint_fraction": config.max_labor_mint_fraction,
        "max_subscription_mint_fraction": config.max_subscription_mint_fraction,
        "max_staking_bonus_fraction": config.max_staking_bonus_fraction,
    }


def dict_to_config(d: dict) -> MycoSystemConfig:
    """Deserialize a dict into a MycoSystemConfig."""
    n = d.get("n_reserve_assets", 3)
    return MycoSystemConfig(
        n_reserve_assets=n,
        vault_names=d.get("vault_names", ["USDC", "ETH", "DAI"][:n]),
        surface_lambdas=np.array(d["surface_lambdas"]) if d.get("surface_lambdas") else None,
        initial_target_weights=np.array(d["initial_target_weights"]) if d.get("initial_target_weights") else None,
        pamm_params=PAMMParams(
            alpha_bar=d.get("pamm_alpha_bar", 10.0),
            xu_bar=d.get("pamm_xu_bar", 0.8),
            theta_bar=d.get("pamm_theta_bar", 0.5),
        ),
        static_fee=d.get("static_fee", 0.003),
        surge_fee_rate=d.get("surge_fee_rate", 0.05),
        imbalance_threshold=d.get("imbalance_threshold", 0.2),
        flow_memory=d.get("flow_memory", 0.999),
        flow_threshold=d.get("flow_threshold", 0.1),
        max_labor_mint_fraction=d.get("max_labor_mint_fraction", 0.1),
        max_subscription_mint_fraction=d.get("max_subscription_mint_fraction", 0.1),
        max_staking_bonus_fraction=d.get("max_staking_bonus_fraction", 0.05),
    )


def _load_from_dict(d: dict) -> None:
    """Set session_state values from a config dict."""
    n = d.get("n_reserve_assets", 3)
    st.session_state["cfg_n_reserve_assets"] = n
    st.session_state["cfg_vault_names"] = d.get("vault_names", ["USDC", "ETH", "DAI"][:n])
    st.session_state["cfg_surface_lambdas"] = d.get("surface_lambdas", [1.0] * n)
    st.session_state["cfg_initial_target_weights"] = d.get("initial_target_weights", [1.0 / n] * n)
    st.session_state["cfg_pamm_alpha_bar"] = d.get("pamm_alpha_bar", 10.0)
    st.session_state["cfg_pamm_xu_bar"] = d.get("pamm_xu_bar", 0.8)
    st.session_state["cfg_pamm_theta_bar"] = d.get("pamm_theta_bar", 0.5)
    st.session_state["cfg_static_fee"] = d.get("static_fee", 0.003)
    st.session_state["cfg_surge_fee_rate"] = d.get("surge_fee_rate", 0.05)
    st.session_state["cfg_imbalance_threshold"] = d.get("imbalance_threshold", 0.2)
    st.session_state["cfg_flow_memory"] = d.get("flow_memory", 0.999)
    st.session_state["cfg_flow_threshold"] = d.get("flow_threshold", 0.1)
    st.session_state["cfg_max_labor_mint_fraction"] = d.get("max_labor_mint_fraction", 0.1)
    st.session_state["cfg_max_subscription_mint_fraction"] = d.get("max_subscription_mint_fraction", 0.1)
    st.session_state["cfg_max_staking_bonus_fraction"] = d.get("max_staking_bonus_fraction", 0.05)
    st.session_state["cfg_active_preset"] = "Custom"


def _ensure_defaults():
    """Initialize session_state with Balanced preset if not yet set."""
    if "cfg_n_reserve_assets" not in st.session_state:
        _apply_preset(DEFAULT_PRESET)


def render():
    st.header("System Configuration")
    _ensure_defaults()

    # --- Preset buttons ---
    st.subheader("Presets")
    cols = st.columns(len(PRESETS))
    for i, name in enumerate(PRESETS):
        with cols[i]:
            if st.button(name, key=f"preset_{name}", use_container_width=True):
                _apply_preset(name)
                st.rerun()

    active = st.session_state.get("cfg_active_preset", DEFAULT_PRESET)
    st.caption(f"Active preset: **{active}**")

    # --- Bonding Surface Geometry ---
    with st.expander("Bonding Surface Geometry", expanded=True):
        n = st.slider(
            "Reserve assets (n)",
            2, 6,
            value=st.session_state.get("cfg_n_reserve_assets", 3),
            key="w_n_assets",
        )
        # If n changed, reset per-asset lists
        prev_n = st.session_state.get("cfg_n_reserve_assets", 3)
        if n != prev_n:
            st.session_state["cfg_n_reserve_assets"] = n
            default_names = ["USDC", "ETH", "DAI", "WBTC", "SOL", "AVAX"]
            st.session_state["cfg_vault_names"] = default_names[:n]
            st.session_state["cfg_surface_lambdas"] = [1.0] * n
            st.session_state["cfg_initial_target_weights"] = [1.0 / n] * n
            st.session_state["cfg_active_preset"] = "Custom"
            st.rerun()

        st.session_state["cfg_n_reserve_assets"] = n

        # Vault names
        st.markdown("**Vault names**")
        vault_names = st.session_state.get("cfg_vault_names", ["USDC", "ETH", "DAI"][:n])
        name_cols = st.columns(n)
        new_names = []
        for i in range(n):
            with name_cols[i]:
                val = st.text_input(
                    f"Asset {i+1}",
                    value=vault_names[i] if i < len(vault_names) else f"Asset{i+1}",
                    key=f"w_vault_{i}",
                )
                new_names.append(val)
        st.session_state["cfg_vault_names"] = new_names

        # Surface lambdas
        st.markdown("**Surface lambdas**")
        lambdas = st.session_state.get("cfg_surface_lambdas", [1.0] * n)
        lam_cols = st.columns(n)
        new_lams = []
        for i in range(n):
            with lam_cols[i]:
                val = st.slider(
                    f"\u03bb_{new_names[i]}",
                    1.0, 10.0,
                    value=float(lambdas[i]) if i < len(lambdas) else 1.0,
                    step=0.5,
                    key=f"w_lambda_{i}",
                )
                new_lams.append(val)
        st.session_state["cfg_surface_lambdas"] = new_lams

        # Target weights
        st.markdown("**Initial target weights** (auto-normalized)")
        weights = st.session_state.get("cfg_initial_target_weights", [1.0 / n] * n)
        wt_cols = st.columns(n)
        raw_weights = []
        for i in range(n):
            with wt_cols[i]:
                val = st.slider(
                    f"w_{new_names[i]}",
                    0.01, 1.0,
                    value=float(weights[i]) if i < len(weights) else 1.0 / n,
                    step=0.01,
                    key=f"w_weight_{i}",
                )
                raw_weights.append(val)
        total_w = sum(raw_weights)
        normalized = [w / total_w for w in raw_weights]
        st.session_state["cfg_initial_target_weights"] = normalized
        st.caption(f"Normalized: {', '.join(f'{w:.2%}' for w in normalized)}")

    # --- P-AMM Redemption Curve ---
    with st.expander("P-AMM Redemption Curve"):
        c1, c2, c3 = st.columns(3)
        with c1:
            alpha = st.slider(
                "\u0101\u0304 (alpha_bar) — discount curvature",
                1.0, 100.0,
                value=st.session_state.get("cfg_pamm_alpha_bar", 10.0),
                step=1.0,
                key="w_alpha_bar",
            )
            st.session_state["cfg_pamm_alpha_bar"] = alpha
        with c2:
            xu = st.slider(
                "x\u0304_U (xu_bar) — no-discount threshold",
                0.0, 1.0,
                value=st.session_state.get("cfg_pamm_xu_bar", 0.8),
                step=0.05,
                key="w_xu_bar",
            )
            st.session_state["cfg_pamm_xu_bar"] = xu
        with c3:
            theta = st.slider(
                "\u03b8\u0304 (theta_bar) — floor redemption rate",
                0.0, 1.0,
                value=st.session_state.get("cfg_pamm_theta_bar", 0.5),
                step=0.05,
                key="w_theta_bar",
            )
            st.session_state["cfg_pamm_theta_bar"] = theta

    # --- Fees & Flow Dampening ---
    with st.expander("Fees & Flow Dampening"):
        c1, c2 = st.columns(2)
        with c1:
            sf = st.slider(
                "static_fee",
                0.0, 0.05,
                value=st.session_state.get("cfg_static_fee", 0.003),
                step=0.001,
                format="%.3f",
                key="w_static_fee",
            )
            st.session_state["cfg_static_fee"] = sf

            sr = st.slider(
                "surge_fee_rate",
                0.0, 0.5,
                value=st.session_state.get("cfg_surge_fee_rate", 0.05),
                step=0.01,
                key="w_surge_fee",
            )
            st.session_state["cfg_surge_fee_rate"] = sr

            it = st.slider(
                "imbalance_threshold",
                0.0, 0.5,
                value=st.session_state.get("cfg_imbalance_threshold", 0.2),
                step=0.05,
                key="w_imbalance_thresh",
            )
            st.session_state["cfg_imbalance_threshold"] = it

        with c2:
            fm = st.slider(
                "flow_memory",
                0.9, 0.9999,
                value=st.session_state.get("cfg_flow_memory", 0.999),
                step=0.0001,
                format="%.4f",
                key="w_flow_memory",
            )
            st.session_state["cfg_flow_memory"] = fm

            ft = st.slider(
                "flow_threshold",
                0.01, 0.5,
                value=st.session_state.get("cfg_flow_threshold", 0.1),
                step=0.01,
                key="w_flow_threshold",
            )
            st.session_state["cfg_flow_threshold"] = ft

    # --- Commitment Caps ---
    with st.expander("Commitment Caps"):
        c1, c2, c3 = st.columns(3)
        with c1:
            lm = st.slider(
                "max_labor_mint_fraction",
                0.0, 0.3,
                value=st.session_state.get("cfg_max_labor_mint_fraction", 0.1),
                step=0.01,
                key="w_labor_cap",
            )
            st.session_state["cfg_max_labor_mint_fraction"] = lm
        with c2:
            sm = st.slider(
                "max_subscription_mint_fraction",
                0.0, 0.3,
                value=st.session_state.get("cfg_max_subscription_mint_fraction", 0.1),
                step=0.01,
                key="w_sub_cap",
            )
            st.session_state["cfg_max_subscription_mint_fraction"] = sm
        with c3:
            sb = st.slider(
                "max_staking_bonus_fraction",
                0.0, 0.15,
                value=st.session_state.get("cfg_max_staking_bonus_fraction", 0.05),
                step=0.01,
                key="w_staking_cap",
            )
            st.session_state["cfg_max_staking_bonus_fraction"] = sb

    # --- Build and store config ---
    config = _build_config()
    st.session_state["myco_config"] = config

    # --- JSON Export / Import ---
    st.divider()
    exp_col, imp_col = st.columns(2)
    with exp_col:
        st.subheader("Export Config")
        config_json = json.dumps(config_to_dict(config), indent=2)
        st.download_button(
            "Download JSON",
            data=config_json,
            file_name="myco_config.json",
            mime="application/json",
        )
    with imp_col:
        st.subheader("Import Config")
        uploaded = st.file_uploader("Upload JSON", type=["json"], key="cfg_upload")
        if uploaded is not None:
            try:
                d = json.loads(uploaded.read())
                _load_from_dict(d)
                st.success("Config loaded successfully!")
                st.rerun()
            except (json.JSONDecodeError, KeyError) as e:
                st.error(f"Invalid config file: {e}")

    # --- Config Summary ---
    st.divider()
    st.subheader("Active Configuration Summary")
    n = config.n_reserve_assets
    names = ", ".join(config.vault_names[:n])
    st.markdown(
        f"**{n} assets** ({names}) | "
        f"P-AMM: \u0101={config.pamm_params.alpha_bar:.1f}, "
        f"x\u0304_U={config.pamm_params.xu_bar:.2f}, "
        f"\u03b8\u0304={config.pamm_params.theta_bar:.2f} | "
        f"Fees: static={config.static_fee:.3f}, surge={config.surge_fee_rate:.2f} | "
        f"Flow: mem={config.flow_memory:.4f}, thresh={config.flow_threshold:.2f} | "
        f"Caps: labor={config.max_labor_mint_fraction:.0%}, "
        f"sub={config.max_subscription_mint_fraction:.0%}, "
        f"stake={config.max_staking_bonus_fraction:.0%}"
    )