from abc import abstractmethod
from typing import Any, Dict, Iterable, List, Optional, Type, TypeVar

from pyjackson import dumps, loads
from sqlalchemy import Column, DateTime, ForeignKey, Integer, String, Text, UniqueConstraint, Float
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship

from algolink.core.objects import DatasetType
from algolink.core.objects.artifacts import ArtifactCollection
from algolink.core.objects.core import (Buildable, EvaluationResults, EvaluationSet, Image, Model, Pipeline,
                                       PipelineStep, Project, RuntimeEnvironment, RuntimeInstance, Task)
from algolink.core.objects.dataset_source import DatasetSource
from algolink.core.objects.metric import Metric
from algolink.core.objects.requirements import Requirements

SQL_OBJECT_FIELD = '_sqlalchemy_object'


def json_column():
    return Column(Text)


def safe_loads(payload, as_class):
    return loads(payload, Optional[as_class])


def sqlobject(obj):
    return getattr(obj, SQL_OBJECT_FIELD, None)


def update_attrs(obj, **attrs):
    for name, value in attrs.items():
        setattr(obj, name, value)


T = TypeVar('T')
S = TypeVar('S', bound='Attaching')


class Attaching:
    id = ...
    name = ...

    def attach(self, obj):
        setattr(obj, SQL_OBJECT_FIELD, self)
        return obj

    @classmethod
    def from_obj(cls: Type[S], obj: T, new=False) -> S:
        kwargs = cls.get_kwargs(obj)
        existing = sqlobject(obj)
        if not new and existing is not None:
            update_attrs(existing, **kwargs)
            return existing
        return cls(**kwargs)

    @classmethod
    @abstractmethod
    def get_kwargs(cls, obj: T) -> dict:
        pass  # pragma: no cover

    @abstractmethod
    def to_obj(self) -> T:
        pass  # pragma: no cover


Base = declarative_base()


class SProject(Base, Attaching):
    __tablename__ = 'projects'
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String, unique=True, nullable=False)
    author = Column(String, unique=False, nullable=False)
    creation_date = Column(DateTime, unique=False, nullable=False)

    tasks: Iterable['STask'] = relationship("STask", back_populates="project")

    def to_obj(self) -> Project:
        p = Project(self.name, id=self.id, author=self.author, creation_date=self.creation_date)
        for task in self.tasks:
            p._tasks.add(task.to_obj())
        return self.attach(p)

    @classmethod
    def get_kwargs(cls, project: Project) -> dict:
        return dict(id=project.id,
                    name=project.name,
                    author=project.author,
                    creation_date=project.creation_date,
                    tasks=[STask.from_obj(t) for t in project.tasks.values()])


class STask(Base, Attaching):
    __tablename__ = 'tasks'

    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String, unique=False, nullable=False)
    author = Column(String, unique=False, nullable=False)
    creation_date = Column(DateTime, unique=False, nullable=False)
    project_id = Column(Integer, ForeignKey('projects.id'), nullable=False)

    project = relationship("SProject", back_populates="tasks")
    models: Iterable['SModel'] = relationship("SModel", back_populates="task")
    pipelines: Iterable['SPipeline'] = relationship("SPipeline", back_populates='task')
    images: Iterable['SImage'] = relationship("SImage", back_populates='task')

    datasets = Column(Text)
    metrics = Column(Text)
    evaluation_sets = Column(Text)

    __table_args__ = (UniqueConstraint('name', 'project_id', name='tasks_name_and_ref'),)

    def to_obj(self) -> Task:
        task = Task(id=self.id,
                    name=self.name,
                    author=self.author,
                    creation_date=self.creation_date,
                    project_id=self.project_id,
                    datasets=safe_loads(self.datasets, Dict[str, DatasetSource]),
                    metrics=safe_loads(self.metrics, Dict[str, Metric]),
                    evaluation_sets=safe_loads(self.evaluation_sets, Dict[str, EvaluationSet]))
        for model in self.models:
            task._models.add(model.to_obj())

        for pipeline in self.pipelines:
            task._pipelines.add(pipeline.to_obj())

        for image in self.images:
            task._images.add(image.to_obj())
        return self.attach(task)

    @classmethod
    def get_kwargs(cls, task: Task) -> dict:
        return dict(id=task.id,
                    name=task.name,
                    author=task.author,
                    creation_date=task.creation_date,
                    project_id=task.project_id,
                    models=[SModel.from_obj(m) for m in task.models.values()],
                    images=[SImage.from_obj(i) for i in task.images.values()],
                    pipelines=[SPipeline.from_obj(p) for p in task.pipelines.values()],
                    datasets=dumps(task.datasets),
                    metrics=dumps(task.metrics),
                    evaluation_sets=dumps(task.evaluation_sets))


class SModel(Base, Attaching):
    __tablename__ = 'models'

    id = Column(Integer, primary_key=True, autoincrement=True)

    name = Column(String, unique=False, nullable=False)
    author = Column(String, unique=False, nullable=False)
    creation_date = Column(DateTime, unique=False, nullable=False)
    wrapper = Column(Text)

    artifact = Column(Text)
    requirements = Column(Text)
    description = Column(Text)
    params = Column(Text)
    task_id = Column(Integer, ForeignKey('tasks.id'), nullable=False)
    task = relationship("STask", back_populates="models")

    evaluations = Column(Text)
    __table_args__ = (UniqueConstraint('name', 'task_id', name='models_name_and_ref'),)

    def to_obj(self) -> Model:
        model = Model(name=self.name,
                      wrapper_meta=safe_loads(self.wrapper, dict),
                      author=self.author,
                      creation_date=self.creation_date,
                      artifact=safe_loads(self.artifact, ArtifactCollection),
                      requirements=safe_loads(self.requirements, Requirements),
                      description=self.description,
                      params=safe_loads(self.params, Dict[str, Any]),
                      id=self.id,
                      task_id=self.task_id,
                      evaluations=safe_loads(self.evaluations, Dict[str, EvaluationResults]))
        return self.attach(model)

    @classmethod
    def get_kwargs(cls, model: Model) -> dict:
        return dict(id=model.id,
                    name=model.name,
                    author=model.author,
                    creation_date=model.creation_date,
                    wrapper=dumps(model.wrapper_meta),
                    artifact=dumps(model.artifact),
                    requirements=dumps(model.requirements),
                    description=model.description,
                    params=dumps(model.params),
                    task_id=model.task_id,
                    evaluations=dumps(model.evaluations))


class SPipeline(Base, Attaching):
    __tablename__ = 'pipelines'

    id = Column(Integer, primary_key=True, autoincrement=True)

    name = Column(String, unique=False, nullable=False)
    author = Column(String, unique=False, nullable=False)
    creation_date = Column(DateTime, unique=False, nullable=False)
    steps = Column(Text)

    input_data = Column(Text)
    output_data = Column(Text)

    task_id = Column(Integer, ForeignKey('tasks.id'), nullable=False)
    task = relationship("STask", back_populates="pipelines")

    evaluations = Column(Text)
    __table_args__ = (UniqueConstraint('name', 'task_id', name='pipelines_name_and_ref'),)

    def to_obj(self) -> Pipeline:
        pipeline = Pipeline(name=self.name,
                            steps=safe_loads(self.steps, List[PipelineStep]),
                            input_data=safe_loads(self.input_data, DatasetType),
                            output_data=safe_loads(self.output_data, DatasetType),
                            author=self.author,
                            creation_date=self.creation_date,
                            id=self.id,
                            task_id=self.task_id,
                            evaluations=safe_loads(self.evaluations, EvaluationResults))
        return self.attach(pipeline)

    @classmethod
    def get_kwargs(cls, pipeline: Pipeline) -> dict:
        return dict(id=pipeline.id,
                    name=pipeline.name,
                    author=pipeline.author,
                    creation_date=pipeline.creation_date,
                    steps=dumps(pipeline.steps),
                    input_data=dumps(pipeline.input_data),
                    output_data=dumps(pipeline.output_data),
                    task_id=pipeline.task_id,
                    evaluations=dumps(pipeline.evaluations))


class SImage(Base, Attaching):
    __tablename__ = 'images'

    id = Column(Integer, primary_key=True, autoincrement=True)

    name = Column(String, unique=False, nullable=False)
    author = Column(String, unique=False, nullable=False)
    creation_date = Column(DateTime, unique=False, nullable=False)

    task_id = Column(Integer, ForeignKey('tasks.id'), nullable=False)
    task = relationship("STask", back_populates="images")

    environment_id = Column(Integer, ForeignKey('environments.id'), nullable=False)

    params = Column(Text)
    source = Column(Text)

    __table_args__ = (UniqueConstraint('name', 'task_id', name='image_name_and_ref'),)

    def to_obj(self) -> Image:
        image = Image(name=self.name,
                      author=self.author,
                      creation_date=self.creation_date,
                      id=self.id,
                      task_id=self.task_id,
                      params=safe_loads(self.params, Image.Params),
                      source=safe_loads(self.source, Buildable),
                      environment_id=self.environment_id)
        return self.attach(image)

    @classmethod
    def get_kwargs(cls, image: Image) -> dict:
        return dict(id=image.id,
                    name=image.name,
                    author=image.author,
                    creation_date=image.creation_date,
                    task_id=image.task_id,
                    params=dumps(image.params),
                    source=dumps(image.source),
                    environment_id=image.environment_id)


class SRuntimeEnvironment(Base, Attaching):
    __tablename__ = 'environments'

    id = Column(Integer, primary_key=True, autoincrement=True)

    name = Column(String, unique=True, nullable=False)
    author = Column(String, unique=False, nullable=False)
    creation_date = Column(DateTime, unique=False, nullable=False)

    params = Column(Text)

    def to_obj(self) -> RuntimeEnvironment:
        environment = RuntimeEnvironment(
            name=self.name,
            author=self.author,
            creation_date=self.creation_date,
            id=self.id,
            params=safe_loads(self.params, RuntimeEnvironment.Params))
        return self.attach(environment)

    @classmethod
    def get_kwargs(cls, environment: RuntimeEnvironment) -> dict:
        return dict(id=environment.id,
                    name=environment.name,
                    author=environment.author,
                    creation_date=environment.creation_date,
                    params=dumps(environment.params))


class SRuntimeInstance(Base, Attaching):
    __tablename__ = 'instances'

    id = Column(Integer, primary_key=True, autoincrement=True)

    name = Column(String, unique=False, nullable=False)
    author = Column(String, unique=False, nullable=False)
    creation_date = Column(DateTime, unique=False, nullable=False)

    image_id = Column(Integer, ForeignKey('images.id'), nullable=False)
    environment_id = Column(Integer, ForeignKey('environments.id'), nullable=False)

    params = Column(Text)

    __table_args__ = (UniqueConstraint('name', 'image_id', 'environment_id', name='instance_name_and_ref'),)

    def to_obj(self) -> RuntimeInstance:
        instance = RuntimeInstance(
            name=self.name,
            author=self.author,
            creation_date=self.creation_date,
            id=self.id,
            image_id=self.image_id,
            environment_id=self.environment_id,
            params=safe_loads(self.params, RuntimeInstance.Params))
        return self.attach(instance)

    @classmethod
    def get_kwargs(cls, instance: RuntimeInstance) -> dict:
        return dict(id=instance.id,
                    name=instance.name,
                    author=instance.author,
                    creation_date=instance.creation_date,
                    image_id=instance.image_id,
                    environment_id=instance.environment_id,
                    params=dumps(instance.params))

class TExperiment(Base, Attaching):
    __tablename__ = 'experiments'
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String, unique=True, nullable=False)
    author = Column(String, unique=False, nullable=False)
    project_remark = Column(String, unique=False, nullable=False)
    nick_name = Column(String, unique=False, nullable=False)
    del_flag = Column(Integer, default = 0)
    creation_date = Column(DateTime, unique=False, nullable=False)

    submodels: Iterable['TSubmodel'] = relationship("TSubmodel", back_populates="experiment")

    def to_obj(self) -> Project:
        p = Project(self.name, id=self.id, author=self.author, creation_date=self.creation_date)
        for task in self.tasks:
            p._tasks.add(task.to_obj())
        return self.attach(p)

    @classmethod
    def get_kwargs(cls, project: Project) -> dict:
        return dict(id=project.id,
                    name=project.name,
                    author=project.author,
                    creation_date=project.creation_date,
                    tasks=[STask.from_obj(t) for t in project.tasks.values()])

class TSubmodel(Base, Attaching):
    __tablename__ = 'submodels'

    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String, unique=False, nullable=False)
    author = Column(String, unique=False, nullable=False)
    sub_model_remark = Column(String, unique=False, nullable=False)
    nick_name = Column(String, unique=False)
    sub_model_sequence = Column(Integer, nullable=False)
    del_flag = Column(Integer, default = 0)

    creation_date = Column(DateTime, unique=False, nullable=False)
    experiment_id = Column(Integer, ForeignKey('experiments.id'), nullable=False)

    experiment = relationship("TExperiment", back_populates="submodels")
    model_metrics: Iterable['TModelmetric'] = relationship("TModelmetric", back_populates="submodel")
    model_params: Iterable['TModelparam'] = relationship("TModelparam", back_populates="submodel")
    best_result: Iterable['TBestresult'] = relationship("TBestresult", back_populates="submodel")
    #models: Iterable['SModel'] = relationship("SModel", back_populates="task")
    #pipelines: Iterable['SPipeline'] = relationship("SPipeline", back_populates='task')
    #images: Iterable['SImage'] = relationship("SImage", back_populates='task')

    #datasets = Column(Text)
    #metrics = Column(Text)
    #evaluation_sets = Column(Text)

    __table_args__ = (UniqueConstraint('name', 'experiment_id', name='submodel_name_and_ref'),)

    def to_obj(self) -> Task:
        task = Task(id=self.id,
                    name=self.name,
                    author=self.author,
                    creation_date=self.creation_date,
                    project_id=self.project_id,
                    datasets=safe_loads(self.datasets, Dict[str, DatasetSource]),
                    metrics=safe_loads(self.metrics, Dict[str, Metric]),
                    evaluation_sets=safe_loads(self.evaluation_sets, Dict[str, EvaluationSet]))
        for model in self.models:
            task._models.add(model.to_obj())

        for pipeline in self.pipelines:
            task._pipelines.add(pipeline.to_obj())

        for image in self.images:
            task._images.add(image.to_obj())
        return self.attach(task)

    @classmethod
    def get_kwargs(cls, task: Task) -> dict:
        return dict(id=task.id,
                    name=task.name,
                    author=task.author,
                    creation_date=task.creation_date,
                    project_id=task.project_id,
                    models=[SModel.from_obj(m) for m in task.models.values()],
                    images=[SImage.from_obj(i) for i in task.images.values()],
                    pipelines=[SPipeline.from_obj(p) for p in task.pipelines.values()],
                    datasets=dumps(task.datasets),
                    metrics=dumps(task.metrics),
                    evaluation_sets=dumps(task.evaluation_sets))

class TModelmetric(Base, Attaching):
    __tablename__ = 'model_metrics'

    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String, unique=False, nullable=False)
    author = Column(String, unique=False, nullable=False)
    metric_type = Column(String, unique=False, nullable=False)
    metric_value = Column(Float, nullable=False)
    epoch = Column(Integer, nullable=False)

    creation_date = Column(DateTime, unique=False, nullable=False)
    sub_model_id = Column(Integer, ForeignKey('submodels.id'), nullable=False)

    submodel = relationship("TSubmodel", back_populates="model_metrics")
    #models: Iterable['SModel'] = relationship("SModel", back_populates="task")
    #pipelines: Iterable['SPipeline'] = relationship("SPipeline", back_populates='task')
    #images: Iterable['SImage'] = relationship("SImage", back_populates='task')

    #datasets = Column(Text)
    #metrics = Column(Text)
    #evaluation_sets = Column(Text)

    #__table_args__ = (UniqueConstraint('name', 'sub_model_id', name='metric_name_and_ref'),)

    def to_obj(self) -> Task:
        task = Task(id=self.id,
                    name=self.name,
                    author=self.author,
                    creation_date=self.creation_date,
                    project_id=self.project_id,
                    datasets=safe_loads(self.datasets, Dict[str, DatasetSource]),
                    metrics=safe_loads(self.metrics, Dict[str, Metric]),
                    evaluation_sets=safe_loads(self.evaluation_sets, Dict[str, EvaluationSet]))
        for model in self.models:
            task._models.add(model.to_obj())

        for pipeline in self.pipelines:
            task._pipelines.add(pipeline.to_obj())

        for image in self.images:
            task._images.add(image.to_obj())
        return self.attach(task)

    @classmethod
    def get_kwargs(cls, task: Task) -> dict:
        return dict(id=task.id,
                    name=task.name,
                    author=task.author,
                    creation_date=task.creation_date,
                    project_id=task.project_id,
                    models=[SModel.from_obj(m) for m in task.models.values()],
                    images=[SImage.from_obj(i) for i in task.images.values()],
                    pipelines=[SPipeline.from_obj(p) for p in task.pipelines.values()],
                    datasets=dumps(task.datasets),
                    metrics=dumps(task.metrics),
                    evaluation_sets=dumps(task.evaluation_sets))

class TModelparam(Base, Attaching):
    __tablename__ = 'model_params'

    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String, unique=False, nullable=False)
    author = Column(String, unique=False, nullable=False)
    param_type = Column(String, unique=False, nullable=False)
    param_value = Column(Float, nullable=False)

    creation_date = Column(DateTime, unique=False, nullable=False)
    sub_model_id = Column(Integer, ForeignKey('submodels.id'), nullable=False)

    submodel = relationship("TSubmodel", back_populates="model_params")

    #models: Iterable['SModel'] = relationship("SModel", back_populates="task")
    #pipelines: Iterable['SPipeline'] = relationship("SPipeline", back_populates='task')
    #images: Iterable['SImage'] = relationship("SImage", back_populates='task')

    #datasets = Column(Text)
    #metrics = Column(Text)
    #evaluation_sets = Column(Text)

    #__table_args__ = (UniqueConstraint('name', 'sub_model_id', name='metric_name_and_ref'),)

    def to_obj(self) -> Task:
        task = Task(id=self.id,
                    name=self.name,
                    author=self.author,
                    creation_date=self.creation_date,
                    project_id=self.project_id,
                    datasets=safe_loads(self.datasets, Dict[str, DatasetSource]),
                    metrics=safe_loads(self.metrics, Dict[str, Metric]),
                    evaluation_sets=safe_loads(self.evaluation_sets, Dict[str, EvaluationSet]))
        for model in self.models:
            task._models.add(model.to_obj())

        for pipeline in self.pipelines:
            task._pipelines.add(pipeline.to_obj())

        for image in self.images:
            task._images.add(image.to_obj())
        return self.attach(task)

    @classmethod
    def get_kwargs(cls, task: Task) -> dict:
        return dict(id=task.id,
                    name=task.name,
                    author=task.author,
                    creation_date=task.creation_date,
                    project_id=task.project_id,
                    models=[SModel.from_obj(m) for m in task.models.values()],
                    images=[SImage.from_obj(i) for i in task.images.values()],
                    pipelines=[SPipeline.from_obj(p) for p in task.pipelines.values()],
                    datasets=dumps(task.datasets),
                    metrics=dumps(task.metrics),
                    evaluation_sets=dumps(task.evaluation_sets))

class TBestresult(Base, Attaching):
    __tablename__ = 'best_result'

    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String, unique=False, nullable=False)
    author = Column(String, unique=False, nullable=False)
    best_value = Column(Float, nullable=False)
    best_epoch = Column(Integer, nullable=False)

    creation_date = Column(DateTime, unique=False, nullable=False)
    sub_model_id = Column(Integer, ForeignKey('submodels.id'), nullable=False)

    submodel = relationship("TSubmodel", back_populates="best_result")
    #models: Iterable['SModel'] = relationship("SModel", back_populates="task")
    #pipelines: Iterable['SPipeline'] = relationship("SPipeline", back_populates='task')
    #images: Iterable['SImage'] = relationship("SImage", back_populates='task')

    #datasets = Column(Text)
    #metrics = Column(Text)
    #evaluation_sets = Column(Text)

    #__table_args__ = (UniqueConstraint('name', 'sub_model_id', name='metric_name_and_ref'),)

    def to_obj(self) -> Task:
        task = Task(id=self.id,
                    name=self.name,
                    author=self.author,
                    creation_date=self.creation_date,
                    project_id=self.project_id,
                    datasets=safe_loads(self.datasets, Dict[str, DatasetSource]),
                    metrics=safe_loads(self.metrics, Dict[str, Metric]),
                    evaluation_sets=safe_loads(self.evaluation_sets, Dict[str, EvaluationSet]))
        for model in self.models:
            task._models.add(model.to_obj())

        for pipeline in self.pipelines:
            task._pipelines.add(pipeline.to_obj())

        for image in self.images:
            task._images.add(image.to_obj())
        return self.attach(task)

    @classmethod
    def get_kwargs(cls, task: Task) -> dict:
        return dict(id=task.id,
                    name=task.name,
                    author=task.author,
                    creation_date=task.creation_date,
                    project_id=task.project_id,
                    models=[SModel.from_obj(m) for m in task.models.values()],
                    images=[SImage.from_obj(i) for i in task.images.values()],
                    pipelines=[SPipeline.from_obj(p) for p in task.pipelines.values()],
                    datasets=dumps(task.datasets),
                    metrics=dumps(task.metrics),
                    evaluation_sets=dumps(task.evaluation_sets))
