import functools
import multiprocessing

from .. import Devutils as dev
import numpy as np
import hashlib

__all__ = [
    "SMILESSupplier",
    "consume_smiles_supplier",
    "match_smiles_supplier"
]

class SMILESSupplier:
    def __init__(self, smiles_file, line_indices=None, name=None,
                 size=int(1e3),
                 split_idx=0,
                 split_char=None,
                 line_parser=None):
        self.name = name
        self.smi = dev.StreamInterface(smiles_file, mode='rb')
        self.line_indices = line_indices
        self._size = size
        self._call_depth = 0
        self._stream = None
        self._cur = None
        self._max_offset = None
        self._offsets:np.ndarray[(None,), int] = None
        self._flexible_offsets = None
        self._assignable_offsets = None
        self._encoding = self.smi.get_encoding()
        self._binary = self.smi.is_binary()
        if isinstance(split_char, str) and self._binary:
            split_char = split_char.encode(self._encoding)
        self.split_char = split_char
        self.split_idx = split_idx
        self._line_parser = line_parser
        self.line_parser = line_parser

    known_suppliers = {
        "zinc20": {
            'smiles_file': 'ZINC20.smi',
            'line_indices': 'ZINC20_idx.npy',
            'split_idx': 0
        },
        "emols": {
            'smiles_file': 'emolecule_sc_2026_01_01.smi',
            'line_indices': 'molecule_sc_2026_01_01_idx.npy',
            'split_idx': 0
        },
        "pubchem":{
            'smiles_file':'pubchem_cid_smi_2026_01.smi',
            'line_indices':'pubchem_cid_smi_2026_01_idx.npy',
            'split_idx':1
        }
    }
    @classmethod
    def from_name(cls, name):
        return cls(**cls.known_suppliers[name], name=name)

    def to_mp_state(self):
        return (
            self.smi._input,
            self.name,
            self.split_idx,
            self.split_char,
            self._line_parser
        )
    @classmethod
    def from_mp_state(cls, state, line_indices=None, **extra):
        smi, name, split_idx, split_char, line_parser = state
        kwargs = {}
        kwargs['line_indices'] = line_indices
        kwargs['name'] = name
        kwargs['split_idx'] = split_idx
        kwargs['split_char'] = split_char
        kwargs['line_parser'] = line_parser
        kwargs.update(extra)

        return cls(
            smi,
            **kwargs
        )

    def __enter__(self):
        self._call_depth += 1
        if self._call_depth == 1:
            if self.line_parser is None:
                self.line_parser = self._default_parser
            self._stream = self.smi.__enter__()
            self._cur = 0
            if self.line_indices is None:
                #TODO: add an array offset style index to this
                #      in case the SMI file is too big for even uint64
                self._max_offset = 0
                self.line_indices = np.zeros(self._size, dtype='uint64')
            if isinstance(self.line_indices, np.ndarray):
                self._offsets = self.line_indices
                self._flexible_offsets = True
                self._assignable_offsets = True
            else:
                self._offsets = np.load(self.line_indices, mmap_mode='r')
                self._flexible_offsets = False
                self._assignable_offsets = False
            if self._max_offset is None:
                self._max_offset = len(self._offsets) - 1
        return self._stream
    def __exit__(self, exc_type, exc_val, exc_tb):
        self._call_depth -= 1
        if self._call_depth == 0:
            self.line_parser = self._line_parser
            self._cur = None
            self._stream.__exit__(exc_type, exc_val, exc_tb)
            self._stream = None
            if self._flexible_offsets:
                self.line_indices = self._offsets
            self._offsets = None

    def __len__(self):
        with self:
            if self._flexible_offsets:
                self.create_line_index(return_index=False)
                return self._max_offset
            else:
                return self._max_offset
    @classmethod
    def _split_smi(cls, line:bytes, encoding='utf-8', split_idx=0, split_char=None, maxsplit=-1):
        if split_char is None:
            return line.split(maxsplit=maxsplit)[split_idx].strip().decode(encoding)
        else:
            return line.split(sep=split_char, maxsplit=maxsplit)[split_idx].strip().decode(encoding)
    def _default_parser(self, line):
        return self._split_smi(line, split_idx=self.split_idx, split_char=self.split_char, encoding=self._encoding)

    @classmethod
    def _consume_next(self, db, parser):
        line = db.readline()
        if len(line) == 0:
            return line
        else:
            return parser(line)
    def find_smi(self, n, block_size=None):
        with self as db:
            #TODO: add ability to stream line indices to avoid reading them into memory
            if n >= self._max_offset:
                self.create_line_index(n, return_index=False)
            db.seek(self._offsets[n])
            if block_size is None:
                self._cur = n
                return self._consume_next(db, self.line_parser)
            else:
                self._cur = n + block_size
                if n + block_size >= self._max_offset:
                    blocks = []
                    for m in range(block_size):
                        if self._assignable_offsets:
                            self._expand_offset_if_needed(n+m)
                            self._offsets[n+m] = db.tell()
                        blocks.append(self._consume_next(db, self.line_parser))
                    return blocks
                else:
                    return [
                        self._consume_next(db, self.line_parser)
                        for _ in range(block_size)
                    ]
    def consume_iter(self, start_at=None, upto=None):
        with self as db:
            if start_at is None:
                start_at = self._cur
            else:
                self.create_line_index(upto=start_at, return_index=False)
            ninds = start_at
            try:
                db.seek(self._offsets[ninds])
                if upto is None:
                    smi = self._consume_next(db, self.line_parser)
                    while len(smi) > 0:
                        ninds += 1
                        if self._assignable_offsets:
                            self._expand_offset_if_needed(ninds)
                            self._offsets[ninds] = db.tell()
                        yield smi
                        smi = self._consume_next(db, self.line_parser)
                else:
                    if ninds >= upto: return
                    smi = self._consume_next(db, self.line_parser)
                    while ninds < upto and len(smi) > 0:
                        ninds += 1
                        if self._assignable_offsets:
                            self._expand_offset_if_needed(ninds)
                            self._offsets[ninds] = db.tell()
                        yield smi
                        smi = self._consume_next(db, self.line_parser)
            finally:
                self._cur = ninds
                self._max_offset = max(self._max_offset, ninds)

    def __next__(self):
        if self._stream is None:
            cls = type(self)
            raise ValueError(f"{cls.__name__} must be opened via `with` before it can be iterated over")
        with self as db:
            db.seek(self._offsets[self._cur])
            return self._consume_next(db, self.line_parser)

    def __iter__(self):
        return self.consume_iter()

    def _expand_offset_if_needed(self, n):
        if n >= len(self._offsets):
            if not self._flexible_offsets:
                raise ValueError(f"{self._max_offset} `line_indices` were passed, but db extends beyond that")
            else:
                new_offsets = np.zeros(2 * len(self._offsets), dtype='uint64')
                new_offsets[:len(self._offsets)] = self._offsets
                self._offsets = new_offsets
    def create_line_index(self, upto=None, return_index=True):
        with self as db:
            if not self._assignable_offsets:
                if return_index:
                    return self._offsets[:self._max_offset]
                else:
                    return None
            ninds = self._max_offset
            db.seek(self._offsets[ninds])
            try:
                if upto is None:
                    l = len(db.readline())
                    while l > 0:
                        ninds += 1
                        self._expand_offset_if_needed(ninds)
                        self._offsets[ninds] = self._offsets[ninds-1] + l
                        if not self._binary:
                            if db.peek(1) == "\r": self._offsets[ninds] += 1
                        l = len(db.readline())
                else:
                    if ninds < upto:
                        l = len(db.readline())
                        while ninds < upto and l > 0:
                            ninds += 1
                            self._expand_offset_if_needed(ninds)
                            self._offsets[ninds] = self._offsets[ninds-1] + l
                            if not self._binary:
                                if db.peek(1) == "\r": self._offsets[ninds] += 1
                            # self._offsets[ninds] = db.tell()
                            l = len(db.readline())
            finally:
                self._max_offset = ninds
            if return_index:
                return self._offsets[:self._max_offset]
            else:
                return None
    @classmethod
    def save_line_index(cls, file, line_index):
        max_offset = line_index[-1]
        for dtype in [np.uint16, np.uint32, np.uint64]:
            if max_offset < np.iinfo(dtype).max:
                line_index = line_index.astype(dtype)
                break
        return np.save(file, line_index)

def _consume_supplier_mp(state, consumer, line_offset, block_size):
    offsets = np.full(block_size, line_offset)
    supplier = SMILESSupplier.from_mp_state(state, line_indices=offsets)
    res = []
    with supplier:
        for smi in supplier.consume_iter(start_at=0, upto=block_size):
            subres = consumer(smi)
            if subres is not None: res.append(subres)
    return res

def consume_smiles_supplier(supplier:SMILESSupplier, consumer, pool=None, start_at=None, upto=None, initializer=None):
    if pool is True:
        pool = multiprocessing.Pool(initializer=initializer)
    elif dev.is_int(pool):
        pool = multiprocessing.Pool(pool, initializer=initializer)

    if pool is None:
        res = []
        with supplier:
            for smi in supplier.consume_iter(start_at=start_at, upto=upto):
                subres = consumer(smi)
                if subres is not None: res.append(subres)

        return res
    else:
        with supplier:
            max_size = len(supplier) if upto is None else upto
            offsets = supplier._offsets # TODO: gross
            nproc = pool._processes
            block_size = max_size // nproc
            num_blocks = int(np.ceil(max_size / block_size))
            block_starts_sizes = [
                (offsets[block_size*i], min([block_size, max_size - (block_size*i+1) + 1]))
                for i in range(num_blocks)
            ]
            #TODO: don't access the `_input` argument directly...
            state = supplier.to_mp_state()
            args = [(state, consumer) + bs for bs in block_starts_sizes]
        res = pool.starmap(_consume_supplier_mp, args)

        return sum(res, [])

def _match_rdkit(matcher, smi):
    from rdkit.Chem import AllChem
    mol = AllChem.MolFromSmiles(smi)
    if mol is None: return None
    if mol.GetSubstructMatch(matcher): return smi

def _disable_rdkit_log(blockage=[]):
    from rdkit.rdBase import BlockLogs
    bl = BlockLogs()
    blockage.append([bl,  bl.__enter__()])

def match_smiles_supplier(supplier:SMILESSupplier, matcher, pool=None,
                          start_at=None, upto=None, quiet=True,
                          out_file=None,
                          initializer=None):
    from rdkit.rdBase import BlockLogs
    smarts_tag = str(matcher)
    if isinstance(matcher, str):
        from .RDKit import RDKitInterface
        AllChem = RDKitInterface.submodule("Chem.AllChem")
        smarts_candidate = AllChem.MolFromSmarts(matcher)
        matcher = functools.partial(_match_rdkit, smarts_candidate)
        if quiet and initializer is None:
            initializer = _disable_rdkit_log

    if out_file is True:
        out_file_bits = [
            "match",
            "{name}" if supplier.name is not None else None,
            "s{start}" if start_at is not None else None,
            "t{to}" if upto is not None else None,
            "{smarts_key}.smi"
        ]
        out_file = "_".join(b for b in out_file_bits if b is not None)

    hash_obj = hashlib.md5(smarts_tag.encode('utf-8'))
    smarts_key = hash_obj.hexdigest()
    if isinstance(out_file, str):
        out_file = out_file.format(
            name=supplier.name,
            start=start_at,
            to=upto,
            smarts_key=smarts_key,
            smarts_tag=smarts_tag)

    if quiet:
        with BlockLogs():
            matches = consume_smiles_supplier(supplier, matcher, pool=pool, start_at=start_at, upto=upto, initializer=initializer)
    else:
        matches = consume_smiles_supplier(supplier, matcher, pool=pool, start_at=start_at, upto=upto, initializer=initializer)

    if out_file is not None:
        with dev.StreamInterface(out_file, mode='w+') as match_output:
            match_output.write(f"SMARTS: {smarts_tag} ({smarts_key})\n")
            for match in matches:
                match_output.write(match + "\n")

    return matches
