import warnings
from typing import Any, Literal, cast

from cognite.client.data_classes.data_modeling import DataModelId, DataModelIdentifier
from cognite.client.utils.useful_types import SequenceNotStr

from cognite.neat._alpha import AlphaFlags
from cognite.neat._client import NeatClient
from cognite.neat._constants import (
    CLASSIC_CDF_NAMESPACE,
    get_default_prefixes_and_namespaces,
)
from cognite.neat._graph import examples as instances_examples
from cognite.neat._graph import extractors
from cognite.neat._graph.transformers import (
    ConvertLiteral,
    LiteralToEntity,
    Transformers,
)
from cognite.neat._graph.transformers._prune_graph import (
    AttachPropertyFromTargetToSource,
    PruneDeadEndEdges,
    PruneInstancesOfUnknownType,
    PruneTypes,
)
from cognite.neat._issues import IssueList
from cognite.neat._issues.errors import NeatValueError
from cognite.neat._issues.warnings import MissingCogniteClientWarning
from cognite.neat._rules import catalog, importers
from cognite.neat._rules.importers import BaseImporter
from cognite.neat._rules.transformers import ClassicPrepareCore
from cognite.neat._utils.reader import NeatReader

from ._state import SessionState
from ._wizard import NeatObjectType, RDFFileType, XMLFileType, object_wizard, rdf_dm_wizard, xml_format_wizard
from .engine import import_engine
from .exceptions import NeatSessionError, session_class_wrapper


@session_class_wrapper
class ReadAPI:
    """Read from a data source into NeatSession graph store."""

    def __init__(self, state: SessionState, verbose: bool) -> None:
        self._state = state
        self._verbose = verbose
        self.cdf = CDFReadAPI(state, verbose)
        self.rdf = RDFReadAPI(state, verbose)
        self.excel = ExcelReadAPI(state, verbose)
        self.csv = CSVReadAPI(state, verbose)
        self.yaml = YamlReadAPI(state, verbose)
        self.xml = XMLReadAPI(state, verbose)
        self.examples = Examples(state)

    def session(self, io: Any) -> None:
        """Reads a Neat Session from a zip file.

        Args:
            io: file path to the Neat Session

        Example:
            ```python
            neat.read.session("path_to_neat_session")
            ```
        """
        reader = NeatReader.create(io)
        path = reader.materialize_path()

        self._state.instances.store.write(extractors.RdfFileExtractor.from_zip(path))


@session_class_wrapper
class BaseReadAPI:
    def __init__(self, state: SessionState, verbose: bool) -> None:
        self._state = state
        self._verbose = verbose


@session_class_wrapper
class CDFReadAPI(BaseReadAPI):
    """Reads from CDF Data Models.
    Use the `.data_model()` method to load a CDF Data Model to the knowledge graph.

    """

    def __init__(self, state: SessionState, verbose: bool) -> None:
        super().__init__(state, verbose)
        self.classic = CDFClassicAPI(state, verbose)

    @property
    def _get_client(self) -> NeatClient:
        if self._state.client is None:
            raise NeatValueError("No client provided. Please provide a client to read a data model.")
        return self._state.client

    def data_model(self, data_model_id: DataModelIdentifier) -> IssueList:
        """Reads a Data Model from CDF to the knowledge graph.

        Args:
            data_model_id: Tuple of strings with the id of a CDF Data Model.
            Notation as follows (<name_of_space>, <name_of_data_model>, <data_model_version>)

        Example:
            ```python
            neat.read.cdf.data_model(("example_data_model_space", "EXAMPLE_DATA_MODEL", "v1"))
            ```
        """

        data_model_id = DataModelId.load(data_model_id)

        if not data_model_id.version:
            raise NeatSessionError("Data model version is required to read a data model.")

        importer = importers.DMSImporter.from_data_model_id(self._get_client, data_model_id)
        return self._state.rule_import(importer)

    def graph(
        self,
        data_model_id: DataModelIdentifier,
        instance_space: str | SequenceNotStr[str] | None = None,
        skip_cognite_views: bool = True,
    ) -> IssueList:
        """Reads a knowledge graph from Cognite Data Fusion (CDF).

        Args:
            data_model_id: Tuple of strings with the id of a CDF Data Model.
            instance_space: The instance spaces to extract. If None, all instance spaces are extracted.
            skip_cognite_views: If True, all Cognite Views are skipped. For example, if you have the CogniteAsset
                view in you data model, it will ont be used to extract instances.

        Returns:
            IssueList: A list of issues that occurred during the extraction.

        """
        return self._graph(data_model_id, instance_space, skip_cognite_views, unpack_json=False)

    def _graph(
        self,
        data_model_id: DataModelIdentifier,
        instance_space: str | SequenceNotStr[str] | None = None,
        skip_cognite_views: bool = True,
        unpack_json: bool = False,
        str_to_ideal_type: bool = False,
    ) -> IssueList:
        extractor = extractors.DMSGraphExtractor.from_data_model_id(
            # We are skipping the Cognite Views
            data_model_id,
            self._get_client,
            instance_space=instance_space,
            skip_cognite_views=skip_cognite_views,
            unpack_json=unpack_json,
            str_to_ideal_type=str_to_ideal_type,
        )
        return self._state.write_graph(extractor)

    def raw(
        self,
        db_name: str,
        table_name: str,
        type: str | None = None,
        foreign_keys: str | SequenceNotStr[str] | None = None,
        unpack_json: bool = False,
        str_to_ideal_type: bool = False,
    ) -> IssueList:
        """Reads a raw table from CDF to the knowledge graph.

        Args:
            db_name: The name of the database
            table_name: The name of the table, this will be assumed to be the type of the instances.
            type: The type of instances in the table. If None, the table name will be used.
            foreign_keys: The name of the columns that are foreign keys. If None, no foreign keys are used.
            unpack_json: If True, the JSON objects will be unpacked into the graph.
            str_to_ideal_type: If True, the string values will be converted to ideal types.

        Returns:
            IssueList: A list of issues that occurred during the extraction.

        Example:
            ```python
            neat.read.cdf.raw("my_db", "my_table", "Asset")
            ```

        """
        extractor = extractors.RAWExtractor(
            self._get_client,
            db_name=db_name,
            table_name=table_name,
            table_type=type,
            foreign_keys=foreign_keys,
            unpack_json=unpack_json,
            str_to_ideal_type=str_to_ideal_type,
        )
        return self._state.instances.store.write(extractor)


@session_class_wrapper
class CDFClassicAPI(BaseReadAPI):
    """Reads from the Classic Data Model from CDF.
    Use the `.graph()` method to load CDF core resources to the knowledge graph.

    """

    @property
    def _get_client(self) -> NeatClient:
        if self._state.client is None:
            raise ValueError("No client provided. Please provide a client to read a data model.")
        return self._state.client

    def graph(
        self,
        root_asset_external_id: str,
        limit_per_type: int | None = None,
        identifier: Literal["id", "externalId"] = "id",
    ) -> IssueList:
        """Reads the classic knowledge graph from CDF.

        The Classic Graph consists of the following core resource type.

        !!! note "Classic Node CDF Resources"
             - Assets
             - TimeSeries
             - Sequences
             - Events
             - Files

        All the classic node CDF resources can have one or more connections to one or more assets. This
        will match a direct relationship in the data modeling of CDF.

        In addition, you have relationships between the classic node CDF resources. This matches an edge
        in the data modeling of CDF.

        Finally, you have labels and data sets that to organize the graph. In which data sets have a similar,
        but different, role as a space in data modeling. While labels can be compared to node types in data modeling,
        used to quickly filter and find nodes/edges.

        This extractor will extract the classic CDF graph into Neat starting from either a data set or a root asset.

        It works as follows:
            1. Extract all core nodes (assets, time series, sequences, events, files) filtered by the given data set or
               root asset.
            2. Extract all relationships starting from any of the extracted core nodes.
            3. Extract all core nodes that are targets of the relationships that are not already extracted.
            4. Extract all labels that are connected to the extracted core nodes/relationships.
            5. Extract all data sets that are connected to the extracted core nodes/relationships.

        Args:
            root_asset_external_id: The external id of the root asset
            limit_per_type: The maximum number of nodes to extract per core node type. If None, all nodes are extracted.
            identifier: The identifier to use for the core nodes. Note selecting "id" can cause issues if the external
                ID of the core nodes is missing. Default is "id".

        Returns:
            IssueList: A list of issues that occurred during the extraction.

        Example:
            ```python
            neat.read.cdf.graph("root_asset_external_id")
            ```
        """
        return self._graph(
            root_asset_external_id, limit_per_type, identifier, reference_timeseries=False, reference_files=False
        )

    def _graph(
        self,
        root_asset_external_id: str,
        limit_per_type: int | None = None,
        identifier: Literal["id", "externalId"] = "id",
        reference_timeseries: bool = False,
        reference_files: bool = False,
        unpack_metadata: bool = False,
        skip_sequence_rows: bool = False,
    ) -> IssueList:
        namespace = CLASSIC_CDF_NAMESPACE
        extractor = extractors.ClassicGraphExtractor(
            self._get_client,
            root_asset_external_id=root_asset_external_id,
            limit_per_type=limit_per_type,
            namespace=namespace,
            prefix="Classic",
            identifier=identifier,
            unpack_metadata=unpack_metadata,
            skip_sequence_rows=skip_sequence_rows,
        )
        self._state.instances.neat_prefix_by_predicate_uri.update(extractor.neat_prefix_by_predicate_uri)
        self._state.instances.neat_prefix_by_type_uri.update(extractor.neat_prefix_by_type_uri)
        extract_issues = self._state.write_graph(extractor)
        if identifier == "externalId":
            self._state.quoted_source_identifiers = True

        self._state.instances.store.transform(
            ConvertLiteral(
                namespace["ClassicTimeSeries"],
                namespace["isString"],
                lambda is_string: "string" if is_string else "numeric",
            )
        )
        self._state.instances.store.transform(
            LiteralToEntity(None, namespace["source"], "ClassicSourceSystem", "name"),
        )
        # Updating the information model.
        prepare_issues = self._state.rule_store.transform(
            ClassicPrepareCore(namespace, reference_timeseries, reference_files)
        )

        all_issues = IssueList(extract_issues + prepare_issues)
        # Update the provenance with all issue.
        object.__setattr__(self._state.instances.store.provenance[-1].target_entity, "issues", all_issues)
        all_issues.action = "Read Classic Graph"
        if all_issues:
            print("Use the .inspect.issues() for more details")

        return all_issues


@session_class_wrapper
class ExcelReadAPI(BaseReadAPI):
    """Reads a Neat Excel Rules sheet to the graph store. The rules sheet may stem from an Information architect,
    or a DMS Architect.

    Args:
        io: file path to the Excel sheet

    Example:
        ```python
        neat.read.excel("information_or_dms_rules_sheet.xlsx")
        ```
    """

    def __init__(self, state: SessionState, verbose: bool) -> None:
        super().__init__(state, verbose)

    def __call__(self, io: Any, enable_manual_edit: bool = False) -> IssueList:
        """Reads a Neat Excel Rules sheet to the graph store. The rules sheet may stem from an Information architect,
        or a DMS Architect.

        Args:
            io: file path to the Excel sheet
            enable_manual_edit: If True, the user will be able to re-import rules which where edit outside NeatSession

        !!! note "Manual Edit Warning"
            This is an alpha feature and is subject to change without notice.
            It is expected to have some limitations and may not work as expected in all cases.
        """
        reader = NeatReader.create(io)
        path = reader.materialize_path()

        if enable_manual_edit:
            warnings.filterwarnings("default")
            AlphaFlags.manual_rules_edit.warn()

        return self._state.rule_import(importers.ExcelImporter(path), enable_manual_edit)


@session_class_wrapper
class YamlReadAPI(BaseReadAPI):
    def __call__(self, io: Any, format: Literal["neat", "toolkit"] = "neat") -> IssueList:
        """Reads a yaml with either neat rules, or several toolkit yaml files to import Data Model(s) into NeatSession.

        Args:
            io: File path to the Yaml file in the case of "neat" yaml, or path to a zip folder or directory with several
                Yaml files in the case of "toolkit".
            format: The format of the yaml file(s). Can be either "neat" or "toolkit".

        Example:
            ```python
            neat.read.yaml("path_to_toolkit_yamls")
            ```
        """
        reader = NeatReader.create(io)
        path = reader.materialize_path()
        importer: BaseImporter
        if format == "neat":
            importer = importers.YAMLImporter.from_file(path, source_name=f"{reader!s}")
        elif format == "toolkit":
            dms_importer = importers.DMSImporter.from_path(path, self._state.client)
            if dms_importer.issue_list.has_warning_type(MissingCogniteClientWarning):
                raise NeatSessionError(
                    "No client provided. You are referencing Cognite containers in your data model, "
                    "NEAT needs a client to lookup the container definitions. "
                    "Please set the client in the session, NeatSession(client=client)."
                )
            importer = dms_importer
        else:
            raise NeatValueError(f"Unsupported YAML format: {format}")
        return self._state.rule_import(importer)


@session_class_wrapper
class CSVReadAPI(BaseReadAPI):
    """Reads a csv that contains a column to use as primary key which will be the unique identifier for the type of
    data you want to read in. Ex. a csv can hold information about assets, and their identifiers are specified in
    a "ASSET_TAG" column.

    Args:
        io: file path or url to the csv
        type: string that specifies what type of data the csv contains. For instance "Asset" or "Equipment"
        primary_key: string name of the column that should be used as the unique identifier for each row of data

    Example:
        ```python
        type_described_in_table = "Turbine"
        column_with_identifier = "UNIQUE_TAG_NAME"
        neat.read.csv("url_or_path_to_csv_file", type=type_described_in_table, primary_key=column_with_identifier)
        ```
    """

    def __call__(self, io: Any, type: str, primary_key: str) -> None:
        warnings.filterwarnings("default")
        AlphaFlags.csv_read.warn()

        engine = import_engine()
        engine.set.format = "csv"
        engine.set.file = NeatReader.create(io).materialize_path()
        engine.set.type = type
        engine.set.primary_key = primary_key
        extractor = engine.create_extractor()

        self._state.instances.store.write(extractor)


@session_class_wrapper
class XMLReadAPI(BaseReadAPI):
    """Reads an XML file that is either of DEXPI or AML format.

    Args:
        io: file path or url to the XML
        format: can be either "dexpi" or "aml" are the currenly supported XML source types.
    """

    def __call__(
        self,
        io: Any,
        format: XMLFileType | None = None,
    ) -> None:
        path = NeatReader.create(io).materialize_path()
        if format is None:
            format = xml_format_wizard()

        if format.lower() == "dexpi":
            return self.dexpi(path)

        if format.lower() == "aml":
            return self.aml(path)

        else:
            raise NeatValueError("Only support XML files of DEXPI format at the moment.")

    def dexpi(self, io: Any) -> None:
        """Reads a DEXPI file into the NeatSession and executes set of predefined transformations.

        Args:
            io: file path or url to the DEXPI file

        Example:
            ```python
            neat.read.xml.dexpi("url_or_path_to_dexpi_file")
            ```

        !!! note "This method bundles several graph transformers which"
            - attach values of generic attributes to nodes
            - create associations between nodes
            - remove unused generic attributes
            - remove associations between nodes that do not exist in the extracted graph
            - remove edges to nodes that do not exist in the extracted graph
        """
        warnings.filterwarnings("default")
        AlphaFlags.dexpi_read.warn()

        path = NeatReader.create(io).materialize_path()
        engine = import_engine()
        engine.set.format = "dexpi"
        engine.set.file = path
        extractor = engine.create_extractor()
        self._state.instances.store.write(extractor)

        DEXPI = get_default_prefixes_and_namespaces()["dexpi"]

        transformers = [
            # Remove any instance which type is unknown
            PruneInstancesOfUnknownType(),
            # Directly connect generic attributes
            AttachPropertyFromTargetToSource(
                target_property=DEXPI.Value,
                target_property_holding_new_property=DEXPI.Name,
                target_node_type=DEXPI.GenericAttribute,
                delete_target_node=True,
            ),
            # Directly connect associations
            AttachPropertyFromTargetToSource(
                target_property=DEXPI.ItemID,
                target_property_holding_new_property=DEXPI.Type,
                target_node_type=DEXPI.Association,
                delete_target_node=True,
            ),
            # Remove unused generic attributes and associations
            PruneTypes([DEXPI.GenericAttribute, DEXPI.Association]),
            # Remove edges to nodes that do not exist in the extracted graph
            PruneDeadEndEdges(),
        ]

        for transformer in transformers:
            self._state.instances.store.transform(cast(Transformers, transformer))

    def aml(self, io: Any):
        """Reads an AML file into NeatSession and executes a set of predefined transformations.

        Args:
            io: file path or url to the AML file

        Example:
            ```python
            neat.read.xml.aml("url_or_path_to_aml_file")
            ```

        !!! note "This method bundles several graph transformers which"
            - attach values of attributes to nodes
            - remove unused attributes
            - remove edges to nodes that do not exist in the extracted graph
        """
        warnings.filterwarnings("default")
        AlphaFlags.aml_read.warn()

        path = NeatReader.create(io).materialize_path()
        engine = import_engine()
        engine.set.format = "aml"
        engine.set.file = path
        extractor = engine.create_extractor()
        self._state.instances.store.write(extractor)

        AML = get_default_prefixes_and_namespaces()["aml"]

        transformers = [
            # Remove any instance which type is unknown
            PruneInstancesOfUnknownType(),
            # Directly connect generic attributes
            AttachPropertyFromTargetToSource(
                target_property=AML.Value,
                target_property_holding_new_property=AML.Name,
                target_node_type=AML.Attribute,
                delete_target_node=True,
            ),
            # Prune unused attributes
            PruneTypes([AML.Attribute]),
            # # Remove edges to nodes that do not exist in the extracted graph
            PruneDeadEndEdges(),
        ]

        for transformer in transformers:
            self._state.instances.store.transform(cast(Transformers, transformer))


@session_class_wrapper
class RDFReadAPI(BaseReadAPI):
    """Reads an RDF source into NeatSession. Supported sources are "ontology" or "imf".

    Args:
        io: file path or url to the RDF source
    """

    def __init__(self, state: SessionState, verbose: bool) -> None:
        super().__init__(state, verbose)

    def ontology(self, io: Any) -> IssueList:
        """Reads an OWL ontology source into NeatSession.

        Args:
            io: file path or url to the OWL file

        Example:
            ```python
            neat.read.rdf.ontology("url_or_path_to_owl_source")
            ```
        """
        warnings.filterwarnings("default")
        AlphaFlags.ontology_read.warn()

        reader = NeatReader.create(io)
        importer = importers.OWLImporter.from_file(reader.materialize_path(), source_name=f"file {reader!s}")
        return self._state.rule_import(importer)

    def imf(self, io: Any) -> IssueList:
        """Reads IMF Types provided as SHACL shapes into NeatSession.

        Args:
            io: file path or url to the IMF file

        Example:
            ```python
            neat.read.rdf.imf("url_or_path_to_imf_source")
            ```
        """
        warnings.filterwarnings("default")
        AlphaFlags.imf_read.warn()

        reader = NeatReader.create(io)
        importer = importers.IMFImporter.from_file(reader.materialize_path(), source_name=f"file {reader!s}")
        return self._state.rule_import(importer)

    def instances(self, io: Any) -> IssueList:
        reader = NeatReader.create(io)
        self._state.instances.store.write(extractors.RdfFileExtractor(reader.materialize_path()))
        return IssueList()

    def __call__(
        self,
        io: Any,
        type: NeatObjectType | None = None,
        source: RDFFileType | None = None,
    ) -> IssueList:
        if type is None:
            type = object_wizard()

        type = type.lower()

        if type == "data model":
            source = source or rdf_dm_wizard("What type of data model is the RDF?")
            source = cast(str, source).lower()  # type: ignore

            if source == "ontology":
                return self.ontology(io)
            elif source == "imf types":
                return self.imf(io)
            else:
                raise ValueError(f"Expected ontology, imf types or instances, got {source}")

        elif type == "instances":
            return self.instances(io)

        else:
            raise NeatSessionError(f"Expected data model or instances, got {type}")


@session_class_wrapper
class Examples:
    """Used as example for reading various sources into NeatSession."""

    def __init__(self, state: SessionState) -> None:
        self._state = state

    @property
    def _get_client(self) -> NeatClient:
        if self._state.client is None:
            raise NeatValueError("No client provided. Please provide a client to read a data model.")
        return self._state.client

    def nordic44(self) -> IssueList:
        """Reads the Nordic 44 knowledge graph into the NeatSession graph store."""
        self._state.instances.store.write(extractors.RdfFileExtractor(instances_examples.nordic44_knowledge_graph))
        return IssueList()

    def pump_example(self) -> IssueList:
        """Reads the Hello World pump example into the NeatSession."""
        importer: importers.ExcelImporter = importers.ExcelImporter(catalog.hello_world_pump)
        return self._state.rule_import(importer)

    def core_data_model(self) -> IssueList:
        """Reads the core data model example into the NeatSession."""

        cdm_v1 = DataModelId.load(("cdf_cdm", "CogniteCore", "v1"))
        importer: importers.DMSImporter = importers.DMSImporter.from_data_model_id(self._get_client, cdm_v1)
        return self._state.rule_import(importer)
