# -*- coding: utf-8 -*-

import typing as T
import dataclasses
from pathlib import Path

import sqlalchemy as sa
import sqlalchemy.orm as orm

from .paths import path_sqlite
from .dataset import BaseDataset, download_sqlite

Base = orm.declarative_base()


{% for dataset_metadata in dataset_metadata_list %}
class {{ dataset_metadata.dataset.class_name }}(Base):
    __tablename__ = "{{ dataset_metadata.dataset.tab }}"

    {{ dataset_metadata.dataset.id_col }}: orm.Mapped[{{ dataset_metadata.id_col_type_spec.dc_type }}] = sa.Column(sa.{{ dataset_metadata.id_col_type_spec.sa_type.__name__ }}, primary_key=True)
    {%- for col, type_spec in dataset_metadata.schema.items() %}
    {%- if col != dataset_metadata.dataset.id_col %}
    {{ col }}: orm.Mapped[T.Optional[{{ type_spec.dc_type }}]] = sa.Column(sa.{{ type_spec.sa_type.__name__ }}, nullable=True)
    {%- endif %}
    {%- endfor %}


@dataclasses.dataclass
class {{ dataset_metadata.dataset.class_name }}Data:
    {{ dataset_metadata.dataset.id_col }}: {{ dataset_metadata.id_col_type_spec.dc_type }} = dataclasses.field()
    {%- for col, type_spec in dataset_metadata.schema.items() %}
    {%- if col != dataset_metadata.dataset.id_col %}
    {{ col }}: T.Optional[{{ type_spec.dc_type }}] = dataclasses.field(default=None)
    {%- endif %}
    {%- endfor %}

{% endfor %}

dataset_mapping = {
    {%- for dataset_metadata in dataset_metadata_list %}
    "{{ dataset_metadata.dataset.tab }}": {
        "orm_class": {{ dataset_metadata.dataset.class_name }},
        "data_class": {{ dataset_metadata.dataset.class_name }}Data,
        "id_col": "{{ dataset_metadata.dataset.id_col }}",
    },
    {%- endfor %}
}

{% for dataset_metadata in dataset_metadata_list %}
@dataclasses.dataclass
class {{ dataset_metadata.dataset.class_name }}Dataset(
    BaseDataset[{{ dataset_metadata.dataset.class_name }}, {{ dataset_metadata.dataset.class_name }}Data]
):
    pass
{% endfor %}


@dataclasses.dataclass
class Lookup:
    """
    The main API to access the acore dataframe data. Useful methods are:

    - :meth:`Lookup.${dataset_name}.get <acore_df.dataset.BaseDataset.get>`
    - :meth:`Lookup.${dataset_name}.get_by_kvs <acore_df.dataset.BaseDataset.get_by_kvs>`
    - :meth:`Lookup.${dataset_name}.df <acore_df.dataset.BaseDataset.df>`
    - :meth:`Lookup.${dataset_name}.row_map <acore_df.dataset.BaseDataset.row_map>`
    - :meth:`Lookup.${dataset_name}.name <acore_df.dataset.BaseDataset>`
    - :meth:`Lookup.${dataset_name}.id_col <acore_df.dataset.BaseDataset>`
    - :meth:`Lookup.${dataset_name}.orm_model <acore_df.dataset.BaseDataset>`
    - :meth:`Lookup.${dataset_name}.orm_table <acore_df.dataset.BaseDataset>`
    - :meth:`Lookup.${dataset_name}.data_class <acore_df.dataset.BaseDataset>`
    - :meth:`Lookup.${dataset_name}.engine <acore_df.dataset.BaseDataset>`
    """
    engine: sa.Engine = dataclasses.field()
    {%- for dataset_metadata in dataset_metadata_list %}
    {{ dataset_metadata.dataset.tab }}: {{ dataset_metadata.dataset.class_name }}Dataset = dataclasses.field()
    {%- endfor %}

    @classmethod
    def new(cls, path_sqlite: Path = path_sqlite):
        if path_sqlite.exists() is False:
            download_sqlite(path_sqlite=path_sqlite)

        engine = sa.create_engine(f"sqlite:///{path_sqlite}")
        return cls(
            engine=engine,
            {%- for dataset_metadata in dataset_metadata_list %}
            {{ dataset_metadata.dataset.tab }}={{ dataset_metadata.dataset.class_name }}Dataset(
                name="{{ dataset_metadata.dataset.tab }}",
                id_col="{{ dataset_metadata.dataset.id_col }}",
                orm_model={{ dataset_metadata.dataset.class_name }},
                orm_table={{ dataset_metadata.dataset.class_name }}.__table__,
                data_class={{ dataset_metadata.dataset.class_name }}Data,
                engine=engine,
            ),
            {%- endfor %}
        )
