from __future__ import annotations

from itertools import groupby
from pathlib import Path

import pandas as pd

from DAJIN2.utils.midsv_handler import revcomp_midsvs

###########################################################
# group by mutation
###########################################################


def annotate_inversion(midsv_tag: list[str]) -> list[str]:
    return ["@" + tag if tag.islower() else tag for tag in midsv_tag]


def group_by_mutation(midsv_tag: list[str]) -> list[list[str]]:
    return [list(group) for _, group in groupby(midsv_tag, key=lambda x: x[0])]


###########################################################
# Report mutations
###########################################################


def _handle_match(group, genome, start, end, header, chromosome) -> tuple[list[list[str]], int, int]:
    end += len(group)
    start = end
    return [None], start, end


def _handle_substitution(group, genome, start, end, header, chromosome) -> tuple[list[list[str]], int, int]:
    result = []
    for g in group:
        ref = g[1]
        mut = g[2]
        result.append([header, genome, chromosome, start, end, f"substitution: {ref}>{mut}"])
        end += 1
        start = end
    return result, start, end


def _handle_deletion(group, genome, start, end, header, chromosome) -> tuple[list[list[str]], int, int]:
    end += len(group) - 1
    size = len(group)
    seq = "".join([g[-1] for g in group])
    result = [header, genome, chromosome, start, end, f"{size}bp deletion: {seq}"]
    end += 1
    start = end
    return [result], start, end


def _handle_insertion(group, genome, start, end, header, chromosome) -> tuple[list[list[str]], int, int]:
    group = group[0]
    size = group.count("|")
    seq_insertion = "".join([g[-1] for g in group.split("|")[:-1]])
    seq_last = group.split("|")[-1]
    result = []
    result.append([header, genome, chromosome, start, end, f"{size}bp insertion: {seq_insertion}"])
    if seq_last.startswith("="):
        pass
    elif seq_last.startswith("-"):
        result.append([header, genome, chromosome, start, end, f"1bp deletion: {seq_last[-1]}"])
    end += 1
    start = end
    return result, start, end


def _handle_inversion(group, genome, start, end, header, chromosome) -> tuple[list[list[str]], int, int]:
    count_deletion = sum(1 for g in group if "-" in g)
    size = len(group) - count_deletion
    end += size - 1
    seq = "".join([g[-1].upper() for g in group if "-" not in g])
    result = [header, genome, chromosome, start, end, f"{size}bp inversion: {seq}"]
    end += 1
    start = end
    return [result], start, end


def _handle_unknown(group, genome, start, end, header, chromosome) -> tuple[list[list[str]], int, int]:
    end += len(group) - 1
    size = len(group)
    result = [header, genome, chromosome, start, end, f"{size}bp unknown bases"]
    end += 1
    start = end
    return [result], start, end


def report_mutations(midsv_tags_grouped: list[list[str]], genome_coordinates: dict | None, header) -> list[list[str]]:
    if genome_coordinates is None:
        genome_coordinates = {}

    genome = genome_coordinates.get("genome")
    chromosome = genome_coordinates.get("chrom")
    start = end = genome_coordinates.get("start", 0)
    handlers = {
        "=": _handle_match,
        "*": _handle_substitution,
        "-": _handle_deletion,
        "+": _handle_insertion,
        "@": _handle_inversion,
        "N": _handle_unknown,
    }
    results = []
    for group in midsv_tags_grouped:
        for prefix, handler in handlers.items():
            if group[0].startswith(prefix):
                result, start, end = handler(group, genome, start, end, header, chromosome)
                if prefix == "=":
                    continue
                results.extend(result)
    return [list(map(str, r)) for r in results]


###########################################################
# main
###########################################################


def export_to_csv(
    tempdir: Path, sample_name: str, genome_coordinates: dict | None, cons_midsv_tags: dict[str, list[str]]
) -> None:
    if genome_coordinates is None:
        genome_coordinates = {}

    results = []

    for key, cons_midsv_tag in cons_midsv_tags.items():
        header = key.replace("|", "_")
        if genome_coordinates.get("strand") == "-":
            cons_midsv_tag = revcomp_midsvs(cons_midsv_tag)
        cons_midsv_tag_inversion = annotate_inversion(cons_midsv_tag)
        cons_midsv_tag_grouped = group_by_mutation(cons_midsv_tag_inversion)
        result = report_mutations(cons_midsv_tag_grouped, genome_coordinates, header)
        results.extend(result)

    col_names = ["Allele ID", "Genome", "Chromosome", "Start", "End", "Mutation"]
    df_results = pd.DataFrame(results, columns=col_names).sort_values(by=["Allele ID", "Start"])

    path_output = Path(tempdir, "report", "MUTATION_INFO", f"{sample_name}.csv")
    df_results.to_csv(path_output, index=False, encoding="utf-8", lineterminator="\n")
