#! /usr/bin/env python
"""
Build a lowest-common-ancestor database with given taxonomy and genome sigs.
"""

import sys
import csv
import os
from collections import defaultdict

from sourmash import sourmash_args
from sourmash.sourmash_args import load_file_as_signatures
from sourmash.logging import notify, error, debug, set_quiet
from . import lca_utils
from .lca_db import LCA_Database
from sourmash.sourmash_args import DEFAULT_LOAD_K


def load_taxonomy_assignments(
    filename,
    *,
    delimiter=",",
    start_column=2,
    use_headers=True,
    force=False,
    split_identifiers=False,
    keep_identifier_versions=False,
):
    """
    Load a taxonomy assignment spreadsheet into a dictionary.

    The 'assignments' dictionary that's returned maps identifiers to
    lineage tuples.
    """
    from sourmash.tax.tax_utils import LineagePair

    # parse spreadsheet!
    # CTB note: can't easily switch to FileInputCSV, because of
    # janky way we do/don't handle headers here. See issue #2198.
    fp = open(filename, newline="")
    r = csv.reader(fp, delimiter=delimiter)
    row_headers = ["identifiers"]
    row_headers += ["_skip_"] * (start_column - 2)
    row_headers += list(lca_utils.taxlist())

    # first check that headers are interpretable.
    if use_headers:
        notify("examining spreadsheet headers...")
        first_row = next(iter(r))

        n_disagree = 0
        for column, value in zip(row_headers, first_row):
            if column == "_skip_":
                continue

            if column.lower() != value.lower():
                notify(f"** assuming column '{value}' is {column} in spreadsheet")
                n_disagree += 1
                if n_disagree > 2:
                    error("whoa, too many assumptions. are the headers right?")
                    error("expecting {}", ",".join(row_headers))
                    if not force:
                        sys.exit(-1)
                    notify("...continue, because --force was specified.")

    # convert into a lineage pair
    assignments = {}
    num_rows = 0
    n_species = 0
    n_strains = 0
    for row in r:
        if row and row[0].strip():  # want non-empty row
            num_rows += 1
            lineage = list(zip(row_headers, row))
            lineage = [x for x in lineage if x[0] != "_skip_"]

            ident = lineage[0][1]
            lineage = lineage[1:]

            # fold, spindle, and mutilate ident?
            if split_identifiers:
                ident = ident.split(" ")[0]

                if not keep_identifier_versions:
                    ident = ident.split(".")[0]

            # clean lineage of null names, replace with 'unassigned'
            lineage = [(a, lca_utils.filter_null(b)) for (a, b) in lineage]
            lineage = [LineagePair(a, b) for (a, b) in lineage]

            # remove end nulls
            while lineage and lineage[-1].name == "unassigned":
                lineage = lineage[:-1]

            # store lineage tuple
            if lineage:
                # check duplicates
                if ident in assignments:
                    if assignments[ident] != tuple(lineage):
                        if not force:
                            raise Exception(f"multiple lineages for identifier {ident}")
                else:
                    assignments[ident] = tuple(lineage)

                    if lineage[-1].rank == "species":
                        n_species += 1
                    elif lineage[-1].rank == "strain":
                        n_species += 1
                        n_strains += 1

    fp.close()

    # this is to guard against a bug that happened once and I can't find
    # any more, when building a large GTDB-based database :) --CTB
    if len(assignments) * 0.2 > n_species and len(assignments) > 50:
        if not force:
            error("")
            error("ERROR: fewer than 20% of lineages have species-level resolution!?")
            error(
                "({} species assignments found, of {} assignments total)",
                n_species,
                len(assignments),
            )
            error("** If this is intentional, re-run the command with -f.")
            sys.exit(-1)

    return assignments, num_rows


def generate_report(
    record_duplicates,
    record_no_lineage,
    record_remnants,
    unused_lineages,
    unused_identifiers,
    filename,
):
    """
    Output a report of anomalies from building the index.
    """
    with open(filename, "w") as fp:
        print(f"Duplicate signatures: {len(record_duplicates)}", file=fp)
        fp.write("\n".join(record_duplicates))
        fp.write("\n")
        print(f"----\nUnused identifiers: {len(unused_identifiers)}", file=fp)
        fp.write("\n".join(unused_identifiers))
        fp.write("\n")
        print(
            f"----\nNo lineage provided for these identifiers: {len(record_no_lineage)}",
            file=fp,
        )
        fp.write("\n".join(record_no_lineage))
        fp.write("\n")
        print(
            f"----\nNo signatures found for these identifiers: {len(record_remnants)}",
            file=fp,
        )
        fp.write("\n".join(record_remnants))
        fp.write("\n")
        print(f"----\nUnused lineages: {len(unused_lineages)}", file=fp)
        for lineage in unused_lineages:
            fp.write(";".join(lca_utils.zip_lineage(lineage)))
            fp.write("\n")


def index(args):
    """
    main function for building an LCA database.
    """
    if args.start_column < 2:
        error("error, --start-column cannot be less than 2")
        sys.exit(-1)

    set_quiet(args.quiet, args.debug)

    args.scaled = int(args.scaled)

    if args.ksize is None:
        args.ksize = DEFAULT_LOAD_K

    moltype = sourmash_args.calculate_moltype(args, default="DNA")
    picklist = sourmash_args.load_picklist(args)

    db_outfile = args.lca_db_out
    if args.database_format == "json":
        if not (
            db_outfile.endswith(".lca.json") or db_outfile.endswith(".lca.json.gz")
        ):  # logic -> db.save
            db_outfile += ".lca.json"
    else:
        assert args.database_format == "sql"
        if not db_outfile.endswith(".lca.sql"):
            db_outfile += ".lca.sql"

    if os.path.exists(db_outfile):
        error(f"ERROR: output file {db_outfile} already exists. Not overwriting.")
        sys.exit(-1)

    notify(f"saving to LCA DB: {format(db_outfile)}")

    notify(
        f"Building LCA database with ksize={args.ksize} scaled={args.scaled} moltype={moltype}."
    )

    # first, load taxonomy spreadsheet
    delimiter = ","
    if args.tabs:
        delimiter = "\t"
    assignments, num_rows = load_taxonomy_assignments(
        args.csv,
        delimiter=delimiter,
        start_column=args.start_column,
        use_headers=not args.no_headers,
        force=args.force,
        split_identifiers=args.split_identifiers,
        keep_identifier_versions=args.keep_identifier_versions,
    )

    notify(
        f"{len(assignments)} distinct identities in spreadsheet out of {num_rows} rows."
    )
    notify(
        f"{len(set(assignments.values()))} distinct lineages in spreadsheet out of {num_rows} rows."
    )

    db = LCA_Database(args.ksize, args.scaled, moltype)

    inp_files = list(args.signatures)
    if args.from_file:
        more_files = sourmash_args.load_pathlist_from_file(args.from_file)
        inp_files.extend(more_files)

    # track duplicates
    md5_to_name = {}

    #
    # main loop, connecting lineage ID to signature.
    #

    n = 0
    total_n = len(inp_files)
    record_duplicates = set()
    record_no_lineage = []
    record_remnants = set(assignments)
    record_used_lineages = set()
    record_used_idents = set()
    n_skipped = 0
    for filename in inp_files:
        n += 1
        it = load_file_as_signatures(
            filename,
            ksize=args.ksize,
            select_moltype=moltype,
            picklist=picklist,
            yield_all_files=args.force,
        )
        for sig in it:
            notify("\r\033[K", end="")
            notify(
                f"\r... loading signature {str(sig)[:30]} ({n} of {total_n}); skipped {n_skipped} so far",
                end="",
            )
            debug(filename, sig)

            # block off duplicates.
            if sig.md5sum() in md5_to_name:
                debug(
                    "WARNING: in file {}, duplicate md5sum: {}; skipping",
                    filename,
                    sig.md5sum(),
                )
                record_duplicates.add(sig.name)
                continue

            md5_to_name[sig.md5sum()] = str(sig)

            # parse identifier, potentially with splitting
            if sig.name:
                ident = sig.name
            else:
                ident = sig.filename

            orig_ident = ident
            if args.split_identifiers:  # hack for NCBI-style names, etc.
                # split on space...
                ident = ident.split(" ")[0]

                if not args.keep_identifier_versions:
                    # ...and on period.
                    ident = ident.split(".")[0]

            lineage = assignments.get(ident)

            # punt if no lineage and --require-taxonomy
            if lineage is None and args.require_taxonomy:
                if args.fail_on_missing_taxonomy:
                    notify(f"ERROR: no taxonomy found for identifier '{ident}'")
                    if args.split_identifiers:
                        notify(f"(Identifier extracted from name: '{orig_ident})')")
                    sys.exit(-1)
                debug("(skipping, because --require-taxonomy was specified)")
                n_skipped += 1
                continue

            # add the signature into the database.
            try:
                db.insert(sig, ident=ident, lineage=lineage)
            except ValueError as e:
                error(
                    "ERROR: cannot insert signature '{}' (md5 {}, loaded from '{}') into database.",
                    sig,
                    sig.md5sum()[:8],
                    filename,
                )
                error("ERROR: {}", str(e))
                sys.exit(-1)

            if lineage:
                # remove from our list of remaining ident -> lineage
                record_remnants.remove(ident)

                # track ident as used
                record_used_idents.add(ident)
                record_used_lineages.add(lineage)

            # track lineage info - either no lineage, or this lineage used.
            else:
                debug("WARNING: no lineage assignment for {}.", ident)
                record_no_lineage.append(ident)

    # end main add signatures loop

    if n_skipped:
        notify(
            f"... loaded {total_n} signatures; skipped {n_skipped} because of --require-taxonomy."
        )
    else:
        notify(f"... loaded {total_n} signatures.")

    # check -- did we find any signatures?
    if n == 0:
        error("ERROR: no signatures found. ??")
        sys.exit(1)

    # check -- did the signatures we found have any hashes?
    if not db.hashvals:
        error("ERROR: no hash values found - are there any signatures?")
        sys.exit(1)
    notify(
        f"loaded {len(db.hashvals)} hashes at ksize={args.ksize} scaled={args.scaled}"
    )

    if picklist:
        sourmash_args.report_picklist(args, picklist)

    # summarize:
    notify(
        f"{len(record_used_lineages)} assigned lineages out of {len(set(assignments.values()))} distinct lineages in spreadsheet."
    )
    unused_lineages = set(assignments.values()) - record_used_lineages

    notify(
        f"{len(record_used_idents)} identifiers used out of {len(set(assignments))} distinct identifiers in spreadsheet."
    )

    assert record_used_idents.issubset(set(assignments))
    unused_identifiers = set(assignments) - record_used_idents

    # now, save!
    db.save(db_outfile, format=args.database_format)

    ## done!

    # output a record of stuff if requested/available:
    if record_duplicates or record_no_lineage or record_remnants or unused_lineages:
        if record_duplicates:
            notify(f"WARNING: {len(record_duplicates)} duplicate signatures.")
        if record_no_lineage:
            notify(
                f"WARNING: no lineage provided for {len(record_no_lineage)} signatures."
            )
        if record_remnants:
            notify(
                f"WARNING: no signatures for {len(record_remnants)} spreadsheet rows."
            )
        if unused_lineages:
            notify(f"WARNING: {len(unused_lineages)} unused lineages.")

        if unused_identifiers:
            notify(f"WARNING: {len(unused_identifiers)} unused identifiers.")

        if args.report:
            notify(f"generating a report and saving in '{args.report}'")
            generate_report(
                record_duplicates,
                record_no_lineage,
                record_remnants,
                unused_lineages,
                unused_identifiers,
                args.report,
            )
        else:
            notify("(You can use --report to generate a detailed report.)")


if __name__ == "__main__":
    sys.exit(index(sys.argv[1:]))
