# src/plotsrv/store.py
from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Callable, Literal

import pandas as pd
from . import config
from .artifacts import Artifact, ArtifactKind, Truncation

IconKey = Literal[
    "unknown",
    "plot",
    "table",
    "text",
    "json",
    "python",
    "markdown",
    "image",
    "html",
    "exception",
]


@dataclass(frozen=True, slots=True)
class ViewMeta:
    """
    Metadata that describes a view shown in the UI dropdown.
    """

    view_id: str
    kind: str  # "none" | "plot" | "table" | "artifact"
    label: str
    section: str | None = None

    icon_key: IconKey = "unknown"


@dataclass(slots=True)
class ViewState:
    """
    Per-view state that the UI can display.

    Each view is independent: plot/table/status.
    """

    kind: str = "none"  # "none" | "plot" | "table"
    icon_key: IconKey = "unknown"
    plot_png: bytes | None = None
    table_df: pd.DataFrame | None = None
    table_html_simple: str | None = None
    status: dict[str, Any] = None  # populated in __post_init__
    table_total_rows: int | None = None
    table_returned_rows: int | None = None
    artifact: Artifact | None = None

    # publish throttling
    last_publish_at: float | None = None  # epoch seconds

    def __post_init__(self) -> None:
        if self.status is None:
            self.status = {
                "last_updated": None,  # ISO string
                "last_duration_s": None,  # float | None
                "last_error": None,  # str | None
            }


def _icon_for_view_kind(
    kind: str, *, artifact_kind: ArtifactKind | None = None
) -> IconKey:
    k = (kind or "none").strip().lower()

    if k == "plot":
        return "plot"
    if k == "table":
        return "table"

    if k == "artifact":
        ak = (artifact_kind or "python").strip().lower()
        if ak in ("text", "json", "python", "markdown", "image", "html"):
            return ak  # type: ignore[return-value]
        if ak == "exception":
            return "exception"
        # fall back
        return "python"

    return "unknown"


# Global store: multi-view

_VIEWS: dict[str, ViewState] = {}
_VIEW_META: dict[str, ViewMeta] = {}

_ACTIVE_VIEW_ID: str = "default"

# ---- Service mode / CLI RunnerService info -----------------------------------

_SERVICE_INFO: dict[str, Any] = {
    "service_mode": False,
    "service_target": None,
    "service_refresh_rate_s": None,
}

_SERVICE_STOP_HOOK: Callable[[], None] | None = None


# Helpers


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _ensure_view(view_id: str) -> ViewState:
    if view_id not in _VIEWS:
        _VIEWS[view_id] = ViewState()
    return _VIEWS[view_id]


def normalize_view_id(
    view_id: str | None, *, section: str | None = None, label: str | None = None
) -> str:
    """
    Normalize incoming view identifiers.

    - If explicit view_id supplied: use it.
    - Otherwise: create view_id from section/label.
    """
    if view_id:
        return str(view_id)

    sec = (section or "default").strip() or "default"
    lab = (label or "default").strip() or "default"
    return f"{sec}:{lab}"


# View registry API


def register_view(
    *,
    view_id: str | None = None,
    section: str | None = None,
    label: str | None = None,
    kind: str = "none",
    icon_key: IconKey | None = None,  # NEW
    activate_if_first: bool = True,
) -> str:
    vid = normalize_view_id(view_id, section=section, label=label)
    st = _ensure_view(vid)

    # allow upgrade of assigned view dropdown menu icon
    if kind in ("plot", "table", "artifact"):
        st.kind = kind

    if icon_key is not None:
        st.icon_key = icon_key

    meta = _VIEW_META.get(vid)
    if meta is None:
        _VIEW_META[vid] = ViewMeta(
            view_id=vid,
            kind=st.kind,
            label=(label or vid),
            section=section,
            icon_key=st.icon_key,
        )
    else:
        _VIEW_META[vid] = ViewMeta(
            view_id=vid,
            kind=st.kind,
            label=(label or meta.label),
            section=(section if section is not None else meta.section),
            icon_key=st.icon_key,
        )

    global _ACTIVE_VIEW_ID
    if (
        activate_if_first
        and (_ACTIVE_VIEW_ID is None or _ACTIVE_VIEW_ID == "default")
        and len(_VIEW_META) == 1
    ):
        _ACTIVE_VIEW_ID = vid

    return vid


def list_views() -> list[ViewMeta]:
    """
    Return UI-ready view metadata.

    Ordering rules:
      - sections listed in config come first, in order
      - remaining sections come afterwards alphabetically
      - labels.<section> listed come first within that section, in order
      - remaining labels come afterwards alphabetically
    """
    metas = list(_VIEW_META.values())

    # Pull in optional configured order
    section_order = config.get_view_order_sections()  # list[str] | None

    section_rank: dict[str, int] = {}
    if section_order:
        for i, s in enumerate(section_order):
            section_rank[s] = i

    def _section_sort_key(sec: str) -> tuple[int, str]:
        # (0, rank) for configured sections, else (1, alpha)
        if sec in section_rank:
            return (0, f"{section_rank[sec]:09d}")
        return (1, sec.lower())

    # Pre-fetch label ranking maps for sections that have them
    label_rank_by_section: dict[str, dict[str, int]] = {}
    for m in metas:
        sec = (m.section or "default").strip() or "default"
        if sec in label_rank_by_section:
            continue
        order = config.get_view_order_labels(sec)
        if order:
            label_rank_by_section[sec] = {lab: i for i, lab in enumerate(order)}

    def _label_sort_key(sec: str, label: str) -> tuple[int, str]:
        ranks = label_rank_by_section.get(sec)
        if ranks and label in ranks:
            return (0, f"{ranks[label]:09d}")
        return (1, label.lower())

    def _key(m: ViewMeta) -> tuple[tuple[int, str], tuple[int, str]]:
        sec = (m.section or "default").strip() or "default"
        return (_section_sort_key(sec), _label_sort_key(sec, m.label))

    return sorted(metas, key=_key)


def set_active_view(view_id: str) -> None:
    global _ACTIVE_VIEW_ID
    _ACTIVE_VIEW_ID = view_id
    _ensure_view(view_id)


def get_active_view_id() -> str:
    return _ACTIVE_VIEW_ID


def get_view_state(view_id: str | None = None) -> ViewState:
    vid = view_id or _ACTIVE_VIEW_ID
    return _ensure_view(vid)


# Backwards-compatible single-view API (uses active view)


def get_kind(view_id: str | None = None) -> str:
    return get_view_state(view_id).kind


def set_plot(png_bytes: bytes, *, view_id: str | None = None) -> None:
    st = get_view_state(view_id)
    vid = view_id or _ACTIVE_VIEW_ID

    st.kind = "plot"
    st.icon_key = _icon_for_view_kind("plot")
    st.plot_png = png_bytes

    st.artifact = Artifact(
        kind="plot",
        obj=png_bytes,
        created_at=datetime.now(timezone.utc),
        view_id=vid,
    )

    st.status["last_updated"] = _now_iso()
    st.status["last_error"] = None

    register_view(
        view_id=vid, kind="plot", icon_key=st.icon_key, activate_if_first=False
    )


def get_plot(*, view_id: str | None = None) -> bytes:
    st = get_view_state(view_id)
    if st.plot_png is None:
        raise LookupError("No plot available")
    return st.plot_png


def has_plot(*, view_id: str | None = None) -> bool:
    st = get_view_state(view_id)
    return st.plot_png is not None


def set_table(
    df: pd.DataFrame,
    html_simple: str | None,
    *,
    view_id: str | None = None,
    total_rows: int | None = None,
    returned_rows: int | None = None,
) -> None:
    st = get_view_state(view_id)
    st.icon_key = _icon_for_view_kind("table")
    vid = view_id or _ACTIVE_VIEW_ID

    st.kind = "table"
    st.table_df = df
    st.table_html_simple = html_simple

    st.table_total_rows = total_rows
    st.table_returned_rows = returned_rows

    st.artifact = Artifact(
        kind="table",
        obj=df,
        created_at=datetime.now(timezone.utc),
        view_id=vid,
    )

    st.status["last_updated"] = _now_iso()
    st.status["last_error"] = None

    register_view(
        view_id=vid, kind="table", icon_key=st.icon_key, activate_if_first=False
    )


def set_artifact(
    *,
    obj: Any,
    kind: ArtifactKind,
    label: str | None = None,
    section: str | None = None,
    view_id: str | None = None,
    truncation: Truncation | None = None,
) -> None:
    st = get_view_state(view_id)
    vid = view_id or _ACTIVE_VIEW_ID

    st.kind = "artifact"
    st.icon_key = _icon_for_view_kind("artifact", artifact_kind=kind)
    st.artifact = Artifact(
        kind=kind,
        obj=obj,
        created_at=datetime.now(timezone.utc),
        label=label,
        section=section,
        view_id=vid,
        truncation=truncation,
    )

    st.status["last_updated"] = _now_iso()
    st.status["last_error"] = None

    register_view(
        view_id=vid, kind="artifact", icon_key=st.icon_key, activate_if_first=False
    )


def has_table(*, view_id: str | None = None) -> bool:
    st = get_view_state(view_id)
    return st.table_df is not None


def has_artifact(*, view_id: str | None = None) -> bool:
    st = get_view_state(view_id)
    return st.artifact is not None


def get_artifact(*, view_id: str | None = None) -> Artifact:
    st = get_view_state(view_id)
    if st.artifact is None:
        raise LookupError("No artifact available")
    return st.artifact


def get_table_df(*, view_id: str | None = None) -> pd.DataFrame:
    st = get_view_state(view_id)
    if st.table_df is None:
        raise LookupError("No table available")
    return st.table_df


def get_table_html_simple(*, view_id: str | None = None) -> str:
    st = get_view_state(view_id)
    if st.table_html_simple is None:
        raise LookupError("No simple HTML table available")
    return st.table_html_simple


def get_table_counts(*, view_id: str | None = None) -> tuple[int | None, int | None]:
    st = get_view_state(view_id)
    return (st.table_total_rows, st.table_returned_rows)


# ------------------------------------------------------------------------------
# Status bookkeeping (per-view)
# ------------------------------------------------------------------------------


def mark_success(*, duration_s: float | None, view_id: str | None = None) -> None:
    st = get_view_state(view_id)
    st.status["last_updated"] = _now_iso()
    st.status["last_duration_s"] = duration_s
    st.status["last_error"] = None


def mark_error(message: str, *, view_id: str | None = None) -> None:
    st = get_view_state(view_id)
    st.status["last_updated"] = _now_iso()
    st.status["last_error"] = message


def get_status(*, view_id: str | None = None) -> dict[str, Any]:
    st = get_view_state(view_id)
    return dict(st.status)


def _parse_iso_utc(s: str | None) -> datetime | None:
    if not s or not isinstance(s, str):
        return None
    try:
        dt = datetime.fromisoformat(s)
    except Exception:
        return None

    if dt.tzinfo is None:
        return dt.replace(tzinfo=timezone.utc)
    return dt.astimezone(timezone.utc)


def get_freshness(*, view_id: str | None = None) -> dict[str, Any]:
    """
    Compute freshness state for a view from its last_updated timestamp and config.

    States:
      - "disabled"
      - "unknown"
      - "ok"
      - "warn"
      - "error"
    """
    vid = view_id or _ACTIVE_VIEW_ID

    enabled = config.get_freshness_enabled()
    expected_every_s = config.get_freshness_expected_every_s(vid)
    warn_after_s = config.get_freshness_warn_after_s(vid)
    overdue_after_s = config.get_freshness_overdue_after_s(vid)

    if not enabled:
        return {
            "enabled": False,
            "state": "disabled",
            "emoji": "",
            "label": "",
            "age_s": None,
            "expected_every_s": expected_every_s,
            "warn_after_s": warn_after_s,
            "overdue_after_s": overdue_after_s,
            "error_after_s": overdue_after_s,  # legacy alias
        }

    if warn_after_s is None and expected_every_s is not None:
        warn_after_s = expected_every_s
    if overdue_after_s is None and warn_after_s is not None:
        overdue_after_s = warn_after_s * 2

    st = get_view_state(vid)
    last_updated_raw = st.status.get("last_updated")
    last_updated_dt = _parse_iso_utc(last_updated_raw)

    if last_updated_dt is None:
        return {
            "enabled": True,
            "state": "unknown",
            "emoji": "⚪",
            "label": "No data yet",
            "age_s": None,
            "expected_every_s": expected_every_s,
            "warn_after_s": warn_after_s,
            "overdue_after_s": overdue_after_s,
            "error_after_s": overdue_after_s,  # legacy alias
        }

    now = datetime.now(timezone.utc)
    age_s = max(0, int((now - last_updated_dt).total_seconds()))

    state = "ok"
    emoji = "✅"
    label = "Fresh"

    if overdue_after_s is not None and age_s >= overdue_after_s:
        state = "error"
        emoji = "❌"
        label = "Overdue"
    elif warn_after_s is not None and age_s >= warn_after_s:
        state = "warn"
        emoji = "⚠️"
        label = "Stale"

    return {
        "enabled": True,
        "state": state,
        "emoji": emoji,
        "label": label,
        "age_s": age_s,
        "expected_every_s": expected_every_s,
        "warn_after_s": warn_after_s,
        "overdue_after_s": overdue_after_s,
        "error_after_s": overdue_after_s,  # legacy alias
    }


# Publish throttling


def should_accept_publish(
    *, view_id: str, update_limit_s: int | None, now_s: float
) -> bool:
    """
    Server-side throttling:
      - if update_limit_s is None: accept
      - else accept only if enough time passed since last publish
    """
    if update_limit_s is None:
        return True

    st = get_view_state(view_id)
    if st.last_publish_at is None:
        st.last_publish_at = now_s
        return True

    if (now_s - st.last_publish_at) >= float(update_limit_s):
        st.last_publish_at = now_s
        return True

    return False


def note_publish(view_id: str, *, now_s: float) -> None:
    st = get_view_state(view_id)
    st.last_publish_at = now_s


# Service info + shutdown control (global)


def set_service_info(
    *, service_mode: bool, target: str | None, refresh_rate_s: int | None
) -> None:
    _SERVICE_INFO["service_mode"] = bool(service_mode)
    _SERVICE_INFO["service_target"] = target
    _SERVICE_INFO["service_refresh_rate_s"] = refresh_rate_s


def get_service_info() -> dict[str, Any]:
    return dict(_SERVICE_INFO)


def set_service_stop_hook(hook: Callable[[], None]) -> None:
    global _SERVICE_STOP_HOOK
    _SERVICE_STOP_HOOK = hook


def clear_service_stop_request() -> None:
    global _SERVICE_STOP_HOOK
    _SERVICE_STOP_HOOK = None


def request_service_stop() -> bool:
    global _SERVICE_STOP_HOOK

    if _SERVICE_STOP_HOOK is None:
        return False

    hook = _SERVICE_STOP_HOOK
    _SERVICE_STOP_HOOK = None
    try:
        hook()
    except Exception:
        pass
    return True


# Reset


def reset() -> None:
    """
    Reset all in-memory state.

    This is mainly used by unit tests to ensure isolation.
    """
    global _VIEWS, _VIEW_META, _ACTIVE_VIEW_ID
    global _SERVICE_INFO, _SERVICE_STOP_HOOK

    _VIEWS = {}
    _VIEW_META = {}
    _ACTIVE_VIEW_ID = "default"

    _SERVICE_INFO = {
        "service_mode": False,
        "service_target": None,
        "service_refresh_rate_s": None,
    }
    _SERVICE_STOP_HOOK = None
