import pandas as pd
from floria_strainer import logger

def parse_haplosets(filename: str, hapq_cut: int = 15) -> tuple[dict, dict]:
    """
    Read a haplosets file and return a dictionary with the values, as well as reads within a haploset with a HAPQ > threshold.

    Parameters
    ----------
    filename : str
        The name of the file to read.

    Returns
    -------
    dict
        A dictionary with the haploset metadata.
    dict
        {
            contig (str): {
                read_name(str): haploset(str)
            }
        }
        A dictionary with the reads withing each haploset passing threshold
    """
    haplosets = {"contig": [], "haploset": [], "HAPQ": []}
    haqp_sup = False
    read_dict = {}
    with open(filename, "r") as f:
        logger.info(f"Reading the haplosets file {filename}")
        for line in f:
            if line.startswith(">"):
                lsplit = line.rstrip()[1:].split("\t")
                hapq = int(lsplit[6].split(":")[1])
                haploset = lsplit[0].split(".")[0]
                contig = lsplit[1].split(":")[1]
                haplosets["contig"].append(contig)
                haplosets["haploset"].append(haploset)
                haplosets["HAPQ"].append(hapq)
                if hapq >= hapq_cut:
                    haqp_sup = True
                    read_dict.setdefault(contig, {})
                else:
                    haqp_sup = False
            elif haqp_sup:
                read_dict[contig][line.split()[0]] = haploset
    return haplosets, read_dict


def parse_vartig_info(filename: str) -> dict:
    """
    Read a vartig_info.txt file and return a dictionary with the values.

    Parameters
    ----------
    filename : str
        The name of the file to read.

    Returns
    -------
    dict
        A dictionary with the values in the file.
    """

    vartig_all = {
        "contig": [],
        "haploset": [],
        "pos": [],
        "cons": [],
        "allele": [],
        "support": [],
    }
    with open(filename, "r") as f:
        logger.info(f"Reading the vartig_info file {filename}")
        for line in f:
            if line.startswith(">"):
                lsplit = line.rstrip()[1:].split("\t")
                hpset = lsplit[0].split(".")[0]
                contig = lsplit[0].split("/")[-1]
                continue
            else:
                try:
                    lsplit = line.split("\t")
                    pos = int(lsplit[0].split(":")[1])
                    cons = lsplit[1]
                    als = lsplit[2].split("|")
                    if cons != "?":
                        for all in als:
                            al, sup = all.split(":")
                            vartig_all["contig"].append(contig)
                            vartig_all["haploset"].append(hpset)
                            vartig_all["pos"].append(int(pos))
                            vartig_all["cons"].append(int(cons))
                            vartig_all["allele"].append(int(al))
                            vartig_all["support"].append(int(sup))
                except ValueError:
                    continue
                continue

    return vartig_all


def parse_floria_contig_ploidy(contig_ploidy: str) -> pd.DataFrame:
    """
    Read a contig_ploidy_info.tsv file generated by Floria and returns it as a DataFrame

    Parameters
    ----------
    contig_ploidy : str
        The name of the file to read.

    Returns
    -------
    pd.DataFrame
        A pandas DataFrame with the values of the file.
    """

    df = pd.read_csv(contig_ploidy, sep="\t")
    return df
