import time
from pathlib import Path

import polars as pl
from phenopackets import Family, Phenopacket

from pheval.utils.file_utils import all_files
from pheval.utils.logger import get_logger
from pheval.utils.phenopacket_utils import (
    GeneIdentifierUpdater,
    PhenopacketRebuilder,
    PhenopacketUtil,
    create_gene_identifier_map,
    phenopacket_reader,
    write_phenopacket,
)

logger = get_logger()


def update_outdated_gene_context(
    phenopacket_path: Path, gene_identifier: str, identifier_map: pl.DataFrame
) -> Phenopacket | Family:
    """
    Update the gene context of the Phenopacket.

    Args:
        phenopacket_path (Path): The path to the Phenopacket file.
        gene_identifier (str): Identifier to update the gene context.
        identifier_map (pl.DataFrame): The gene identifier map used for updating.

    Returns:
        Union[Phenopacket, Family]: The updated Phenopacket or Family.
    Notes:
        This function updates the gene context within the Phenopacket or Family instance.
        The gene_identifier parameter should be chosen from ensembl_id, hgnc_id, or entrez_id
        to update to the current gene identifier in the Phenopacket. We recommend using the ENSEMBL namespace
        to describe the gene identifiers.
    """
    phenopacket = phenopacket_reader(phenopacket_path)
    interpretations = PhenopacketUtil(phenopacket).interpretations()
    updated_interpretations = GeneIdentifierUpdater(
        identifier_map=identifier_map, gene_identifier=gene_identifier
    ).update_genomic_interpretations_gene_identifier(interpretations, phenopacket_path)
    return PhenopacketRebuilder(phenopacket).update_interpretations(updated_interpretations)


def create_updated_phenopacket(
    gene_identifier: str,
    phenopacket_path: Path,
    output_dir: Path,
    identifier_map: pl.DataFrame = None,
) -> None:
    """
    Update the gene context within the interpretations for a Phenopacket and writes the updated Phenopacket.

    Args:
        gene_identifier (str): Identifier used to update the gene context.
        phenopacket_path (Path): The path to the input Phenopacket file.
        output_dir (Path): The directory where the updated Phenopacket will be written.
        identifier_map (pl.DataFrame): The gene identifier map used for updating.
    Notes:
        The gene_identifier parameter should be chosen from ensembl_id, hgnc_id, or entrez_id
        to update to the current gene identifier in the Phenopacket. We recommend using the ENSEMBL namespace
        to describe the gene identifiers.
    """
    identifier_map = create_gene_identifier_map() if identifier_map is None else identifier_map
    updated_phenopacket = update_outdated_gene_context(phenopacket_path, gene_identifier, identifier_map)
    write_phenopacket(updated_phenopacket, output_dir.joinpath(phenopacket_path.name))


def create_updated_phenopackets(gene_identifier: str, phenopacket_dir: Path, output_dir: Path) -> None:
    """
    Update the gene context within the interpretations for a directory of Phenopackets
    and writes the updated Phenopackets.

    Args:
        gene_identifier (str): Identifier used to update the gene context.
        phenopacket_dir (Path): The path to the input Phenopacket directory.
        output_dir (Path): The directory where the updated Phenopackets will be written.
    Notes:
        The gene_identifier parameter should be chosen from ensembl_id, hgnc_id, or entrez_id
        to update to the current gene identifier in the Phenopacket. We recommend using the ENSEMBL namespace
        to describe the gene identifiers.
    """
    identifier_map = create_gene_identifier_map()
    for phenopacket_path in all_files(phenopacket_dir):
        logger.info(f"Updating gene context for: {phenopacket_path.name}")
        updated_phenopacket = update_outdated_gene_context(phenopacket_path, gene_identifier, identifier_map)
        write_phenopacket(updated_phenopacket, output_dir.joinpath(phenopacket_path.name))


def update_phenopackets(gene_identifier: str, phenopacket_path: Path, phenopacket_dir: Path, output_dir: Path) -> None:
    """
    Update the gene identifiers in either a single phenopacket or a directory of phenopackets.

    Args:
        gene_identifier (str): The gene identifier to be updated.
        phenopacket_path (Path): The path to a single Phenopacket file.
        phenopacket_dir (Path): The directory containing multiple Phenopacket files.
        output_dir (Path): The output directory to save the updated Phenopacket files.
    Notes:
        The gene_identifier parameter should be chosen from ensembl_id, hgnc_id, or entrez_id
        to update to the current gene identifier in the Phenopacket. We recommend using the ENSEMBL namespace
        to describe the gene identifiers.
    """
    start_time = time.perf_counter()
    logger.info("Updating phenopackets.")
    output_dir.mkdir(exist_ok=True)
    logger.info(f"Created directory {output_dir}.")
    logger.info(f"Gene identifier set to: {gene_identifier}.")
    if phenopacket_path is not None:
        logger.info(f"Updating {phenopacket_path}.")
        create_updated_phenopacket(gene_identifier, phenopacket_path, output_dir)
    elif phenopacket_dir is not None:
        logger.info(f"Updating {len(all_files(phenopacket_dir))} phenopackets in {phenopacket_dir}.")
        create_updated_phenopackets(gene_identifier, phenopacket_dir, output_dir)
    logger.info(f"Updating finished! Total time: {time.perf_counter() - start_time:.2f} seconds.")
