from __future__ import annotations

from collections.abc import Iterable
from typing import Any

import h5py
import numpy as np
import pandas as pd
from anndata import AnnData, read_h5ad
from scipy import sparse

from readseurat.convert.seurat.model import Assay, CustomS3, CustomS4, SeuratObject
from readseurat.convert.utils import (
    _build_var_frame,
    _choose_primary_matrix,
    _ensure_cells_by_features,
    _ensure_features_by_cells,
    _logical_to_int,
    _write_factor_group,
    _write_string_dataset,
)


def convert_to_anndata(
    seurat: SeuratObject,
    assay_name: str | None = None,
    use_counts: bool = False,
    include_layers: bool = True,
) -> AnnData:
    if assay_name is None:
        assay_name = seurat.active_assay

    cells = seurat.cell_names
    meta = seurat.meta_data
    if meta.index.name is None or len(meta.index) != len(cells):
        meta = meta.copy()
        meta.index = cells

    assay = seurat.assays[assay_name]
    x_matrix = _choose_primary_matrix(assay, use_counts)
    x_matrix = _ensure_cells_by_features(x_matrix, len(cells), len(assay.features))

    var = _build_var_frame(assay)
    adata = AnnData(X=x_matrix, obs=meta, var=var)

    if include_layers:
        _add_layers(adata, assay, len(cells))

    _add_reductions(adata, seurat, len(cells))
    _add_graphs(adata, seurat, len(cells))
    _add_images_misc_tools(adata, seurat)

    return adata


def _add_layers(adata: AnnData, assay: Assay, n_cells: int) -> None:
    if assay.counts is not None:
        counts = _ensure_cells_by_features(assay.counts, n_cells, adata.n_vars)
        if counts.shape == adata.shape:
            adata.layers["counts"] = counts

    if assay.scale_data is not None:
        scale_data = _ensure_cells_by_features(assay.scale_data, n_cells, adata.n_vars)
        if scale_data.shape == adata.shape:
            adata.layers["scale.data"] = scale_data
        else:
            scaled_features = assay.scaled_features
            if scaled_features is not None:
                scale_data = _ensure_cells_by_features(assay.scale_data, n_cells, len(scaled_features))
                adata.uns["scale.data"] = {
                    "features": scaled_features,
                    "matrix": scale_data,
                }


def _add_reductions(adata: AnnData, seurat: SeuratObject, n_cells: int) -> None:
    for name, reduction in seurat.reductions.items():
        emb = reduction.cell_embeddings
        if emb.shape[0] != n_cells and emb.shape[1] == n_cells:
            emb = emb.T
        if emb.shape[0] == n_cells:
            adata.obsm[f"X_{name}"] = emb
            adata.uns[f"reduction_{name}"] = {
                "active_assay": reduction.active_assay,
                "key": reduction.key,
                "global": reduction.global_,
                "misc": reduction.misc,
                "jackstraw": reduction.jackstraw,
            }


def _add_graphs(adata: AnnData, seurat: SeuratObject, n_cells: int) -> None:
    for name, graph in seurat.graphs.items():
        matrix = graph.matrix
        if matrix.shape == (n_cells, n_cells):
            adata.obsp[name] = matrix
        elif matrix.T.shape == (n_cells, n_cells):  # type: ignore[attr-defined]
            adata.obsp[name] = matrix.T  # type: ignore[attr-defined]
        adata.uns[f"graph_{name}"] = {"assay_used": graph.assay_used}


def _add_images_misc_tools(adata: AnnData, seurat: SeuratObject) -> None:
    if seurat.images:
        adata.uns["images"] = seurat.images
    adata.uns["misc"] = seurat.misc
    adata.uns["tools"] = seurat.tools
    adata.uns["commands"] = seurat.commands


def convert_to_h5seurat(
    adata: AnnData | str,
    output_path: str,
    assay_name: str = "RNA",
    project: str = "SeuratProject",
    version: str = "5.0.0",
    key: str | None = None,
) -> None:
    """Convert an AnnData object or .h5ad file to h5Seurat.

    Args:
        adata: AnnData object or path to .h5ad file.
        output_path: Output .h5seurat path.
        assay_name: Name for the assay group.
        project: Project attribute value.
        version: Seurat version string to store.
        key: Assay key attribute (defaults to f"{assay_name}_").
    """
    if isinstance(adata, str):
        adata_obj = read_h5ad(adata)
    else:
        adata_obj = adata

    if key is None:
        key = f"{assay_name}_"

    _write_h5seurat(
        adata_obj,
        output_path=output_path,
        assay_name=assay_name,
        project=project,
        version=version,
        key=key,
    )


def _write_h5seurat(
    adata: AnnData,
    output_path: str,
    assay_name: str,
    project: str,
    version: str,
    key: str,
) -> None:
    cell_names = np.asarray(adata.obs_names, dtype=object)
    features = np.asarray(adata.var_names, dtype=object)

    with h5py.File(output_path, "w") as f:
        f.attrs["project"] = _utf8(project)
        f.attrs["active.assay"] = _utf8(assay_name)
        f.attrs["version"] = _utf8(version)

        _write_string_dataset(f, "cell.names", cell_names)

        _write_data_frame_node(f, "meta.data", pd.DataFrame(adata.obs))

        # Write active.ident - required by Seurat's LoadH5Seurat
        _write_active_ident(f, adata)

        assays_group = f.create_group("assays")
        assay_group = assays_group.create_group(assay_name)
        assay_group.attrs["key"] = _utf8(key)
        _write_string_dataset(assay_group, "features", features)

        data_matrix = _ensure_features_by_cells(adata.X, adata.n_obs, adata.n_vars)
        _write_matrix(assay_group, "data", data_matrix)

        if "counts" in adata.layers:
            counts = _ensure_features_by_cells(adata.layers["counts"], adata.n_obs, adata.n_vars)
            _write_matrix(assay_group, "counts", counts)

        _write_scale_data(assay_group, adata, features)
        _write_meta_features(assay_group, pd.DataFrame(adata.var))
        _write_variable_features(assay_group, pd.DataFrame(adata.var), features)

        _write_reductions(f, adata, assay_name)
        _write_graphs(f, adata, assay_name)
        _write_images_misc_tools(f, adata)


def _write_active_ident(root: h5py.Group, adata: AnnData) -> None:
    """Write active.ident as a factor dataset.

    Uses seurat_clusters if available, otherwise leiden, louvain, or defaults to 0.
    """
    n_cells = adata.n_obs

    # Try to find cluster information
    if "seurat_clusters" in adata.obs.columns:
        ident = pd.Series(adata.obs["seurat_clusters"])
    elif "leiden" in adata.obs.columns:
        ident = pd.Series(adata.obs["leiden"])
    elif "louvain" in adata.obs.columns:
        ident = pd.Series(adata.obs["louvain"])
    else:
        # Default to all cells in cluster 0
        ident = pd.Series(np.zeros(n_cells, dtype=int), index=adata.obs.index)

    # Write as factor group
    factor_group = root.create_group("active.ident")
    _write_factor_group(factor_group, ident)


def _write_scale_data(group: h5py.Group, adata: AnnData, features: np.ndarray) -> None:
    if "scale.data" in adata.layers:
        scale = _ensure_features_by_cells(adata.layers["scale.data"], adata.n_obs, adata.n_vars)
        _write_matrix(group, "scale.data", scale, dense_only=True)
        _write_string_dataset(group, "scaled.features", features)
        return
    if "scale.data" in adata.uns:
        payload = adata.uns["scale.data"]
        if isinstance(payload, dict) and "features" in payload and "matrix" in payload:
            scaled_features = np.asarray(payload["features"], dtype=object)
            scale = _ensure_features_by_cells(payload["matrix"], adata.n_obs, len(scaled_features))
            _write_matrix(group, "scale.data", scale, dense_only=True)
            _write_string_dataset(group, "scaled.features", scaled_features)


def _write_meta_features(group: h5py.Group, var: pd.DataFrame) -> None:
    if var is None or var.empty:
        return
    meta = var.copy()
    _write_data_frame_node(group, "meta.features", meta)


def _write_variable_features(
    group: h5py.Group,
    var: pd.DataFrame,
    features: np.ndarray,
) -> None:
    if "variable_features" in var.columns:
        mask = var["variable_features"].astype(bool).to_numpy()
        variable = features[mask]
        _write_string_dataset(group, "variable.features", variable)


def _write_reductions(
    root: h5py.Group,
    adata: AnnData,
    assay_name: str,
) -> None:
    if not adata.obsm:
        return
    reductions = root.create_group("reductions")
    for key, value in adata.obsm.items():
        name = key[2:] if key.startswith("X_") else key
        group = reductions.create_group(name)
        group.attrs["active.assay"] = np.array([_utf8(assay_name)], dtype=object)
        group.attrs["key"] = _utf8(f"{name}_")
        group.attrs["global"] = np.array([1], dtype=np.int8)
        matrix = _ensure_cells_by_features(value, adata.n_obs, value.shape[1])
        # Convert to dense array if sparse
        if sparse.issparse(matrix):
            dense_matrix = sparse.spmatrix.toarray(matrix)  # type: ignore[attr-defined]
        else:
            dense_matrix = np.asarray(matrix)
        _write_dense_matrix(group, "cell.embeddings", dense_matrix)


def _write_graphs(
    root: h5py.Group,
    adata: AnnData,
    assay_name: str,
) -> None:
    if not adata.obsp:
        return
    graphs = root.create_group("graphs")
    for name, matrix in adata.obsp.items():
        group = graphs.create_group(name)
        group.attrs["assay.used"] = _utf8(assay_name)
        mat = _ensure_cells_by_cells(matrix, adata.n_obs)
        _write_sparse_matrix(group, mat)


def _write_images_misc_tools(root: h5py.Group, adata: AnnData) -> None:
    images = root.create_group("images")
    if isinstance(adata.uns.get("images"), dict):
        for name, payload in adata.uns["images"].items():
            _write_list_group(images.create_group(name), payload)

    misc = root.create_group("misc")
    if isinstance(adata.uns.get("misc"), dict):
        _write_list_group(misc, adata.uns["misc"])

    tools = root.create_group("tools")
    if isinstance(adata.uns.get("tools"), dict):
        _write_list_group(tools, adata.uns["tools"])

    commands = root.create_group("commands")
    if isinstance(adata.uns.get("commands"), dict):
        _write_list_group(commands, adata.uns["commands"])


def _write_data_frame_node(group: h5py.Group, name: str, df: pd.DataFrame) -> None:
    df = df.copy()
    if df.index is not None and df.index.name is None:
        df.index.name = "row.names"

    if df.shape[1] == 0:
        subgroup = group.create_group(name)
        _write_data_frame_group(subgroup, df)
        return

    if _requires_group(df):
        subgroup = group.create_group(name)
        _write_data_frame_group(subgroup, df)
    else:
        _write_data_frame_dataset(group, name, df)


def _write_data_frame_dataset(group: h5py.Group, name: str, df: pd.DataFrame) -> None:
    logicals = [
        col_name for col_name, series in df.items() if pd.api.types.is_bool_dtype(series) or _is_pandas_boolean(series)
    ]
    dtype = []
    arrays: dict[str, Any] = {}
    for col_name, series in df.items():
        col_name_str = str(col_name)
        col = series.to_numpy()
        if col_name_str in logicals:
            arrays[col_name_str] = _logical_to_int(col)
            dtype.append((col_name_str, np.int8))
        elif pd.api.types.is_string_dtype(series) or series.dtype == object:
            arrays[col_name_str] = np.asarray(col, dtype=object)
            dtype.append((col_name_str, h5py.string_dtype("utf-8")))
        else:
            arrays[col_name_str] = col
            dtype.append((col_name_str, col.dtype))

    dset = group.create_dataset(name, shape=(len(df),), dtype=np.dtype(dtype))
    for name in arrays:
        dset[name] = arrays[name]
    if logicals:
        dset.attrs["logicals"] = np.array([_utf8(name) for name in logicals], dtype=object)

    if df.index is not None:
        row_names = np.asarray(df.index.astype(str), dtype=object)
        group.create_dataset(
            f"{name}.row.names",
            data=row_names,
            dtype=h5py.string_dtype("utf-8"),
        )


def _write_data_frame_group(group: h5py.Group, df: pd.DataFrame) -> None:
    colnames = [str(name) for name in df.columns]
    group.attrs["colnames"] = np.array([_utf8(name) for name in colnames], dtype=h5py.string_dtype("utf-8"))
    logicals: list[str] = []
    for name, series in df.items():
        name_str = str(name)
        if isinstance(series.dtype, pd.CategoricalDtype):
            _write_factor_group(group.create_group(name_str), series)
        else:
            col = series.to_numpy()
            if pd.api.types.is_bool_dtype(series) or _is_pandas_boolean(series):
                logicals.append(name_str)
                col = _logical_to_int(col)
                dset = group.create_dataset(name_str, data=col, dtype=np.int8)
                dset.attrs["s3class"] = _utf8("logical")
            elif pd.api.types.is_string_dtype(series) or series.dtype == object:
                _write_string_dataset(group, name_str, np.asarray(col, dtype=object))
            else:
                group.create_dataset(name_str, data=col)

    if logicals:
        group.attrs["logicals"] = np.array([_utf8(n) for n in logicals], dtype=object)

    if df.index is not None:
        _write_string_dataset(group, "row.names", np.asarray(df.index.astype(str), dtype=object))


def _write_matrix(group: h5py.Group, name: str, matrix, dense_only: bool = False) -> None:
    if sparse.issparse(matrix):
        if dense_only:
            matrix = matrix.toarray()
            _write_dense_matrix(group, name, matrix)
        else:
            sub = group.create_group(name)
            _write_sparse_matrix(sub, matrix)
        return
    _write_dense_matrix(group, name, np.asarray(matrix))


def _write_dense_matrix(group: h5py.Group, name: str, matrix: np.ndarray) -> None:
    mat = np.asarray(matrix)
    if mat.ndim != 2:
        raise ValueError("Dense matrix must be 2D")
    group.create_dataset(name, data=mat.T)


def _write_sparse_matrix(group: h5py.Group, matrix) -> None:
    mat = sparse.csc_matrix(matrix)
    group.create_dataset("indices", data=mat.indices.astype(np.int64))
    group.create_dataset("indptr", data=mat.indptr.astype(np.int64))
    group.create_dataset("data", data=mat.data)
    group.attrs["dims"] = np.array(mat.shape, dtype=np.int64)


def _write_list_group(group: h5py.Group, payload: Any) -> None:
    if isinstance(payload, CustomS4):
        group.attrs["s4class"] = _utf8(payload.s4class)
        for name, value in payload.data.items():
            _write_list_node(group, name, value)
        return
    if isinstance(payload, CustomS3):
        group.attrs["s3class"] = _utf8(payload.s3class)
        for name, value in payload.data.items():
            _write_list_node(group, name, value)
        return
    if isinstance(payload, dict):
        for name, value in payload.items():
            _write_list_node(group, name, value)
        return
    if isinstance(payload, list):
        for idx, value in enumerate(payload):
            _write_list_node(group, str(idx), value)
        return
    _write_scalar_dataset(group, "value", payload)


def _write_list_node(group: h5py.Group, name: str, value: Any) -> None:
    if isinstance(value, dict) or isinstance(value, list) or isinstance(value, CustomS3) or isinstance(value, CustomS4):
        sub = group.create_group(name)
        _write_list_group(sub, value)
        return
    if sparse.issparse(value):
        sub = group.create_group(name)
        _write_sparse_matrix(sub, value)
        return
    if isinstance(value, pd.DataFrame):
        _write_data_frame_node(group, name, value)
        return
    if isinstance(value, np.ndarray) and value.ndim == 2:
        _write_dense_matrix(group, name, value)
        return
    if isinstance(value, (list, tuple, np.ndarray)):
        _write_vector_dataset(group, name, value)
        return
    _write_scalar_dataset(group, name, value)


def _write_vector_dataset(group: h5py.Group, name: str, values: Iterable[Any]) -> None:
    values = list(values)
    if not values:
        group.create_dataset(name, data=np.array([], dtype=np.float32))
        return
    if isinstance(values[0], str):
        _write_string_dataset(group, name, values)
        return
    group.create_dataset(name, data=np.asarray(values))


def _write_scalar_dataset(group: h5py.Group, name: str, value: Any) -> None:
    if isinstance(value, str):
        dtype = h5py.string_dtype("utf-8")
        group.create_dataset(name, data=_utf8(value), dtype=dtype)
        return
    group.create_dataset(name, data=value)


def _requires_group(df: pd.DataFrame) -> bool:
    for _, series in df.items():
        if isinstance(series.dtype, pd.CategoricalDtype):
            return True
    return False


def _is_pandas_boolean(series: pd.Series) -> bool:
    return isinstance(series.dtype, pd.BooleanDtype)


def _ensure_cells_by_cells(matrix, n_cells: int):
    if sparse.issparse(matrix):
        if matrix.shape == (n_cells, n_cells):
            return matrix.tocsc()
        if matrix.T.shape == (n_cells, n_cells):
            return matrix.T.tocsc()
        return matrix.tocsc()
    arr = np.asarray(matrix)
    if arr.shape == (n_cells, n_cells):
        return sparse.csc_matrix(arr)
    if arr.T.shape == (n_cells, n_cells):
        return sparse.csc_matrix(arr.T)
    return sparse.csc_matrix(arr)


def _utf8(value: Any) -> str:
    if isinstance(value, bytes):
        return value.decode("utf-8")
    return str(value)
