from functools import partial
from ...utilities import default_for_serialisation
import rapidjson as json
dumper = partial(json.dumps, default=default_for_serialisation)
from ...parsers.blast_utils import BlastOpener
from .xml_utils import _get_query_for_blast, _get_target_for_blast
from xml.parsers.expat import ExpatError
import xml
from ...utilities.dbutils import connect
from sqlalchemy.orm.session import Session
import msgpack
import tempfile
import os
from .query import Query
from .target import Target
from .hit import prepare_hit, InvalidHit
from .xml_utils import get_multipliers, get_off_by_one
from .utils import load_into_db
import multiprocessing as mp
from ...utilities.log_utils import create_null_logger
import pandas as pd
import zlib
import struct


struct_row = struct.Struct(">LLL")


def xml_pickler(configuration, filename, default_header,
                cache=None,
                max_target_seqs=10):
    valid, _, exc = BlastOpener(filename).sniff(default_header=default_header)
    engine = connect(configuration)
    session = Session(bind=engine)

    if not valid:
        err = "Invalid BLAST file: %s" % filename
        raise TypeError(err)
    directory = os.path.dirname(filename)
    with tempfile.NamedTemporaryFile(suffix=".db", dir=directory, mode="wb", delete=False) as dbhandle:
        dbname = dbhandle.name
        if not isinstance(cache, dict) or set(cache.keys()) != {"query", "target"}:
            cache = dict()
            cache["query"] = dict((item.query_name, item.query_id) for item in session.query(Query))
            cache["target"] = dict((item.target_name, item.target_id) for item in session.query(Target))

        rows = b""
        try:
            with BlastOpener(filename) as opened:
                try:
                    qmult, tmult = None, None
                    for query_counter, record in enumerate(opened, start=1):
                        if qmult is None:
                            qmult, tmult = get_multipliers(record)
                            off_by_one = get_off_by_one(record)
                        hits, hsps, cache = objectify_record(
                            record, [], [], cache,
                            qmult=qmult, tmult=tmult, off_by_one=off_by_one)

                        record = {"hits": hits, "hsps": hsps}
                        jrecord = dumper(record)
                        write_start = dbhandle.tell()
                        write_length = dbhandle.write(msgpack.dumps(jrecord))
                        rows += struct_row.pack(query_counter, write_start, write_length)

                except ExpatError as err:
                    raise ExpatError("{} is an invalid BLAST file, sending back anything salvageable.\n{}".format(filename,
                                                                                                                  err))
        except xml.etree.ElementTree.ParseError as err:
            # logger.error("%s is an invalid BLAST file, sending back anything salvageable", filename)
            raise xml.etree.ElementTree.ParseError(
                "{} is an invalid BLAST file, sending back anything salvageable.\n{}".format(filename, err))
        except ValueError as err:
            # logger.error("Invalid BLAST entry")
            raise ValueError(
                "{} is an invalid BLAST file, sending back anything salvageable.\n{}".format(filename, err))

        return dbname, zlib.compress(rows)


def _serialise_xmls(self):
    # Load sequences in DB, precache IDs

    if isinstance(self.xml, str):
        self.xml = [self.xml]
    else:
        assert isinstance(self.xml, (list, set))

    hits, hsps = [], []
    hit_counter, record_counter = 0, 0

    for filename in self.xml:
        valid, header, exc = BlastOpener(filename).sniff()
        if valid is True:
            self.header = header
        else:
            self.logger.error(exc)
            self.xml.remove(filename)

    cache = {"query": self.queries, "target": self.targets}
    if self._blast_loading_debug is False and (self.single_thread is True or self.procs == 1):
        for filename in self.xml:
            valid, _, exc = BlastOpener(filename).sniff(default_header=self.header)
            if not valid:
                self.logger.error(exc)
                continue
            try:
                self.logger.debug("Analysing %s", filename)
                qmult, tmult = None, None
                with BlastOpener(filename) as opened:
                    for record in opened:
                        if qmult is None:
                            qmult, tmult = get_multipliers(record)
                            off_by_one = get_off_by_one(record)
                        record_counter += 1
                        if record_counter > 0 and record_counter % 10000 == 0:
                            self.logger.info("Parsed %d queries", record_counter)
                        current = len(hits)
                        hits, hsps, cache = objectify_record(
                            record, hits, hsps,
                            qmult=qmult,
                            tmult=tmult,
                            cache=cache,
                            logger=self.logger, off_by_one=off_by_one)
                        hit_counter += len(hits) - current
                        hits, hsps = load_into_db(self, hits, hsps, force=False)
                self.logger.debug("Finished %s", filename)
            except ExpatError:
                self.logger.error("%s is an invalid BLAST file, saving what's available", filename)
        _, _ = load_into_db(self, hits, hsps, force=True)
    elif self._blast_loading_debug is True or self.procs > 1:
        self.logger.warning("Creating a pool with %d processes", min(self.procs, len(self.xml)))
        results = []
        if self._blast_loading_debug is True:
            for num, filename in enumerate(self.xml):
                results.append(xml_pickler(self.configuration,
                                           filename, self.header,
                                           cache=None,
                                           max_target_seqs=self._max_target_seqs))
        else:
            pool = mp.Pool(self.procs)
            for num, filename in enumerate(self.xml):
                args = (self.configuration, filename, self.header)
                kwds = {"max_target_seqs": self._max_target_seqs, "cache": None}
                pool.apply_async(xml_pickler, args=args, kwds=kwds, callback=results.append)
            pool.close()
            pool.join()

        for dbfile, rows in results:
            with open(dbfile, "rb") as dbhandle:
                for query_counter, write_start, write_length in struct_row.iter_unpack(zlib.decompress(rows)):
                    dbhandle.seek(write_start)
                    result = json.loads(msgpack.loads(dbhandle.read(write_length), raw=False, use_list=True))
                    __hits = result["hits"]
                    __hsps = result["hsps"]
                    assert (not __hsps and not __hits) or __hsps != __hits, (__hits, __hsps)
                    hit_counter += len(__hits)
                    hits.extend(__hits)
                    hsps.extend(__hsps)
                    hits, hsps = load_into_db(self, hits, hsps, force=False)
                    if record_counter > 0 and record_counter % 10000 == 0:
                        self.logger.debug("Parsed %d queries", record_counter)
            os.remove(dbfile)

        self.logger.debug("Finished sending off the data for serialisation")
        _, _ = load_into_db(self, hits, hsps, force=True)

    self.logger.info("Loaded %d alignments for %d queries",
                     hit_counter, record_counter)

    self.logger.info("Finished loading blast hits")


def objectify_record(record, hits, hsps, cache,
                     logger=create_null_logger(),
                     qmult=1, tmult=1, off_by_one=False):
    """
    Private method to serialise a single record into the DB.

    :param record: The BLAST record to load into the DB.
    :type record: Bio.SearchIO.QueryResult
    :param hits: Cache of hits to load into the DB.
    :type hits: list

    :param hsps: Cache of hsps to load into the DB.
    :type hsps: list

    :param cache: cache of match query_name/query_id
    :type cache: dict

    :param logger: logger
    :type logger: logging.Logger

    :param qmult: query multiplier (3 for BLASTX or TBLASTX, 1 otherwise)
    :type qmult: int

    :param qmult: query multiplier (1 for TBLASTX or TBLASTN, 1 otherwise)
    :type qmult: int

    :param off_by_one: DIAMOND before 0.9.31 has a bug which causes the query end to be off by 1.
    :type off_by_one: bool

    :returns: hits, hsps, cache
    :rtype: (list, list, dict)
    """

    if len(record.hits) == 0:
        return hits, hsps, cache

    current_query, name, cache["query"] = _get_query_for_blast(record, cache["query"])
    alignments = dict((ccc, (hit.id, max(_.bitscore for _ in hit.hsps))) for ccc, hit in enumerate(record.hits))
    ids, scores = list(zip(*alignments.values()))
    phits = pd.DataFrame({"idx": list(alignments.keys()), "target_name": ids, "bitscore": scores})
    for hit_num, t in enumerate(phits.sort_values(["bitscore", "target_name"],
                                                  ascending=[False, True]).itertuples(name=None)):
        idx, target_name, bitscore = t[1:]
        alignment = record.hits[idx]
        logger.debug("Started the hit %s vs. %s", name, record.hits[idx].id)
        current_target, cache["target"] = _get_target_for_blast(alignment, cache["target"])
        hit_dict_params = dict()
        (hit_dict_params["query_multiplier"],
         hit_dict_params["target_multiplier"]) = (qmult, tmult)
        hit_evalue = min(_.evalue for _ in record.hits[idx].hsps)
        hit_bs = bitscore
        hit_dict_params["hit_number"] = hit_num + 1
        hit_dict_params["evalue"] = hit_evalue
        hit_dict_params["bits"] = hit_bs
        try:
            hit, hit_hsps = prepare_hit(alignment, current_query,
                                        current_target,
                                        query_length=record.seq_len,
                                        off_by_one=off_by_one,
                                        **hit_dict_params)
        except InvalidHit as exc:
            logger.error(exc)
            continue
        hit["query_aligned_length"] = min(record.seq_len, hit["query_aligned_length"])
        hits.append(hit)
        hsps.extend(hit_hsps)

    return hits, hsps, cache
