"""Kernel-side preview builders for the OmicVerse JupyterLab plugin."""

from __future__ import annotations

import ast
from collections import OrderedDict
from typing import Any, Dict, Mapping, Optional

import pandas as pd
from IPython import get_ipython
from IPython.core.formatters import JSONFormatter

DATAFRAME_MIME_TYPE = "application/vnd.omicverse.dataframe+json"
ANNDATA_MIME_TYPE = "application/vnd.omicverse.anndata+json"
_PREVIEW_REGISTRY: "OrderedDict[str, Any]" = OrderedDict()
_PREVIEW_REGISTRY_LIMIT = 128


def _json_safe_frame(frame: pd.DataFrame, max_rows: int, max_cols: int) -> Dict[str, Any]:
    preview = frame.iloc[:max_rows, :max_cols].copy()
    preview = preview.astype(object).where(pd.notna(preview), None)
    return {
        "columns": [str(col) for col in preview.columns],
        "index": [str(idx) for idx in preview.index],
        "data": preview.values.tolist(),
    }


def _df_payload(
    frame: pd.DataFrame,
    name: Optional[str] = None,
    max_rows: int = 40,
    max_cols: int = 12,
) -> Dict[str, Any]:
    return {
        "type": "dataframe",
        "name": name,
        "shape": [int(frame.shape[0]), int(frame.shape[1])],
        "dtypes": {str(col): str(dtype) for col, dtype in frame.dtypes.items()},
        "table": _json_safe_frame(frame, max_rows=max_rows, max_cols=max_cols),
    }


def _repr_payload(
    value: Any,
    name: Optional[str] = None,
    max_chars: int = 4000,
) -> Dict[str, Any]:
    try:
        text = repr(value)
    except Exception:
        text = "<unavailable>"
    return {
        "type": "content",
        "name": name,
        "content": text[:max_chars],
    }


def _matrix_payload(
    value: Any,
    name: Optional[str] = None,
    max_rows: int = 12,
    max_cols: int = 12,
) -> Dict[str, Any]:
    try:
        import numpy as np
    except Exception:
        return _repr_payload(value, name=name)

    try:
        from scipy import sparse
    except Exception:
        sparse = None

    array = value
    if sparse is not None and sparse.issparse(array):
        array = array[:max_rows, :max_cols].toarray()
    elif hasattr(array, "shape") and len(getattr(array, "shape", ())) >= 2:
        array = array[:max_rows, :max_cols]

    if not isinstance(array, np.ndarray):
        try:
            array = np.asarray(array)
        except Exception:
            return _repr_payload(value, name=name)

    if array.ndim == 1:
        array = array.reshape(-1, 1)
    if array.ndim != 2:
        return {
            "type": "array",
            "name": name,
            "shape": [int(dim) for dim in getattr(value, "shape", array.shape)],
            "dtype": str(getattr(value, "dtype", array.dtype)),
            "content": repr(value)[:4000],
        }

    table = {
        "columns": [str(i) for i in range(array.shape[1])],
        "index": [str(i) for i in range(array.shape[0])],
        "data": array.tolist(),
    }
    return {
        "type": "array",
        "name": name,
        "shape": [int(dim) for dim in getattr(value, "shape", array.shape)],
        "dtype": str(getattr(value, "dtype", array.dtype)),
        "table": table,
    }


def _pack_keys(keys: Any, limit: int) -> Dict[str, Any]:
    values = [str(key) for key in list(keys)]
    return {
        "keys": values,
        "total": len(values),
        "more": max(0, len(values) - limit),
    }


def _is_embedding_like(value: Any) -> bool:
    shape = getattr(value, "shape", None)
    if shape is None:
        return False
    try:
        return len(shape) >= 2 and int(shape[1]) >= 2
    except Exception:
        return False


def _pack_embedding_keys(mapping: Mapping[str, Any], limit: int) -> Dict[str, Any]:
    values = [str(key) for key, value in mapping.items() if _is_embedding_like(value)]
    return {
        "keys": values,
        "total": len(values),
        "more": max(0, len(values) - limit),
    }


def _mapping_previews(
    mapping: Mapping[str, Any],
    slot_name: str,
    max_items: int,
    max_rows: int,
    max_cols: int,
) -> Dict[str, Any]:
    previews: Dict[str, Any] = {}
    for key in list(mapping.keys())[:max_items]:
        label = f'{slot_name}["{key}"]'
        value = mapping[key]
        if isinstance(value, pd.DataFrame):
            previews[str(key)] = _df_payload(value, name=label, max_rows=max_rows, max_cols=max_cols)
        elif hasattr(value, "shape"):
            previews[str(key)] = _matrix_payload(value, name=label, max_rows=min(max_rows, 10), max_cols=min(max_cols, 10))
        else:
            previews[str(key)] = _repr_payload(value, name=label)
    return previews


def _anndata_payload(
    value: Any,
    name: Optional[str] = None,
    max_rows: int = 24,
    max_cols: int = 10,
    key_limit: int = 18,
    nested_preview_limit: int = 3,
) -> Dict[str, Any]:
    obs_pack = _pack_keys(value.obs.columns, key_limit)
    var_pack = _pack_keys(value.var.columns, key_limit)
    uns_pack = _pack_keys(value.uns.keys() if getattr(value, "uns", None) else [], key_limit)
    obsm_pack = _pack_keys(value.obsm.keys() if getattr(value, "obsm", None) else [], key_limit)
    embedding_pack = _pack_embedding_keys(value.obsm, key_limit) if getattr(value, "obsm", None) else _pack_keys([], key_limit)
    layers_pack = _pack_keys(value.layers.keys() if getattr(value, "layers", None) else [], key_limit)

    return {
        "type": "anndata",
        "name": name,
        "summary": {
            "shape": [int(value.n_obs), int(value.n_vars)],
            "obs_columns": obs_pack["keys"],
            "obs_columns_total": obs_pack["total"],
            "obs_columns_more": obs_pack["more"],
            "var_columns": var_pack["keys"],
            "var_columns_total": var_pack["total"],
            "var_columns_more": var_pack["more"],
            "uns_keys": uns_pack["keys"],
            "uns_keys_total": uns_pack["total"],
            "uns_keys_more": uns_pack["more"],
            "obsm_keys": obsm_pack["keys"],
            "obsm_keys_total": obsm_pack["total"],
            "obsm_keys_more": obsm_pack["more"],
            "embedding_keys": embedding_pack["keys"],
            "embedding_keys_total": embedding_pack["total"],
            "embedding_keys_more": embedding_pack["more"],
            "layers": layers_pack["keys"],
            "layers_total": layers_pack["total"],
            "layers_more": layers_pack["more"],
        },
        "previews": {
            "obs": _df_payload(value.obs, name=f"{name}.obs" if name else ".obs", max_rows=max_rows, max_cols=max_cols),
            "var": _df_payload(value.var, name=f"{name}.var" if name else ".var", max_rows=max_rows, max_cols=max_cols),
            "uns": _mapping_previews(value.uns, f"{name}.uns" if name else "uns", nested_preview_limit, max_rows, max_cols)
            if getattr(value, "uns", None)
            else {},
            "obsm": _mapping_previews(value.obsm, f"{name}.obsm" if name else "obsm", nested_preview_limit, max_rows, max_cols)
            if getattr(value, "obsm", None)
            else {},
            "layers": _mapping_previews(
                value.layers,
                f"{name}.layers" if name else "layers",
                nested_preview_limit,
                max_rows,
                max_cols,
            )
            if getattr(value, "layers", None)
            else {},
        },
    }


def _register_preview_value(value: Any) -> str:
    token = f"obj:{id(value):x}"
    _PREVIEW_REGISTRY[token] = value
    _PREVIEW_REGISTRY.move_to_end(token)
    while len(_PREVIEW_REGISTRY) > _PREVIEW_REGISTRY_LIMIT:
        _PREVIEW_REGISTRY.popitem(last=False)
    return token


def _resolve_reference_or_expression(target: str, namespace: Mapping[str, Any]) -> Any:
    if target in _PREVIEW_REGISTRY:
        return _PREVIEW_REGISTRY[target]
    if str(target).startswith("obj:"):
        raise KeyError("This preview is stale after a kernel restart. Re-run the cell to refresh the AnnData output.")
    return resolve_expression(target, namespace)


_OV_SC_COLOR = [
    "#1F577B", "#A56BA7", "#E0A7C8", "#E069A6", "#941456",
    "#FCBC10", "#EF7B77", "#279AD7", "#F0EEF0", "#EAEFC5",
    "#7CBB5F", "#368650", "#A499CC", "#5E4D9A", "#78C2ED",
    "#866017", "#9F987F", "#E0DFED", "#01A0A7", "#75C8CC",
    "#F0D7BC", "#D5B26C", "#D5DA48", "#B6B812", "#9DC3C3",
    "#A89C92", "#FEE00C", "#FEF2A1",
]

_OV_CET = [
    "#d60000", "#8c3bff", "#018700", "#00acc6", "#97ff00", "#ff7ed1", "#6b004f", "#ffa52f",
    "#00009c", "#857067", "#004942", "#4f2a00", "#00fdcf", "#bcb6ff", "#95b379", "#bf03b8",
    "#2466a1", "#280041", "#dbb3af", "#fdf490", "#4f445b", "#a37c00", "#ff7066", "#3f806e",
    "#82000c", "#a37bb3", "#344d00", "#9ae4ff", "#eb0077", "#2d000a", "#5d90ff", "#00c61f",
    "#5701aa", "#001d00", "#9a4600", "#959ea5", "#9a425b", "#001f31", "#c8c300", "#ffcfff",
    "#00bd9a", "#3615ff", "#2d2424", "#df57ff", "#bde6bf", "#7e4497", "#524f3b", "#d86600",
    "#647438", "#c17287", "#6e7489", "#809c03", "#bd8a64", "#623338", "#cacdda", "#6beb82",
    "#213f69", "#a17eff", "#fd03ca", "#75bcfd", "#d8c382", "#cda3cd", "#6d4f00", "#006974",
    "#469e5d", "#93c6bf", "#f9ff00", "#bf5444", "#00643b", "#5b4fa8", "#521f64", "#4f5eff",
    "#7e8e77", "#b808f9", "#8a91c3", "#b30034", "#87607e", "#9e0075", "#ffddc3", "#500800",
    "#1a0800", "#4b89b5", "#00dfdf", "#c8fff9", "#2f3415", "#ff2646", "#ff97aa", "#03001a",
    "#c860b1", "#c3a136", "#7c4f3a", "#f99e77", "#566464", "#d193ff", "#2d1f69", "#411a34",
    "#af9397", "#629e99", "#bcdd7b", "#ff5d93", "#0f2823", "#b8bdac", "#743b64", "#0f000c",
    "#7e6ebc", "#9e6b3b", "#ff4600", "#7e0087", "#ffcd3d", "#2f3b42", "#fda5ff", "#89013d",
]


def _color_to_hex(color: Any) -> str:
    try:
        import matplotlib.colors as mcolors

        return mcolors.to_hex(color)
    except Exception:
        return str(color)


def _default_discrete_colors(n_categories: int) -> list[str]:
    if n_categories <= len(_OV_SC_COLOR):
        base = _OV_SC_COLOR
    elif n_categories <= 56:
        base = _OV_CET[:56]
    else:
        base = _OV_CET[:112]
    return [base[index % len(base)] for index in range(n_categories)]


def _get_uns_colors_for_labels(adata: Any, col_name: str, labels: list[str]) -> Optional[list[str]]:
    key = f"{col_name}_colors"
    if adata is None or key not in getattr(adata, "uns", {}):
        return None

    try:
        uns_colors = list(adata.uns[key])
        original = adata.obs[col_name]
        if hasattr(original, "cat"):
            original_labels = list(original.cat.categories)
            color_map = {
                str(label): _color_to_hex(uns_colors[index % len(uns_colors)])
                for index, label in enumerate(original_labels)
            }
            resolved = [color_map.get(label) for label in labels]
            if all(color is not None for color in resolved):
                return [str(color) for color in resolved]
    except Exception:
        pass

    try:
        return [_color_to_hex(adata.uns[key][index % len(adata.uns[key])]) for index in range(len(labels))]
    except Exception:
        return None


def _embedding_candidates(basis: str) -> list[str]:
    basis = str(basis)
    candidates = [basis]
    if basis.startswith("X_"):
        candidates.extend([basis[2:], basis.upper(), basis.lower()])
    else:
        candidates.extend([f"X_{basis}", basis.upper(), basis.lower()])
    return [candidate for i, candidate in enumerate(candidates) if candidate and candidate not in candidates[:i]]


def _extract_embedding(adata: Any, basis: str) -> tuple[str, Any]:
    if not getattr(adata, "obsm", None):
        raise KeyError("AnnData object has no obsm embeddings")
    for candidate in _embedding_candidates(basis):
        if candidate in adata.obsm:
            return candidate, adata.obsm[candidate]
    raise KeyError(f'Embedding "{basis}" not found in adata.obsm')


def _sample_indices(n_obs: int, max_points: int) -> list[int]:
    if n_obs <= max_points:
        return list(range(n_obs))
    step = max(1, n_obs // max_points)
    sampled = list(range(0, n_obs, step))
    if len(sampled) > max_points:
        sampled = sampled[:max_points]
    return sampled


def _hover_texts(obs_names: list[str], color_label: Optional[str], color_values: list[Any]) -> list[str]:
    if not color_label:
        return obs_names
    texts = []
    for obs_name, color_value in zip(obs_names, color_values):
        label = "NA" if color_value is None else str(color_value)
        texts.append(f"{obs_name}<br>{color_label}: {label}")
    return texts


def plot_embedding_payload(
    target: str,
    basis: str = "X_umap",
    color_by: Optional[str] = None,
    namespace: Optional[Mapping[str, Any]] = None,
    max_points: int = 50000,
) -> Dict[str, Any]:
    if namespace is None:
        ipython = get_ipython()  # type: ignore[name-defined]
        if ipython is None:
            raise RuntimeError("No active IPython shell was found")
        namespace = ipython.user_ns

    adata = _resolve_reference_or_expression(target, namespace)
    if adata.__class__.__name__ != "AnnData":
        raise TypeError("Target is not an AnnData object")

    try:
        import numpy as np
    except Exception as exc:  # pragma: no cover - numpy is a runtime dependency for this feature
        raise RuntimeError("numpy is required for embedding previews") from exc

    key, embedding = _extract_embedding(adata, basis)
    coords = np.asarray(embedding)
    if coords.ndim != 2 or coords.shape[1] < 2:
        raise ValueError(f'Embedding "{key}" must be a 2D matrix with at least two columns')

    sampled_idx = _sample_indices(int(coords.shape[0]), max_points=max_points)
    sampled = coords[sampled_idx, :2]
    x = [float(value) for value in sampled[:, 0].tolist()]
    y = [float(value) for value in sampled[:, 1].tolist()]
    obs_names = [str(adata.obs_names[i]) for i in sampled_idx]

    payload: Dict[str, Any] = {
        "type": "embedding",
        "name": target,
        "ref": _register_preview_value(adata),
        "basis": key,
        "total_points": int(coords.shape[0]),
        "shown_points": len(sampled_idx),
        "sampled": len(sampled_idx) < int(coords.shape[0]),
        "x": x,
        "y": y,
    }

    if not color_by:
        payload["color"] = {"mode": "none"}
        payload["hover"] = obs_names
        return payload

    column_name = color_by[4:] if color_by.startswith("obs:") else color_by
    if column_name not in adata.obs.columns:
        payload["color"] = {"mode": "none"}
        payload["hover"] = obs_names
        payload["warning"] = f'obs column "{column_name}" was not found'
        return payload

    series = adata.obs.iloc[sampled_idx][column_name]
    if pd.api.types.is_numeric_dtype(series) and not pd.api.types.is_bool_dtype(series):
        numeric = pd.to_numeric(series, errors="coerce")
        values = [None if pd.isna(value) else float(value) for value in numeric.tolist()]
        finite = [value for value in values if value is not None]
        payload["color"] = {
            "mode": "continuous",
            "column": str(column_name),
            "values": values,
            "min": min(finite) if finite else None,
            "max": max(finite) if finite else None,
        }
        payload["hover"] = _hover_texts(obs_names, str(column_name), values)
        return payload

    categories = series if pd.api.types.is_categorical_dtype(series) else series.astype("category")
    labels = [str(label) for label in categories.cat.categories.tolist()]
    codes = categories.cat.codes.astype(int).tolist()
    values = ["NA" if pd.isna(raw_value) else str(raw_value) for raw_value in series.astype(object).tolist()]
    if any(code < 0 for code in codes):
        labels = [*labels, "NA"]
        na_code = len(labels) - 1
        codes = [na_code if code < 0 else code for code in codes]
    base_palette = _get_uns_colors_for_labels(adata, str(column_name), labels[:-1] if labels and labels[-1] == "NA" else labels)
    if base_palette is None:
        base_palette = _default_discrete_colors(len(labels) - 1 if labels and labels[-1] == "NA" else len(labels))

    payload["color"] = {
        "mode": "categorical",
        "column": str(column_name),
        "labels": labels,
        "codes": codes,
        "palette": base_palette + (["#94a3b8"] if labels and labels[-1] == "NA" else []),
    }
    payload["hover"] = _hover_texts(obs_names, str(column_name), values)
    return payload


def _resolve_anndata_slot_value(adata: Any, slot: str, key: Optional[str] = None) -> Any:
    slot = str(slot)
    if slot in {"obs", "var"}:
        frame = getattr(adata, slot)
        if key is None:
            return frame
        if key not in frame.columns:
            raise KeyError(f'Column "{key}" not found in adata.{slot}')
        return frame[key]

    if slot in {"uns", "obsm", "layers"}:
        mapping = getattr(adata, slot)
        if key is None:
            return mapping
        if key not in mapping:
            raise KeyError(f'Key "{key}" not found in adata.{slot}')
        return mapping[key]

    raise KeyError(f'Unsupported AnnData slot "{slot}"')


def preview_anndata_slot(
    target: str,
    slot: str,
    key: Optional[str] = None,
    namespace: Optional[Mapping[str, Any]] = None,
    **kwargs: Any,
) -> Dict[str, Any]:
    if namespace is None:
        ipython = get_ipython()  # type: ignore[name-defined]
        if ipython is None:
            raise RuntimeError("No active IPython shell was found")
        namespace = ipython.user_ns

    adata = _resolve_reference_or_expression(target, namespace)
    if adata.__class__.__name__ != "AnnData":
        raise TypeError("Target is not an AnnData object")

    value = _resolve_anndata_slot_value(adata, slot, key)
    label = f'adata.{slot}' if key is None else f'adata.{slot}["{key}"]'
    return preview_value(value, name=label, **kwargs)


def preview_value(
    value: Any,
    name: Optional[str] = None,
    max_rows: int = 40,
    max_cols: int = 12,
    key_limit: int = 18,
    nested_preview_limit: int = 3,
) -> Dict[str, Any]:
    if isinstance(value, pd.DataFrame):
        payload = _df_payload(value, name=name, max_rows=max_rows, max_cols=max_cols)
    elif isinstance(value, pd.Series):
        series_name = str(value.name) if value.name is not None else "value"
        payload = _df_payload(
            value.to_frame(name=series_name),
            name=name,
            max_rows=max_rows,
            max_cols=1,
        )
    elif value.__class__.__name__ == "AnnData":
        payload = _anndata_payload(
            value,
            name=name,
            max_rows=min(max_rows, 24),
            max_cols=min(max_cols, 10),
            key_limit=key_limit,
            nested_preview_limit=nested_preview_limit,
        )
    elif hasattr(value, "shape"):
        payload = _matrix_payload(value, name=name, max_rows=min(max_rows, 12), max_cols=min(max_cols, 12))
    else:
        payload = _repr_payload(value, name=name)

    payload["ref"] = _register_preview_value(value)
    return payload


def _resolve_node(node: ast.AST, namespace: Mapping[str, Any]) -> Any:
    if isinstance(node, ast.Name):
        if node.id not in namespace:
            raise KeyError(f'Variable "{node.id}" not found')
        return namespace[node.id]

    if isinstance(node, ast.Attribute):
        base = _resolve_node(node.value, namespace)
        if node.attr.startswith("_"):
            raise KeyError("Private attributes are not allowed")
        return getattr(base, node.attr)

    if isinstance(node, ast.Subscript):
        base = _resolve_node(node.value, namespace)
        key = _resolve_subscript_key(node.slice)
        return base[key]

    raise KeyError("Only names, attributes, and string/integer subscripts are allowed")


def _resolve_subscript_key(node: ast.AST) -> Any:
    if isinstance(node, ast.Constant) and isinstance(node.value, (str, int)):
        return node.value
    if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub) and isinstance(node.operand, ast.Constant):
        if isinstance(node.operand.value, int):
            return -node.operand.value
    raise KeyError("Only string and integer subscripts are allowed")


def resolve_expression(expression: str, namespace: Mapping[str, Any]) -> Any:
    expression = expression.strip()
    if not expression:
        raise KeyError("Expression is empty")
    tree = ast.parse(expression, mode="eval")
    return _resolve_node(tree.body, namespace)


def preview_variable(
    expression: str,
    namespace: Optional[Mapping[str, Any]] = None,
    **kwargs: Any,
) -> Dict[str, Any]:
    if namespace is None:
        ipython = get_ipython()  # type: ignore[name-defined]
        if ipython is None:
            raise RuntimeError("No active IPython shell was found")
        namespace = ipython.user_ns

    value = resolve_expression(expression, namespace)
    return preview_value(value, name=expression, **kwargs)


def _ensure_json_formatter(ipython: Any, mime_type: str) -> JSONFormatter:
    formatter = ipython.display_formatter.formatters.get(mime_type)
    if formatter is None:
        formatter = JSONFormatter(parent=ipython.display_formatter)
        formatter.format_type = mime_type
        formatter.enabled = True
        ipython.display_formatter.formatters[mime_type] = formatter
    return formatter


def _suppress_pandas_html(ipython: Any) -> None:
    html_formatter = ipython.display_formatter.formatters.get("text/html")
    if html_formatter is None:
        return

    # Force pandas objects to skip their default HTML repr so JupyterLab
    # selects the OmicVerse custom MIME renderer instead of the stock table.
    html_formatter.for_type_by_name(
        "pandas.core.frame",
        "DataFrame",
        lambda value: None,
    )
    html_formatter.for_type_by_name(
        "pandas.core.series",
        "Series",
        lambda value: None,
    )


def enable_formatters(ipython: Optional[Any] = None, **kwargs: Any) -> bool:
    if ipython is None:
        ipython = get_ipython()  # type: ignore[name-defined]
    if ipython is None:
        return False

    _suppress_pandas_html(ipython)

    dataframe_formatter = _ensure_json_formatter(ipython, DATAFRAME_MIME_TYPE)
    dataframe_formatter.for_type_by_name(
        "pandas.core.frame",
        "DataFrame",
        lambda value: preview_value(value, **kwargs),
    )
    dataframe_formatter.for_type_by_name(
        "pandas.core.series",
        "Series",
        lambda value: preview_value(value, **kwargs),
    )

    try:
        import anndata

        anndata_formatter = _ensure_json_formatter(ipython, ANNDATA_MIME_TYPE)
        anndata_formatter.for_type(
            anndata.AnnData,
            lambda value: preview_value(value, **kwargs),
        )
    except Exception:
        pass

    return True


def load_ipython_extension(ipython: Optional[Any]) -> None:
    enable_formatters(ipython)
