from pathlib import Path
from anndata import read_mtx, AnnData
from anndata.utils import make_index_unique
import pandas as pd
import tarfile


def read_10x_mtx(path, var_names = "gene_symbols", make_unique = True, gex_only = True, prefix = None):
    """\
    Read 10x-Genomics-formatted mtx directory.

    Parameters
    ----------
    path
        Path to directory for `.mtx` and `.tsv` files,
        e.g. './filtered_gene_bc_matrices/hg19/'.
    var_names
        The variables index.
    make_unique
        Whether to make the variables index unique by appending '-1',
        '-2' etc. or not.
    gex_only
        Only keep 'Gene Expression' data and ignore other feature types,
        e.g. 'Antibody Capture', 'CRISPR Guide Capture', or 'Custom'
    prefix
        Any prefix before `matrix.mtx`, `genes.tsv` and `barcodes.tsv`. For instance,
        if the files are named `patientA_matrix.mtx`, `patientA_genes.tsv` and
        `patientA_barcodes.tsv` the prefix is `patientA_`.
        (Default: no prefix)

    Returns
    -------
    An :class:`~anndata.AnnData` object
    """
    path = Path(path)
    prefix = "" if prefix is None else prefix
    is_legacy = (path / f"{prefix}genes.tsv").is_file()
    adata = _read_10x_mtx(
        path,
        var_names=var_names,
        make_unique=make_unique,
        prefix=prefix,
        is_legacy=is_legacy,
    )
    if is_legacy or not gex_only:
        return adata
    gex_rows = adata.var["feature_types"] == "Gene Expression"
    return adata[:, gex_rows].copy()


def _read_10x_mtx(path, var_names = "gene_symbols", make_unique = True, cache = False, cache_compression = None, prefix = "", is_legacy = False):
    """
    Read mex from output from Cell Ranger v2- or v3+
    """
    suffix = "" if is_legacy else ".gz"
    adata = read_mtx(path / f"{prefix}matrix.mtx{suffix}").T  # transpose the data
    genes = pd.read_csv(
        path / f"{prefix}{'genes' if is_legacy else 'features'}.tsv{suffix}",
        header=None,
        sep="\t",
    )
    if var_names == "gene_symbols":
        var_names_idx = pd.Index(genes[1].values)
        if make_unique:
            var_names_idx = make_index_unique(var_names_idx)
        adata.var_names = var_names_idx
        adata.var["gene_ids"] = genes[0].values
    elif var_names == "gene_ids":
        adata.var_names = genes[0].values
        adata.var["gene_symbols"] = genes[1].values
    else:
        raise ValueError("`var_names` needs to be 'gene_symbols' or 'gene_ids'")
    if not is_legacy:
        adata.var["feature_types"] = genes[2].values
    barcodes = pd.read_csv(path / f"{prefix}barcodes.tsv{suffix}", header=None)
    adata.obs_names = barcodes[0].values
    return adata


def detect_compression_type(file_path):
    """Detects if a file is gzip or bzip2 based on its header."""
    with open(file_path, 'rb') as f:
        file_start = f.read(2)
        if file_start == b'\x1f\x8b':  # gzip magic number
            return 'gz'
        elif file_start == b'BZ':  # bzip2 magic number
            return 'bz2'
        else:
            raise ValueError("Unknown or unsupported compression format")


def decompress_tar(tar_fn, outdir='.'):                  
    compression_type = detect_compression_type(tar_fn)
    mode = f'r:{compression_type}'
    with tarfile.open(tar_fn, mode) as tar:
        tar.extractall(outdir)
