"""
`deriva_ml_base.py` is the core module for the Deriva ML project.  This module implements the DerivaML class, which is
the primary interface to the Deriva based catalogs.  The module also implements the Feature and Vocabulary functions
in the DerivaML.

DerivaML and its associated classes all depend on a catalog that implements a `deriva-ml` schema with tables and
relationships that follow a specific data model.

"""

import getpass
import logging
from datetime import datetime
from itertools import chain
from pathlib import Path
from typing import Optional, Any, Iterable, TYPE_CHECKING
from deriva.core import (
    ErmrestCatalog,
    get_credential,
    urlquote,
    DEFAULT_SESSION_CONFIG,
)
from deriva.core.datapath import DataPathException, _ResultSet as ResultSet
from deriva.core.datapath import _CatalogWrapper
from deriva.core.ermrest_catalog import ResolveRidResult
from deriva.core.ermrest_model import FindAssociationResult, Key, Table
from deriva.core.hatrac_store import HatracStore
from pydantic import validate_call, ConfigDict

from .execution_configuration import ExecutionConfiguration
from .feature import Feature, FeatureRecord
from .dataset import Dataset
from .deriva_model import DerivaModel
from .upload import (
    table_path,
    execution_rids,
    execution_metadata_dir,
    upload_directory,
    UploadAssetDirectory,
)
from .deriva_definitions import ColumnDefinition
from .deriva_definitions import ExecMetadataVocab
from .deriva_definitions import (
    RID,
    Status,
    FileUploadState,
    DerivaMLException,
    ML_SCHEMA,
    VocabularyTerm,
)

if TYPE_CHECKING:
    from .execution import Execution


class DerivaML(Dataset):
    """Base class for ML operations on a Deriva catalog.

    This class is intended to be used as a base class on which more domain specific interfaces are built.

    Attributes:
        host_name: Hostname of the Deriva server.
        catalog_id: Catalog ID. Either and identifier, or a catalog name.
        domain_schema: Schema name for domain specific tables and relationships.
        model: ERMRest model for the catalog
    """

    def __init__(
        self,
        hostname: str,
        catalog_id: str | int,
        domain_schema: str = None,
        project_name: str = None,
        cache_dir: Optional[str] = None,
        working_dir: Optional[str] = None,
        model_version: str = "1",
        ml_schema: str = ML_SCHEMA,
        logging_level=logging.INFO,
    ):
        """Create and initialize a DerivaML instance.

        This method will connect to a catalog, and initialize local configuration for the ML execution.
        This class is intended to be used as a base class on which domain-specific interfaces are built.

        Args:
            hostname: Hostname of the Deriva server.
            catalog_id: Catalog ID. Either and identifier, or a catalog name.
            domain_schema: Schema name for domain specific tables and relationships.
            cache_dir: Directory path for caching data downloaded from the Deriva server as bdbag.
            working_dir: Directory path for storing data used by or generated by any computations.
            model_version: A string that indicates the version model.  Typically passed in via
        """
        self.credential = get_credential(hostname)
        self.catalog = ErmrestCatalog(
            "https",
            hostname,
            catalog_id,
            self.credential,
            session_config=self._get_session_config(),
        )
        self.model = DerivaModel(
            self.catalog.getCatalogModel(), domain_schema=domain_schema
        )

        default_workdir = self.__class__.__name__ + "_working"
        self.working_dir = (
            Path(working_dir) / getpass.getuser()
            if working_dir
            else Path.home() / "deriva-ml"
        ) / default_workdir
        self.working_dir.mkdir(parents=True, exist_ok=True)
        self.cache_dir = (
            Path(cache_dir) if cache_dir else Path.home() / "deriva-ml" / "cache"
        )
        self.cache_dir.mkdir(parents=True, exist_ok=True)

        # Initialize dataset class.
        super().__init__(self.model, self.cache_dir)

        self.host_name = hostname
        self.catalog_id = catalog_id
        self.ml_schema = ml_schema
        self.version = model_version
        self.configuration = None

        self.domain_schema = self.model.domain_schema
        self.project_name = project_name or self.domain_schema

        self.start_time = datetime.now()
        self.status = Status.pending.value
        self._logger = logging.getLogger("deriva_ml")
        self._logger.setLevel(logging_level)

        logging.basicConfig(
            level=logging_level,
            format="%(asctime)s - %(name)s.%(levelname)s - %(message)s",
        )

        # Set logging level for dervia library
        deriva_logger = logging.getLogger("deriva")
        deriva_logger.setLevel(logging_level)

        if "dirty" in self.version:
            logging.info(
                f"Loading dirty model.  Consider commiting and tagging: {self.version}"
            )

    @staticmethod
    def _get_session_config():
        """ """
        session_config = DEFAULT_SESSION_CONFIG.copy()
        session_config.update(
            {
                # our PUT/POST to ermrest is idempotent
                "allow_retry_on_all_methods": True,
                # do more retries before aborting
                "retry_read": 8,
                "retry_connect": 5,
                # increase delay factor * 2**(n-1) for Nth retry
                "retry_backoff_factor": 5,
            }
        )
        return session_config

    @property
    def pathBuilder(self) -> _CatalogWrapper:
        """Get a new instance of a pathBuilder object."""
        return self.catalog.getPathBuilder()

    @property
    def domain_path(self):
        """Get a new instance of a pathBuilder object to the domain schema"""

        return self.pathBuilder.schemas[self.domain_schema]

    def table_path(self, table: str | Table) -> Path:
        """Return a local file path in which to place a CSV to add values to a table on upload.

        Args:
          table: return:
          table: str | Table:

        Returns:
            Path to a CSV file in which to add values to a table on upload.
        """
        return table_path(
            self.working_dir,
            schema=self.domain_schema,
            table=self.model.get_table(table).name,
        )

    def asset_dir(
        self, table: str | Table, prefix: str | Path = None
    ) -> UploadAssetDirectory:
        """Return a local file path in which to place a files for an asset table.  T

        Args:
            table: Location of where to place files.  Defaults to execution_assets_path.
            prefix: Root path to asset directory.

        Returns:
            Path to the directory in which asset files should be placed.
        """
        table = self.model.get_table(table)
        if not self.model.is_asset(table):
            raise DerivaMLException(f"The table {table} is not an asset table.")

        prefix = Path(prefix) if prefix else self.working_dir
        return UploadAssetDirectory(
            model=self.model,
            prefix=prefix,
            schema=table.schema.name,
            table=table.name,
        )

    def download_dir(self, cached: bool = False) -> Path:
        """Location where downloaded files are placed.

        Args:
          cached: bool:  (Default value = False)

        Returns:

        """
        return self.cache_dir if cached else self.working_dir

    def chaise_url(self, table: RID | Table) -> str:
        """Return a Chaise URL to the specified table.

        Args:
            table: Table or RID to be visited
            table: str | Table:

        Returns:
            URL to the table in Chaise format.
        """
        try:
            table = self.model.get_table(table)
            uri = self.catalog.get_server_uri().replace(
                "ermrest/catalog/", "chaise/recordset/#"
            )
        except DerivaMLException:
            # Perhaps we have a RID....
            uri = self.cite(table)
        return f"{uri}/{urlquote(table.schema.name)}:{urlquote(table.name)}"

    def cite(self, entity: dict | str) -> str:
        """Return a citation URL for the provided entity.

        Args:
            entity: A dict that contains the column values for a specific entity or a RID.

        Returns:
            The URI for the provided entity.

        Raises:
            DerivaMLException: if provided RID does not exist.
        """
        if entity.startswith(f"https://{self.host_name}/id/{self.catalog_id}/"):
            # Already got a citation...
            return entity
        try:
            self.resolve_rid(
                rid := entity if isinstance(entity, str) else entity["RID"]
            )
            return f"https://{self.host_name}/id/{self.catalog_id}/{rid}@{self.catalog.latest_snapshot().snaptime}"
        except KeyError as e:
            raise DerivaMLException(f"Entity {e} does not have RID column")
        except DerivaMLException as _e:
            raise DerivaMLException("Entity RID does not exist")

    def user_list(self) -> list[dict[str, str]]:
        """List of users in the catalog

        Args:

        Returns:
          A list of dictionaries containing user information.

        """
        user_path = self.pathBuilder.public.ERMrest_Client.path
        return [
            {"ID": u["ID"], "Full_Name": u["Full_Name"]}
            for u in user_path.entities().fetch()
        ]

    def resolve_rid(self, rid: RID) -> ResolveRidResult:
        """Return a named tuple with information about the specified RID.

        Args:
            rid: RID of the object of interest

        Returns:
            ResolveRidResult which has information about the specified RID.

        Raises:
          DerivaMLException: if the RID doesn't exist.
        """
        try:
            return self.catalog.resolve_rid(rid, self.model.model)
        except KeyError as _e:
            raise DerivaMLException(f"Invalid RID {rid}")

    def retrieve_rid(self, rid: RID) -> dict[str, Any]:
        """Return a dictionary that represents the values of the specified RID.

        Args:
            rid: RID of the object of interest

        Returns:
          A dictionary that represents the values of the specified RID.

        Raises:
          DerivaMLException: if the RID doesn't exist.
        """
        return self.resolve_rid(rid).datapath.entities().fetch()[0]

    def add_page(self, title: str, content: str) -> None:
        """

        Args:
          title: str:
          content: str:

        Returns:

        """
        self.pathBuilder.www.tables[self.domain_schema].insert(
            [{"Title": title, "Content": content}]
        )

    def create_vocabulary(
        self, vocab_name: str, comment: str = "", schema: Optional[str] = None
    ) -> Table:
        """Create a controlled vocabulary table with the given vocab name.

        Args:
            vocab_name: Name of the controlled vocabulary table.
            comment: Description of the vocabulary table. (Default value = '')
            schema: Schema in which to create the controlled vocabulary table.  Defaults to domain_schema.
            vocab_name: str:

        Returns:
            An ERMRest table object for the newly created vocabulary table.
        """
        schema = schema or self.domain_schema
        return self.model.schemas[schema].create_table(
            Table.define_vocabulary(
                vocab_name, f"{self.project_name}:{{RID}}", comment=comment
            )
        )

    def create_asset(
        self,
        asset_name: str,
        column_defs: Iterable[ColumnDefinition] = None,
        comment: str = "",
        schema: str = None,
    ) -> Table:
        """Create an asset table with the given asset name.

        Args:
            asset_name: Name of the asset table.
            column_defs: Iterable of ColumnDefinition objects to provide additional metadata for asset.
            comment: Description of the asset table. (Default value = '')
            schema: Schema in which to create the asset table.  Defaults to domain_schema.
            asset_name: str:
            schema: str:  (Default value = None)

        Returns:
            Table object for the asset table.
        """
        column_defs = column_defs or []
        schema = schema or self.domain_schema
        asset_table = self.model.schemas[schema].create_table(
            Table.define_asset(
                schema,
                asset_name,
                column_defs=[c.model_dump() for c in column_defs],
                comment=comment,
            )
        )
        return asset_table

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def create_feature(
        self,
        target_table: Table | str,
        feature_name: str,
        terms: list[Table | str] = None,
        assets: list[Table | str] = None,
        metadata: Iterable[ColumnDefinition | Table | Key | str] = None,
        optional: Optional[list[str]] = None,
        comment: str = "",
    ) -> type[FeatureRecord]:
        """Create a new feature that can be associated with a table.

        The feature can associate a controlled vocabulary term, an asset, or any other values with a
        specific instance of an object and  execution.

        Args:
            feature_name: Name of the new feature to be defined
            target_table: table name or object on which the feature is to be associated
            terms: List of controlled vocabulary terms that will be part of the feature value
            assets: List of asset table names or objects that will be part of the feature value
            metadata: List of other value types that are associated with the feature
            optional: List of columns that are optional in the feature
            comment: return: A Feature class that can be used to create instances of the feature.

        Returns:
            A Feature class that can be used to create instances of the feature.

        Raises:
            DerivaException: If the feature cannot be created.
        """

        terms = terms or []
        assets = assets or []
        metadata = metadata or []
        optional = optional or []

        def normalize_metadata(m: Key | Table | ColumnDefinition | str):
            """

            Args:
              m: Key | Table | ColumnDefinition | str:

            Returns:

            """
            if isinstance(m, str):
                return self.model.get_table(m)
            elif isinstance(m, ColumnDefinition):
                return m.model_dump()
            else:
                return m

        # Make sure that the provided assets or terms are actually assets or terms.
        if not all(map(self.model.is_asset, assets)):
            raise DerivaMLException(f"Invalid create_feature asset table.")
        if not all(map(self.model.is_vocabulary, terms)):
            raise DerivaMLException(f"Invalid create_feature asset table.")

        # Get references to the necessary tables and make sure that the
        # provided feature name exists.
        target_table = self.model.get_table(target_table)
        execution = self.model.schemas[self.ml_schema].tables["Execution"]
        feature_name_table = self.model.schemas[self.ml_schema].tables["Feature_Name"]
        feature_name_term = self.add_term(
            "Feature_Name", feature_name, description=comment
        )
        atable_name = f"Execution_{target_table.name}_{feature_name_term.name}"

        # Now create the association table that implements the feature.
        atable = self.model.schemas[self.domain_schema].create_table(
            target_table.define_association(
                table_name=atable_name,
                associates=[execution, target_table, feature_name_table],
                metadata=[
                    normalize_metadata(m) for m in chain(assets, terms, metadata)
                ],
                comment=comment,
            )
        )
        # Now set optional terms.
        for c in optional:
            atable.columns[c].alter(nullok=True)
        atable.columns["Feature_Name"].alter(default=feature_name_term.name)
        return self.feature_record_class(target_table, feature_name)

    def feature_record_class(
        self, table: str | Table, feature_name: str
    ) -> type[FeatureRecord]:
        """Create a pydantic model for entries into the specified feature table.

        For information on how to
        See the pydantic documentation for more details about the pydantic model.

        Args:
            table: table name or object on which the feature is to be associated
            feature_name: name of the feature to be created
            table: str | Table:
            feature_name: str:

        Returns:
            A Feature class that can be used to create instances of the feature.
        """
        return self.lookup_feature(table, feature_name).feature_record_class()

    def delete_feature(self, table: Table | str, feature_name: str) -> bool:
        """

        Args:
          table: Table | str:
          feature_name: str:

        Returns:
        """
        table = self.model.get_table(table)
        try:
            feature = next(
                f for f in self.find_features(table) if f.feature_name == feature_name
            )
            feature.feature_table.drop()
            return True
        except StopIteration:
            return False

    def lookup_feature(self, table: str | Table, feature_name: str) -> Feature:
        """Lookup the named feature associated with the provided table.

        Args:
            table: param feature_name:
            table: str | Table:
            feature_name: str:

        Returns:
            A Feature class that represents the requested feature.

        Raises:
          DerivaMLException: If the feature cannot be found.
        """
        table = self.model.get_table(table)
        try:
            return [
                f for f in self.find_features(table) if f.feature_name == feature_name
            ][0]
        except IndexError:
            raise DerivaMLException(
                f"Feature {table.name}:{feature_name} doesn't exist."
            )

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def find_features(self, table: Table | str) -> Iterable[Feature]:
        """List the names of the features in the specified table.

        Args:
            table: The table to find features for.
            table: Table | str:

        Returns:
            An iterable of FeatureResult instances that describe the current features in the table.
        """
        table = self.model.get_table(table)

        def is_feature(a: FindAssociationResult) -> bool:
            """

            Args:
              a: FindAssociationResult:

            Returns:

            """
            # return {'Feature_Name', 'Execution'}.issubset({c.name for c in a.table.columns})
            return {
                "Feature_Name",
                "Execution",
                a.self_fkey.foreign_key_columns[0].name,
            }.issubset({c.name for c in a.table.columns})

        return [
            Feature(a, self.model)
            for a in table.find_associations(min_arity=3, max_arity=3, pure=False)
            if is_feature(a)
        ]

    @validate_call
    def add_features(self, features: Iterable[FeatureRecord]) -> int:
        """Add a set of new feature values to the catalog.

        Args:
          features: Iterable[FeatureRecord]:

        Returns:
            Number of attributes added
        """
        features = list(features)
        feature_table = features[0].feature.feature_table
        feature_path = self.pathBuilder.schemas[feature_table.schema.name].tables[
            feature_table.name
        ]
        entries = feature_path.insert(f.model_dump() for f in features)
        return len(entries)

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def list_feature_values(self, table: Table | str, feature_name: str) -> ResultSet:
        """Return a datapath ResultSet containing all values of a feature associated with a table.

        Args:
            table: param feature_name:
            table: Table | str:
            feature_name: str:

        Returns:

        """
        table = self.model.get_table(table)
        feature = self.lookup_feature(table, feature_name)
        pb = self.catalog.getPathBuilder()
        return (
            pb.schemas[feature.feature_table.schema.name]
            .tables[feature.feature_table.name]
            .entities()
            .fetch()
        )

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def add_term(
        self,
        table: str | Table,
        term_name: str,
        description: str,
        synonyms: Optional[list[str]] = None,
        exists_ok: bool = True,
    ) -> VocabularyTerm:
        """Creates a new control vocabulary term in the control vocabulary table.

        Args:

        Args:
            table: The name of the control vocabulary table.
            term_name: The name of the new control vocabulary.
            description: The description of the new control vocabulary.
            synonyms: Optional list of synonyms for the new control vocabulary. Defaults to an empty list.
            exists_ok: Optional flag indicating whether to allow creation if the control vocabulary name
                already exists. Defaults to True.

        Returns:
          The RID of the newly created control vocabulary.

        Raises:
          DerivaException: If the control vocabulary name already exists and exist_ok is False.
        """
        synonyms = synonyms or []
        table = self.model.get_table(table)
        pb = self.catalog.getPathBuilder()
        if not (self.model.is_vocabulary(table)):
            raise DerivaMLException(f"The table {table} is not a controlled vocabulary")

        schema_name = table.schema.name
        table_name = table.name
        try:
            term_id = VocabularyTerm.model_validate(
                pb.schemas[schema_name]
                .tables[table_name]
                .insert(
                    [
                        {
                            "Name": term_name,
                            "Description": description,
                            "Synonyms": synonyms,
                        }
                    ],
                    defaults={"ID", "URI"},
                )[0]
            )
        except DataPathException:
            term_id = self.lookup_term(table, term_name)
            if not exists_ok:
                raise DerivaMLException(f"{term_name} already exists")
            # Check vocabulary
        return term_id

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def lookup_term(self, table: str | Table, term_name: str) -> VocabularyTerm:
        """Given a term name, return the vocabulary record.  Can provide either the term name
         or a synonym for the term.  Generate an exception if the term is not in the vocabulary.

        Args:
            table: The name of the controlled vocabulary table or a ERMRest table object.
            term_name: The name of the term to look up.

        Returns:
          The entry the associated term or synonym.

        Raises:
          DerivaException: If the schema or vocabulary table doesn't exist, or if the term is not
            found in the vocabulary.
        """
        vocab_table = self.model.get_table(table)
        if not self.model.is_vocabulary(vocab_table):
            raise DerivaMLException(f"The table {table} is not a controlled vocabulary")
        schema_name, table_name = vocab_table.schema.name, vocab_table.name
        schema_path = self.catalog.getPathBuilder().schemas[schema_name]
        for term in schema_path.tables[table_name].entities():
            if term_name == term["Name"] or (
                term["Synonyms"] and term_name in term["Synonyms"]
            ):
                return VocabularyTerm.model_validate(term)
        raise DerivaMLException(f"Term {term_name} is not in vocabulary {table_name}")

    def list_vocabulary_terms(self, table: str | Table) -> list[VocabularyTerm]:
        """Return a list of terms that are in a vocabulary table.

        Args:
            table: The name of the controlled vocabulary table or a ERMRest table object.
            table: str | Table:

        Returns:
            The list of terms that are in a vocabulary table.

        Raises:
            DerivaMLException: If the schema or vocabulary table doesn't exist, or if the table is not
                a controlled vocabulary.
        """
        pb = self.catalog.getPathBuilder()
        table = self.model.get_table(table)
        if not (self.model.is_vocabulary(table)):
            raise DerivaMLException(f"The table {table} is not a controlled vocabulary")

        return [
            VocabularyTerm(**v)
            for v in pb.schemas[table.schema.name].tables[table.name].entities().fetch()
        ]

    def download_asset(self, asset_url: str, dest_filename: str) -> Path:
        """Download an asset from a URL and place it in a local directory.

        Args:
            asset_url: URL of the asset.
            dest_filename: Destination filename.

        Returns:
            A  Path object to the downloaded asset.
        """
        hs = HatracStore("https", self.host_name, self.credential)
        hs.get_obj(path=asset_url, destfilename=dest_filename)
        return Path(dest_filename)

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def upload_assets(
        self,
        assets_dir: str | Path | UploadAssetDirectory,
    ) -> dict[Any, FileUploadState] | None:
        """Upload assets from a directory. This routine assumes that the current upload specification includes a
        configuration for the specified directory.  Every asset in the specified directory is uploaded

        Args:
            assets_dir: Directory containing the assets to upload.

        Returns:
            Results of the upload operation.

        Raises:
            DerivaMLException: If there is an issue uploading the assets.
        """

        def path_to_asset(path: str) -> str:
            components = path.split("/")
            return components[
                components.index("asset") + 2
            ]  # Look for asset in the path to find the name

        if isinstance(assets_dir, UploadAssetDirectory):
            assets_dir = assets_dir.path

        if not self.model.is_asset(Path(assets_dir).name):
            raise DerivaMLException("Directory does not have name of an asset table.")
        results = upload_directory(self.model, assets_dir)
        return {path_to_asset(p): r for p, r in results.items()}

    def _update_status(
        self, new_status: Status, status_detail: str, execution_rid: RID
    ):
        """Update the status of an execution in the catalog.

        Args:
            new_status: New status.
            status_detail: Details of the status.
            execution_rid: Resource Identifier (RID) of the execution.
            new_status: Status:
            status_detail: str:
             execution_rid: RID:

        Returns:

        """
        self.status = new_status.value
        self.pathBuilder.schemas[self.ml_schema].Execution.update(
            [
                {
                    "RID": execution_rid,
                    "Status": self.status,
                    "Status_Detail": status_detail,
                }
            ]
        )

    # @validate_call
    def create_execution(self, configuration: ExecutionConfiguration) -> "Execution":
        """Create an execution object

        Given an execution configuration, initialize the local compute environment to prepare for executing an
        ML or analytic routine.  This routine has a number of side effects.

        1. The datasets specified in the configuration are downloaded and placed in the cache-dir. If a version is
        not specified in the configuration, then a new minor version number is created for the dataset and downloaded.

        Args:
            configuration: ExecutionConfiguration:

        Returns:
            An execution object.
        """
        from .execution import Execution

        return Execution(configuration, self)

    # @validate_call
    def restore_execution(self, execution_rid: Optional[RID] = None) -> "Execution":
        from .execution import Execution

        # Find path to execution
        if not execution_rid:
            e_rids = execution_rids(self.working_dir)
            if len(e_rids) != 1:
                raise DerivaMLException(f"Multiple execution RIDs were found {e_rids}.")

            execution_rid = e_rids[0]
        cfile = (
            execution_metadata_dir(
                self.working_dir,
                exec_rid=execution_rid,
                metadata_type=ExecMetadataVocab.execution_config.value,
            )
            / "configuration.json"
        )
        configuration = ExecutionConfiguration.load_configuration(cfile)
        return Execution(configuration, self, reload=execution_rid)
