"""
DCC Extractor module
    This module takes care of extracting data from XML DCC files to  Spark DataFrames.
    There are two kinds of DCC files: experiment and specimen files representing 4 different data
    entities, experiment and line in the experiment files; mouse and embryo in the specimen files.

    The files are expected to be organized by data source in the following directory structure:
    <DCC_XML_PATH>/<DATASOURCE>/*<FILE_TYPE>*.xml
    Each directory containing the raw XML from the DCC.
"""
from typing import List

import py4j
from pyspark.sql import DataFrame, SparkSession
from pyspark.sql.functions import input_file_name, udf, explode_outer, lit, when, col
from pyspark.sql.types import StringType

from impc_etl import logger
from impc_etl.shared.exceptions import UnsupportedFileTypeError, NoDataFoundError


def extract_dcc_xml_files(
    spark_session: SparkSession, dcc_xml_path: str, file_type: str
) -> DataFrame:
    """
    Extracts the DCC XML files into a Spark DataFrame.
    Takes in a SparkSession, the path on the XML files provided by the DCC in the right directory structure,
    a file_type (can be 'specimen' or 'experiment'). Returns a DataFrame containing all the data on the XML files.
    """
    if file_type not in ["experiment", "specimen"]:
        raise UnsupportedFileTypeError

    dcc_xml_path = (
        dcc_xml_path + "/" if not dcc_xml_path.endswith("/") else dcc_xml_path
    )
    # Loading data for 3i, EuroPhenome and Pain Working Group.
    # These projects contain a small amount of data so the files are directly
    # in the project path
    other_projects_path = f"{dcc_xml_path}*/*{file_type}*.xml"

    # Loading IMPC XML files
    # The IMPC directory structure is more nested because it has the potential of containing lots of files
    impc_path = f"{dcc_xml_path}*/*/*/*{file_type}*.xml"

    logger.info(
        f"loading DCC data source from paths '{other_projects_path}' '{impc_path}'"
    )
    try:
        dcc_df = (
            spark_session.read.format("com.databricks.spark.xml")
            .options(rowTag="centre", samplingRatio="1", nullValue="", mode="FAILFAST")
            .load(",".join([other_projects_path, impc_path]))
        )

        logger.info(f"adding _dataSource column")
        dcc_df = dcc_df.withColumn("_sourceFile", lit(input_file_name()))
        data_source_extract = udf(
            lambda x: x.split("/")[-4]
            if x.split("/")[-4] == "impc"
            else x.split("/")[-2],
            StringType(),
        )
        phenotyping_data_status_extract = udf(lambda x: x.split("/")[-3], StringType())
        dcc_df = dcc_df.withColumn("_dataSource", data_source_extract("_sourceFile"))
        dcc_df = dcc_df.withColumn(
            "_sourcePhenotypingStatus",
            when(
                col("_dataSource") == "impc",
                phenotyping_data_status_extract("_sourceFile"),
            ).otherwise(lit(None)),
        )

        logger.info(
            f"finished load of DCC data source from path '{other_projects_path}' '{impc_path}'"
        )
    except py4j.protocol.Py4JJavaError as e:
        if "InvalidInputException" in str(e):
            raise FileNotFoundError
        else:
            raise e
    return dcc_df


def get_entity_by_type(
    dcc_df: DataFrame, entity_type: str, centre_columns: List[str]
) -> DataFrame:
    """
    Takes a DCC DataFrame and obtains a given the entity type
    adding a '_type' column to the DataFrame, Takes in a DataFrame generated by
    `impc_etl.jobs.extract.dcc_extractor_helper.extract_dcc_xml_files`, an entity_type
    (can be 'line', 'experiment', 'mouse' or 'embryo'), the list of centre_columns
    that the output DataFrame should maintain and returns A DataFrame containing only the data for the
    specified entity type.
    """
    if entity_type not in dcc_df.columns:
        print("No entries for " + entity_type + " in the provided DCC XML files")
        raise NoDataFoundError
    centre_columns.append(dcc_df[entity_type])
    entity_df = dcc_df.where(dcc_df[entity_type].isNotNull()).select(
        centre_columns + [entity_type]
    )
    entity_df = (
        entity_df.withColumn("tmp", explode_outer(entity_df[entity_type]))
        .select(["tmp.*"] + centre_columns)
        .withColumn("_type", lit(entity_type))
        .drop(entity_type)
    )
    return entity_df
