from __future__ import annotations
from dataclasses import dataclass
import logging
import os
import json
from typing import Any, List, Dict, Optional, Union, Type, TYPE_CHECKING

from modelon.impact.client.entities.interfaces.workspace import WorkspaceInterface
from modelon.impact.client.sal.service import Service
from modelon.impact.client.experiment_definition.interfaces.definition import (
    BaseExperimentDefinition,
)
from modelon.impact.client.entities.custom_function import CustomFunction
from modelon.impact.client.operations.workspace.exports import (
    WorkspaceExportOperation,
    Export,
)
from modelon.impact.client.operations.workspace.imports import WorkspaceImportOperation
from modelon.impact.client.operations.workspace.conversion import (
    WorkspaceConversionOperation,
)
from modelon.impact.client.operations.project_import import ProjectImportOperation
from modelon.impact.client.operations.experiment import ExperimentOperation
from modelon.impact.client.operations.external_result_import import (
    ExternalResultImportOperation,
)
from modelon.impact.client.entities.model import Model
from modelon.impact.client.entities.external_result import ExternalResult
from modelon.impact.client.entities.model_executable import ModelExecutable
from modelon.impact.client.entities.experiment import Experiment
from modelon.impact.client.entities.project import Project, VcsUri

if TYPE_CHECKING:
    from modelon.impact.client.operations.base import BaseOperation

logger = logging.getLogger(__name__)

ExperimentDefinition = Union[
    Type[BaseExperimentDefinition],
    Dict[str, Any],
]


@dataclass
class Reference:
    id: str


@dataclass
class ReleasedProjectReference:
    id: str
    name: str
    version: Optional[str] = None
    build: Optional[str] = None


@dataclass
class VcsReference:
    id: str
    vcs_uri: str


def _get_project_entry_reference(
    reference: Any,
) -> Union[ReleasedProjectReference, VcsReference, Reference]:
    if "name" in reference:
        return ReleasedProjectReference(
            id=reference.get('id'),
            name=reference.get('name'),
            version=reference.get('version'),
            build=reference.get('build'),
        )
    elif "vcsUri" in reference:
        return VcsReference(id=reference.get('id'), vcs_uri=reference.get('vcsUri'))
    else:
        return Reference(id=reference.get('id'))


class ProjectEntry:
    def __init__(self, data: Any) -> None:
        self._data = data

    @property
    def reference(self) -> Union[ReleasedProjectReference, VcsReference, Reference]:
        return _get_project_entry_reference(self._data.get('reference'))

    @property
    def id(self) -> str:
        return self.reference.id

    @property
    def disabled(self) -> bool:
        return self._data.get('disabled')

    @property
    def disabled_content(self) -> bool:
        return self._data.get('disabledContent')


class WorkspaceDefinition:
    def __init__(self, data: Any) -> None:
        self._data = data

    @property
    def name(self) -> str:
        return self._data.get('name')

    @property
    def format(self) -> str:
        return self._data.get('format')

    @property
    def description(self) -> str:
        return self._data.get('description')

    @property
    def created_by(self) -> str:
        return self._data.get('createdBy')

    @property
    def default_project_id(self) -> str:
        return self._data.get('defaultProjectId')

    @property
    def projects(self) -> List[ProjectEntry]:
        projects = self._data.get('projects', [])
        return [ProjectEntry(project) for project in projects]

    @property
    def dependencies(self) -> List[ProjectEntry]:
        dependencies = self._data.get('dependencies', [])
        return [ProjectEntry(dependency) for dependency in dependencies]

    def to_file(self, path: str) -> str:
        os.makedirs(path, exist_ok=True)
        definition_path = os.path.join(path, self.name + ".json")
        with open(definition_path, "w", encoding='utf-8') as f:
            json.dump(self._data, f, indent=4)
        return definition_path

    @classmethod
    def from_file(cls, path: str) -> WorkspaceDefinition:
        with open(path) as json_file:
            data = json.load(json_file)
        return cls(data)

    def to_dict(self) -> Dict[str, Any]:
        return self._data


class Workspace(WorkspaceInterface):
    """Class containing Workspace functionalities."""

    def __init__(
        self,
        workspace_id: str,
        workspace_definition: Union[WorkspaceDefinition, Dict[str, Any]],
        service: Service,
    ):
        self._workspace_id = workspace_id
        self._workspace_definition = (
            WorkspaceDefinition(workspace_definition)
            if isinstance(workspace_definition, dict)
            else workspace_definition
        )
        self._sal = service

    def __eq__(self, obj: object) -> bool:
        return isinstance(obj, Workspace) and obj._workspace_id == self._workspace_id

    def __repr__(self) -> str:
        return f"Workspace with id '{self._workspace_id}'"

    @property
    def id(self) -> str:
        """Workspace id."""
        return self._workspace_id

    @property
    def size(self) -> float:
        """Workspace size in bytes."""
        return self._sal.workspace.workspace_get(workspace_id=self.id, size_info=True)[
            "sizeInfo"
        ]["total"]

    @property
    def definition(self) -> WorkspaceDefinition:
        return self._workspace_definition

    def get_custom_function(self, name: str) -> CustomFunction:
        """Returns a CustomFunction class object.

        Args:
            name: The name of the custom function.

        Returns:
            The CustomFunction class object.

        Example::

            workspace.get_custom_function('dynamic')

        """
        custom_function = self._sal.custom_function.custom_function_get(
            self._workspace_id, name
        )
        return CustomFunction(
            self._workspace_id,
            custom_function["name"],
            custom_function["parameters"],
            self._sal,
        )

    def get_custom_functions(self) -> List[CustomFunction]:
        """Returns a list of CustomFunctions class objects.

        Returns:
            A list of CustomFunction class objects.

        Example::

            workspace.get_custom_functions()

        """
        custom_functions = self._sal.custom_function.custom_functions_get(
            self._workspace_id
        )
        return [
            CustomFunction(
                self._workspace_id,
                custom_function["name"],
                custom_function["parameters"],
                self._sal,
            )
            for custom_function in custom_functions["data"]["items"]
        ]

    def delete(self) -> None:
        """Deletes a workspace.

        Example::

            workspace.delete()

        """
        self._sal.workspace.workspace_delete(self._workspace_id)

    def upload_result(
        self,
        path_to_result: str,
        label: Optional[str] = None,
        description: Optional[str] = None,
    ) -> ExternalResultImportOperation:
        """Uploads a '.mat' result file to the workspace.

        Args:
            path_to_result: The path for the result file to be imported.
            label: The label of the result file. Default: None.
            description: The description of the result file. Default: None.

        Example::

            workspace.upload_result('A.mat')
            workspace.upload_result('B.mat', label = "result_for_PID.mat",
            description = "This is a result file for PID controller")

        """
        resp = self._sal.external_result.result_upload(
            self._workspace_id, path_to_result, label=label, description=description
        )
        return ExternalResultImportOperation[ExternalResult](
            resp["data"]["location"], self._sal, ExternalResult.from_operation
        )

    def export(self, publish: bool = False) -> WorkspaceExportOperation:
        """Exports the workspace as a binary compressed archive. Similar to
        :obj:`~modelon.impact.client.entities.workspace.Workspace.download`,
        but gives more control for getting the workspace async.
        Returns an modelon.impact.client.operations.workspace.exports
        .WorkspaceExportOperation class object. The binary archive is stored
        to cloud storage if the publish argument is set to True.

        Args:
            publish: To export the workspace and save it to cloud storage.

        Returns:
            An WorkspaceExportOperation class object.

        Example::

            path = workspace.export().wait().download_as('/home/workspace.zip')
            path = workspace.export(publish=True)
        """
        resp = self._sal.workspace.workspace_export_setup(self._workspace_id, publish)
        return WorkspaceExportOperation[Workspace](
            resp["data"]["location"], self._sal, Export.from_operation
        )

    def download(self, path: str) -> str:
        """Downloads the workspace as a binary compressed archive. Returns the
        local path to the downloaded workspace archive. Similar to
        :obj:`~modelon.impact.client.entities.workspace.Workspace.export`, but
        does the entire setup and download in one go.

        Args:

            path: The local path to store the downloaded workspace.

        Returns:
            Local path to the downloaded workspace archive.

        Example::

            workspace.download(path)

        """
        ws_path = os.path.join(path, self._workspace_id + ".zip")
        ops = self.export().wait()
        return ops.download_as(ws_path)

    def get_model(self, class_name: str, project: Optional[Project] = None) -> Model:
        """Returns a Model class object.

        Args:
            class_name: The Modelica class path of the model.
            project: Project class object

        Returns:
            Model class object.

        Example::

            workspace.get_model(class_name)

        """
        project = project or self.get_default_project()
        return Model(class_name, self._workspace_id, project.id, self._sal)

    def get_fmus(self) -> List[ModelExecutable]:
        """Returns a list of ModelExecutable class objects.

        Returns:
            List of ModelExecutable class objects.

        Example::

            workspace.get_fmus()

        """
        resp = self._sal.workspace.fmus_get(self._workspace_id)
        return [
            ModelExecutable(self._workspace_id, item["id"], self._sal, item)
            for item in resp["data"]["items"]
        ]

    def get_fmu(self, fmu_id: str) -> ModelExecutable:
        """Returns a ModelExecutable class object.

        Returns:
            ModelExecutable class object.

        Example::

            workspace.get_fmu(fmu_id)

        """
        resp = self._sal.workspace.fmu_get(self._workspace_id, fmu_id)
        return ModelExecutable(self._workspace_id, resp["id"], self._sal, resp)

    def get_experiments(self, class_path: Optional[str] = None) -> List[Experiment]:
        """Returns a list of Experiment class objects.

        Parameters:

            class_path --
                The modelica class path. If given, only the experiments
                generated for model with the specified class path
                are returned.

        Returns:
            List of Experiment class objects.

        Example::

            workspace.get_experiments()

        """
        resp = self._sal.workspace.experiments_get(self._workspace_id, class_path)
        return [
            Experiment(self._workspace_id, item["id"], self._sal, item)
            for item in resp["data"]["items"]
        ]

    def get_experiment(self, experiment_id: str) -> Experiment:
        """Returns an Experiment class object.

        Args:
            experiment_id:
                The ID of the experiment.

        Returns:
            Experiment class object.

        Example::

            workspace.get_experiment(experiment_id)

        """
        resp = self._sal.workspace.experiment_get(self._workspace_id, experiment_id)
        return Experiment(self._workspace_id, resp["id"], self._sal, resp)

    def create_experiment(
        self,
        definition: ExperimentDefinition,
        user_data: Optional[Dict[str, Any]] = None,
    ) -> Experiment:
        """Creates an experiment. Returns an Experiment class object.

        Args:
            definition: A parametrized experiment definition class of type
                SimpleModelicaExperimentDefinition or SimpleFMUExperimentDefinition.
            user_data:
                Optional dictionary object with custom data to attach to the experiment.

        Returns:
            Experiment class object.

        Example::

            workspace.create_experiment(definition)

        """
        if isinstance(
            definition,
            BaseExperimentDefinition,
        ):
            definition_as_dict = definition.to_dict()
        elif isinstance(definition, dict):
            definition_as_dict = definition
        else:
            raise TypeError(
                "Definition object must either be a dictionary or an instance of either"
                "SimpleModelicaExperimentDefinition or SimpleFMUExperimentDefinition"
                "class!"
            )
        resp = self._sal.workspace.experiment_create(
            self._workspace_id, definition_as_dict, user_data
        )
        return Experiment(self._workspace_id, resp["experiment_id"], self._sal)

    def execute(
        self,
        definition: ExperimentDefinition,
        user_data: Optional[Dict[str, Any]] = None,
    ) -> ExperimentOperation:
        """Executes an experiment.

        Args:
            definition:
                An experiment definition class instance of
                SimpleFMUExperimentDefinition
                or SimpleModelicaExperimentDefinition
                or a dictionary object containing the definition.
            user_data:
                Optional dictionary object with custom data to attach to the experiment.


        Returns:
            An ExperimentOperation class object.

        Example::

            experiment_ops = workspace.execute(definition)
            experiment_ops.cancel()
            experiment_ops.status
            experiment_ops.wait()

        """
        exp_id = self.create_experiment(definition, user_data).id
        return ExperimentOperation[Experiment](
            self._workspace_id,
            self._sal.experiment.experiment_execute(self._workspace_id, exp_id),
            self._sal,
            Experiment.from_operation,
        )

    def get_projects(
        self, vcs_info: bool = True, include_disabled: bool = False
    ) -> List[Project]:
        """Return the list of projects for a workspace.

        Args:
            vcs_info: If True, the versioning details are returned for the
            project(if under version control).
            include_disabled: If True, projects disabled in the workspace
            are also listed.

        Returns:
            A list of Project class objects.

        Example::

            projects = workspace.get_projects()

        """
        resp = self._sal.workspace.projects_get(
            self._workspace_id, vcs_info=vcs_info, include_disabled=include_disabled
        )
        projects = [
            Project(
                item["id"],
                item['definition'],
                item["projectType"],
                VcsUri.from_dict(item["vcsUri"]) if item.get("vcsUri") else None,
                self._sal,
            )
            for item in resp["data"]["items"]
        ]
        return projects

    def get_dependencies(
        self, vcs_info: bool = True, include_disabled: bool = False
    ) -> List[Project]:
        """Return the list of project dependencies for a workspace.

        Args:
            vcs_info: If True, the versioning details are returned for the
            project(if under version control).
            include_disabled: If True, projects disabled in the workspace
            are also listed.

        Returns:
            A list of Project class object.

        Example::

            dependencies = workspace.get_dependencies()

        """
        resp = self._sal.workspace.dependencies_get(
            self._workspace_id, vcs_info, include_disabled
        )
        return [
            Project(
                item["id"],
                item['definition'],
                item["projectType"],
                VcsUri.from_dict(item["vcsUri"]) if item.get("vcsUri") else None,
                self._sal,
            )
            for item in resp["data"]["items"]
        ]

    def create_project(self, name: str) -> Project:
        """Creates a new project in the workspace.

        Returns:
            A Project class object.

        Example::

            project = workspace.create_project("test")

        """
        resp = self._sal.workspace.project_create(self._workspace_id, name)
        return Project(
            resp["id"],
            resp['definition'],
            resp["projectType"],
            VcsUri.from_dict(resp["vcsUri"]) if resp.get("vcsUri") else None,
            self._sal,
        )

    def get_default_project(self) -> Project:
        """Return the default project for a workspace.

        Returns:
            An Project class object.

        Example::

            project = workspace.get_default_project()

        """
        if not self._workspace_definition.default_project_id:
            raise ValueError(
                f'No default project exists for the workspace {self._workspace_id}!'
            )
        resp = self._sal.project.project_get(
            self._workspace_definition.default_project_id,
            vcs_info=True,
            size_info=False,
        )
        return Project(
            resp["id"],
            resp["definition"],
            resp["projectType"],
            VcsUri.from_dict(resp["vcsUri"]) if resp.get("vcsUri") else None,
            self._sal,
        )

    def get_shared_definition(self, strict: bool = False) -> WorkspaceDefinition:
        """Return the shared definition for a workspace.

        Args:
            strict: If True, the version control URIs for workspace projects and
                dependencies are to specific commits or not.

        Returns:
            An WorkspaceDefinition class object.

        Example::

            definition = workspace.get_shared_definition()

        """
        return WorkspaceDefinition(
            self._sal.workspace.shared_definition_get(
                self._workspace_id, strict=strict
            )["definition"]
        )

    def import_project_from_zip(self, path_to_project: str) -> ProjectImportOperation:
        """Imports a Project from a compressed(.zip) project file and adds it
        to the workspace. Returns the project class object.

        Args:
            path_to_project: The path for the compressed project(.zip) to be uploaded.

        Returns:
            A ProjectImportOperation class object.

        Example::

            workspace.import_project_from_zip(path_to_project).wait()

        """
        resp = self._sal.workspace.import_project_from_zip(
            self._workspace_id, path_to_project
        )
        return ProjectImportOperation[Project](
            resp["data"]["location"], self._sal, Project.from_operation
        )

    def import_dependency_from_zip(
        self, path_to_dependency: str
    ) -> ProjectImportOperation:
        """Imports a Project dependency from a compressed(.zip) project file
        and adds it to the workspace. Returns the project class object.

        Args:
            path_to_dependency: The path for the compressed project(.zip) to be
            uploaded.

        Returns:
            An ProjectImportOperation class object.

        Example::

            workspace.import_dependency_from_zip(path_to_project).wait()

        """
        resp = self._sal.workspace.import_dependency_from_zip(
            self._workspace_id, path_to_dependency
        )
        return ProjectImportOperation[Project](
            resp["data"]["location"], self._sal, Project.from_operation
        )

    @classmethod
    def from_import_operation(
        cls, operation: BaseOperation[Workspace], **kwargs: Any
    ) -> Workspace:
        assert isinstance(operation, WorkspaceImportOperation)
        return cls(**kwargs, service=operation._sal)

    @classmethod
    def from_conversion_operation(
        cls, operation: BaseOperation[Workspace], **kwargs: Any
    ) -> Workspace:
        assert isinstance(operation, WorkspaceConversionOperation)
        return cls(**kwargs, service=operation._sal)
