from __future__ import annotations

from typing import Any

import h5py
import numpy as np
import pandas as pd
from scipy import sparse

from readseurat.convert.errors import MissingNodeError, TypeMismatchError
from readseurat.convert.seurat.model import (
    Assay,
    CustomS3,
    CustomS4,
    Graph,
    Reduction,
    SeuratObject,
)


def parse_h5seurat_to_seurat_object(
    path: str,
    row_major: bool = True,
    validate: bool = True,
) -> SeuratObject:
    with h5py.File(path, "r") as file:
        if validate:
            _validate_root(file)
        project = _get_attr_str(file, "project") or "SeuratProject"
        active_assay = _get_attr_str(file, "active.assay") or "RNA"
        version = _get_attr_str(file, "version") or "unknown"
        return SeuratObject(
            project=project,
            active_assay=active_assay,
            version=version,
            cell_names=_get_cells(file),
            meta_data=_get_meta_data(file, row_major),
            assays={name: _get_assay(file, row_major, name) for name in _list_group(file, "assays")},
            reductions={name: _get_reduction(file, row_major, name) for name in _list_group(file, "reductions")},
            graphs={name: _get_graph(file, row_major, name) for name in _list_group(file, "graphs")},
            images={name: _get_image(file, row_major, name) for name in _list_group(file, "images")},
            misc=_get_misc(file, row_major),
            tools=_get_tools(file, row_major),
            commands=_get_commands(file, row_major),
        )


def _get_cells(file: h5py.File) -> np.ndarray:
    return _read_string_dataset(_require_dataset(file, "cell.names"))


def _get_meta_data(file: h5py.File, row_major: bool) -> pd.DataFrame:
    node = _require_node(file, "meta.data")
    return _read_data_frame(file, row_major, node)


def _get_assay(file: h5py.File, row_major: bool, name: str) -> Assay:
    group = _require_group(file, f"assays/{name}")

    # Seurat v4: features stored as an HDF5 string dataset
    # Seurat v5: features stored as int32; gene names are in meta.data/_index
    features_dset = _require_dataset(file, f"assays/{name}/features")
    if h5py.check_string_dtype(features_dset.dtype):
        features = _read_string_dataset(features_dset)
    else:
        meta_index_path = f"assays/{name}/meta.data/_index"
        if meta_index_path in file:
            features = _read_string_dataset(file[meta_index_path])
        else:
            features = np.arange(features_dset.shape[-1]).astype(str)

    # Seurat v4: data/counts/scale.data at assay root
    # Seurat v5: data/counts/scale.data under layers/
    data_path = f"assays/{name}/data"
    if data_path not in file:
        data_path = f"assays/{name}/layers/data"
    data_node = _require_node(file, data_path)
    data = _read_matrix(row_major, data_node)

    # v4 stores key as an HDF5 attribute; v5 stores it as a dataset
    key = _get_attr_str(group, "key", required=False)
    if key is None and "key" in group:
        key_dset = group["key"]
        if isinstance(key_dset, h5py.Dataset) and key_dset.size > 0:
            key = _decode_string(key_dset.asstr()[()].flat[0])
    if not key:
        key = f"{name}_"

    counts = _read_optional_matrix(row_major, group, "counts")
    if counts is None:
        layers = group.get("layers")
        if isinstance(layers, h5py.Group):
            counts = _read_optional_matrix(row_major, layers, "counts")

    scale_data = _read_optional_dense(row_major, group, "scale.data")
    if scale_data is None:
        layers = group.get("layers")
        if isinstance(layers, h5py.Group):
            scale_data = _read_optional_dense(row_major, layers, "scale.data")

    scaled_features = _read_optional_string(group, "scaled.features")
    if scaled_features is None:
        scaled_features = _read_optional_string(group, "scale.features")

    # Seurat v4: meta.features compound dataset; Seurat v5: meta.data group
    meta_features = _read_optional_data_frame(file, row_major, group, "meta.features")
    if meta_features is None:
        meta_features = _read_optional_data_frame(file, row_major, group, "meta.data")

    variable_features = _read_optional_string(group, "variable.features")
    misc = _read_optional_list(row_major, group, "misc")
    return Assay(
        name=name,
        features=features,
        data=data,
        key=key,
        counts=counts,
        scale_data=scale_data,
        scaled_features=scaled_features,
        meta_features=meta_features,
        variable_features=variable_features,
        misc=misc,
    )


def _get_reduction(file: h5py.File, row_major: bool, name: str) -> Reduction:
    group = _require_group(file, f"reductions/{name}")
    active_assay = _get_attr_str_list(group, "active.assay")
    key = _get_attr_str(group, "key") or f"{name}_"
    global_val = _get_attr_logical(group, "global")
    cell_embeddings = _read_dense_matrix(row_major, _require_dataset(file, f"reductions/{name}/cell.embeddings"))
    feature_loadings = _read_optional_dense(row_major, group, "feature.loadings")
    feature_loadings_projected = _read_optional_dense(row_major, group, "feature.loadings.projected")
    misc = _read_optional_list(row_major, group, "misc")
    jackstraw = _read_optional_list(row_major, group, "jackstraw")
    return Reduction(
        name=name,
        active_assay=active_assay,
        key=key,
        global_=global_val,
        cell_embeddings=cell_embeddings,
        feature_loadings=feature_loadings,
        feature_loadings_projected=feature_loadings_projected,
        misc=misc,
        jackstraw=jackstraw,
    )


def _get_graph(file: h5py.File, row_major: bool, name: str) -> Graph:
    group = _require_group(file, f"graphs/{name}")
    matrix = _read_sparse_matrix(group)
    assay_used = _get_attr_str(group, "assay.used", required=False)
    return Graph(name=name, matrix=matrix, assay_used=assay_used)


def _get_image(file: h5py.File, row_major: bool, name: str) -> Any:
    group = _require_group(file, f"images/{name}")
    return _read_list_group(row_major, group)


def _get_misc(file: h5py.File, row_major: bool) -> dict[str, Any]:
    return _read_list_group(row_major, _require_group(file, "misc"))


def _get_tools(file: h5py.File, row_major: bool) -> dict[str, Any]:
    return _read_list_group(row_major, _require_group(file, "tools"))


def _get_commands(file: h5py.File, row_major: bool) -> dict[str, Any]:
    return _read_list_group(row_major, _require_group(file, "commands"))


def _validate_root(file: h5py.File) -> None:
    for attr in ("project", "active.assay", "version"):
        if attr not in file.attrs:
            raise MissingNodeError(f"Missing required attribute: {attr}")
    for dset in ("cell.names", "meta.data"):
        if dset not in file:
            raise MissingNodeError(f"Missing required dataset/group: {dset}")
    for group in ("assays", "misc", "tools", "commands"):
        if group not in file:
            raise MissingNodeError(f"Missing required group: {group}")
    active_assay = _get_attr_str(file, "active.assay")
    if f"assays/{active_assay}" not in file:
        raise MissingNodeError(f"Active assay not found: assays/{active_assay}")


def _list_group(file: h5py.File, name: str) -> list[str]:
    if name not in file:
        return []
    group = file[name]
    if not isinstance(group, h5py.Group):
        raise TypeMismatchError(f"Expected group at {name}")
    return list(group.keys())


def _require_node(file: h5py.File, path: str) -> h5py.Group | h5py.Dataset:
    if path not in file:
        raise MissingNodeError(f"Missing required node: {path}")
    return file[path]


def _require_group(file: h5py.File, path: str) -> h5py.Group:
    node = _require_node(file, path)
    if not isinstance(node, h5py.Group):
        raise TypeMismatchError(f"Expected group at {path}")
    return node


def _require_dataset(file: h5py.File, path: str) -> h5py.Dataset:
    node = _require_node(file, path)
    if not isinstance(node, h5py.Dataset):
        raise TypeMismatchError(f"Expected dataset at {path}")
    return node


def _get_attr_str(
    node: h5py.File | h5py.Group | h5py.Dataset,
    name: str,
    required: bool = True,
) -> str | None:
    attrs = node.attrs
    if name not in attrs:
        if required:
            raise MissingNodeError(f"Missing required attribute: {name}")
        return None
    return _decode_string(attrs[name])


def _get_attr_str_list(node: h5py.Group | h5py.Dataset, name: str) -> list[str]:
    if name not in node.attrs:
        raise MissingNodeError(f"Missing required attribute: {name}")
    return _decode_string_list(node.attrs[name])


def _get_attr_logical(node: h5py.Group, name: str) -> bool:
    if name not in node.attrs:
        raise MissingNodeError(f"Missing required attribute: {name}")
    value = node.attrs[name]
    if isinstance(value, (np.ndarray, list)):
        value = value[0]
    if isinstance(value, (bytes, str)):
        value = _decode_string(value)
    return bool(value)


def _get_attr_str_list_optional(node: h5py.Group | h5py.Dataset, name: str) -> list[str]:
    if name not in node.attrs:
        return []
    return _decode_string_list(node.attrs[name])


def _read_string_dataset(dset: h5py.Dataset) -> np.ndarray:
    data = dset.asstr()[()]
    if isinstance(data, np.ndarray):
        return data
    return np.array([data])


def _read_matrix(row_major: bool, node: h5py.Group | h5py.Dataset):
    if isinstance(node, h5py.Dataset):
        return _read_dense_matrix(row_major, node)
    if isinstance(node, h5py.Group):
        return _read_sparse_matrix(node)
    raise TypeMismatchError("Matrix node must be dataset or group")


def _read_dense_matrix(row_major: bool, dset: h5py.Dataset) -> np.ndarray:
    data = dset[()]
    if data.ndim != 2:
        raise TypeMismatchError("Dense matrix dataset must be 2D")
    if row_major:
        return data.T
    return data


def _read_sparse_matrix(group: h5py.Group) -> sparse.csc_matrix:
    for name in ("indices", "indptr", "data"):
        if name not in group:
            raise MissingNodeError(f"Sparse matrix missing dataset: {name}")
    indices = group["indices"][()]
    indptr = group["indptr"][()]
    data = group["data"][()]
    if "dims" in group.attrs:
        dims = tuple(int(v) for v in group.attrs["dims"])
        shape = (dims[0], dims[1])
    else:
        n_cols = len(indptr) - 1
        n_rows = int(indices.max()) + 1 if len(indices) else 0
        shape = (n_rows, n_cols)
    return sparse.csc_matrix((data, indices, indptr), shape=shape)


def _read_factor_group(group: h5py.Group) -> pd.Categorical:
    levels = _read_string_dataset(group["levels"])
    values = group["values"][()]
    codes = np.array(values, dtype=int) - 1
    return pd.Categorical.from_codes(codes, categories=levels)


def _read_logical(values: np.ndarray | int | float) -> pd.arrays.BooleanArray:
    mapping = {0: False, 1: True, 2: pd.NA}
    if isinstance(values, np.ndarray):
        iterable = values.tolist()
    else:
        iterable = [values]
    converted = [mapping.get(int(v), pd.NA) for v in iterable]
    return pd.array(converted, dtype="boolean")  # type: ignore[return-value]


def _read_data_frame(file: h5py.File, row_major: bool, node: h5py.Group | h5py.Dataset) -> pd.DataFrame:
    if isinstance(node, h5py.Dataset):
        return _read_data_frame_dataset(file, node)
    if isinstance(node, h5py.Group):
        return _read_data_frame_group(row_major, node)
    raise TypeMismatchError("Data frame node must be dataset or group")


def _read_data_frame_dataset(file: h5py.File, dset: h5py.Dataset) -> pd.DataFrame:
    arr = dset[()]
    if arr.dtype.fields is None:
        raise TypeMismatchError("Compound dataset required for data frame")
    data: dict[str, Any] = {}
    logicals = _get_attr_str_list_optional(dset, "logicals")
    for name in arr.dtype.names:
        col = arr[name]
        if name in logicals:
            data[name] = _read_logical(col)
        else:
            data[name] = _decode_column(col)
    df = pd.DataFrame(data)
    row_names_path = f"{dset.name}.row.names"
    if row_names_path in file:
        row_names = _read_string_dataset(file[row_names_path])
        df.index = row_names
    return df


def _read_data_frame_group(row_major: bool, group: h5py.Group) -> pd.DataFrame:
    colnames = _get_attr_str_list_optional(group, "colnames")
    logicals = _get_attr_str_list_optional(group, "logicals")
    keys = colnames if colnames else list(group.keys())
    data: dict[str, Any] = {}
    for key in keys:
        if key in ("row.names", "_index"):
            continue
        node = group[key]
        if isinstance(node, h5py.Group) and _is_factor_group(node):
            data[key] = _read_factor_group(node)
        elif isinstance(node, h5py.Dataset):
            col = node[()]
            if key in logicals:
                data[key] = _read_logical(col)
            else:
                data[key] = _decode_column(col)
        else:
            data[key] = _read_list_group(row_major, node) if isinstance(node, h5py.Group) else node[()]
    df = pd.DataFrame(data)
    if "row.names" in group:
        row_names = _read_string_dataset(group["row.names"])
        df.index = row_names
    elif "_index" in group:
        row_names = _read_string_dataset(group["_index"])
        df.index = row_names
    return df


def _read_list_group(row_major: bool, group: h5py.Group) -> Any:
    s3class = _get_attr_str(group, "s3class", required=False)
    s4class = _get_attr_str(group, "s4class", required=False)
    data = {name: _read_node(row_major, group[name]) for name in group.keys()}
    if s4class:
        return CustomS4(s4class=s4class, data=data)
    if s3class:
        return CustomS3(s3class=s3class, data=data)
    return data


def _read_node(row_major: bool, node: h5py.Group | h5py.Dataset):
    if isinstance(node, h5py.Dataset):
        return _read_dataset(row_major, node)
    if _is_factor_group(node):
        return _read_factor_group(node)
    return _read_list_group(row_major, node)


def _read_dataset(row_major: bool, dset: h5py.Dataset):
    if dset.dtype.kind in {"S", "O"}:
        return _read_string_dataset(dset)
    if dset.ndim == 2:
        return _read_dense_matrix(row_major, dset)
    data = dset[()]
    if _get_attr_str(dset, "s3class", required=False) == "logical":
        return _read_logical(data)
    return _decode_column(data)


def _read_optional_matrix(row_major: bool, group: h5py.Group, name: str):
    if name not in group:
        return None
    return _read_matrix(row_major, group[name])


def _read_optional_dense(row_major: bool, group: h5py.Group, name: str) -> np.ndarray | None:
    if name not in group:
        return None
    node = group[name]
    if not isinstance(node, h5py.Dataset):
        raise TypeMismatchError(f"Expected dataset at {group.name}/{name}")
    return _read_dense_matrix(row_major, node)


def _read_optional_string(group: h5py.Group, name: str) -> np.ndarray | None:
    if name not in group:
        return None
    node = group[name]
    if not isinstance(node, h5py.Dataset):
        raise TypeMismatchError(f"Expected dataset at {group.name}/{name}")
    return _read_string_dataset(node)


def _read_optional_data_frame(file: h5py.File, row_major: bool, group: h5py.Group, name: str) -> pd.DataFrame | None:
    if name not in group:
        return None
    return _read_data_frame(file, row_major, group[name])


def _read_optional_list(row_major: bool, group: h5py.Group, name: str) -> dict[str, Any] | None:
    if name not in group:
        return None
    node = group[name]
    if not isinstance(node, h5py.Group):
        raise TypeMismatchError(f"Expected group at {group.name}/{name}")
    return _read_list_group(row_major, node)


def _is_factor_group(group: h5py.Group) -> bool:
    return "levels" in group and "values" in group


def _decode_string(value: Any) -> str:
    if isinstance(value, bytes):
        return value.decode("utf-8")
    if isinstance(value, np.ndarray):
        if value.shape == ():
            return _decode_string(value[()])
        if value.size == 1:
            return _decode_string(value[0])
    return str(value)


def _decode_string_list(value: Any) -> list[str]:
    if isinstance(value, np.ndarray):
        return [_decode_string(v) for v in value.tolist()]
    if isinstance(value, (list, tuple)):
        return [_decode_string(v) for v in value]
    return [_decode_string(value)]


def _decode_column(col: Any) -> Any:
    if isinstance(col, np.ndarray) and col.dtype.kind in {"S", "O"}:
        return np.array([_decode_string(v) for v in col])
    if isinstance(col, np.ndarray) and col.shape == ():
        return _decode_string(col[()])
    if isinstance(col, bytes):
        return _decode_string(col)
    return col
