{% macro init_arg_type(field) %}
{%- if field.conversion is none -%}
    Optional[{{ field.type_for("user")|replace("NonNegativeInt", "int") }}]
{%- else -%}
    Optional[{{ field.conversion.arg_type }}]
{%- endif -%}
{% endmacro %}

{% macro field_assignment(field) %}
{%- if field.conversion is none -%}
    self._{{ field.name }} = {{ field.name }}
{%- else -%}
    self._{{ field.name }} = {{ field.conversion.func }}({{ field.name }})
{%- endif -%}
{% endmacro %}

{% macro field_type_spec(field) %}
{%- set ty = field.type_for("user")|replace("NonNegativeInt", "int") -%}
{%- if ty == "None" -%}
type(None)
{%- else -%}
{{ ty }}
{%- endif -%}
{% endmacro %}

{{ banner }}
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2023 SciCat Project (https://github.com/SciCatProject/scitacean)
# flake8: noqa

"""Base class for Dataset."""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any, Dict, List, Literal, Optional, Tuple, Type, TypeVar, Union

import dateutil.parser

from ._base_model import DatasetType
from ._internal.dataclass_wrapper import dataclass_optional_args
from .datablock import OrigDatablock
from .filesystem import RemotePath
from .model import (
    construct,
    Attachment,
    BaseModel,
    BaseUserModel,
    DownloadDataset,
    DownloadLifecycle,
    Lifecycle,
    Relationship,
    Technique,
)
from .pid import PID


M = TypeVar("M", bound=BaseModel)


def _parse_datetime(x: Optional[Union[datetime, str]]) -> Optional[datetime]:
    if isinstance(x, datetime) or x is None:
        return x
    if x == "now":
        return datetime.now(tz=timezone.utc)
    return dateutil.parser.parse(x)


def _parse_pid(pid: Optional[Union[str, PID]]) -> Optional[PID]:
    if pid is None:
        return pid
    return PID.parse(pid)


def _parse_remote_path(path: Optional[Union[str, RemotePath]]) -> Optional[RemotePath]:
    if path is None:
        return path
    return RemotePath(path)


def _validate_checksum_algorithm(algorithm: Optional[str]) -> Optional[str]:
    if algorithm is None:
        return algorithm
    import hashlib

    if algorithm not in hashlib.algorithms_available:
        raise ValueError(f"Checksum algorithm not recognized: {algorithm}")
    return algorithm


class DatasetBase:
    @dataclass_optional_args(frozen=True, kw_only=True, slots=True)
    class Field:
        name: str
        description: str
        read_only: bool
        required: bool
        scicat_name: str
        type: type
        used_by_derived: bool
        used_by_raw: bool

        def used_by(self, dataset_type: DatasetType) -> bool:
            return (
                self.used_by_raw
                if dataset_type == DatasetType.RAW
                else self.used_by_derived
            )

    _FIELD_SPEC = [
        Field(
            name="type",
            description="Characterize type of dataset, either 'raw' or 'derived'. Autofilled when choosing the proper inherited models.",
            read_only=False,
            required=True,
            scicat_name="type",
            type=DatasetType,
            used_by_derived=True,
            used_by_raw=True,
        ),
{%- for field in spec.user_dset_fields(manual=False)|sort(attribute="name") -%}
        Field(
            name="{{ field.name }}",
            description="{{ field.description }}",
            read_only={{ not field.upload }},
            required={{ field.required }},
            scicat_name="{{ field.scicat_name }}",
            type={{ field_type_spec(field) }},
            used_by_derived={{ field.used_by_derived or not field.upload }},
            used_by_raw={{ field.used_by_raw or not field.upload }},
        ),
{%- endfor -%}
    ]

    __slots__ = (
{% for field in spec.user_dset_fields(manual=False)|sort(attribute="name") %}
        "_{{ field.name }}",
{% endfor %}
        "_meta",
        "_type",
        "_default_checksum_algorithm",
        "_orig_datablocks",
        "_attachments",
    )

    def __init__(self,
        type: Union[DatasetType, Literal["raw", "derived"]],
{%- for field in spec.user_dset_fields(manual=False)|sort(attribute="name")|selectattr("upload") %}
        {{ field.name }}: {{ init_arg_type(field) }} = {{ field.default }},
{%- endfor %}
        meta: Optional[Dict[str, Any]] = None,
        checksum_algorithm: Optional[str] = "blake2b",
    ) -> None:
        self._type = DatasetType(type)
{%- for field in spec.user_dset_fields(manual=False)|sort(attribute="name")|selectattr("upload") %}
        {{ field_assignment(field) }}
{%- endfor %}
{%- for field in spec.user_dset_fields(manual=False)|sort(attribute="name")|rejectattr("upload") %}
        self._{{ field.name }} = None
{%- endfor %}
        self._meta = meta or {}
        self._default_checksum_algorithm = _validate_checksum_algorithm(
            checksum_algorithm
        )
        self._orig_datablocks: List[OrigDatablock] = []
        self._attachments: Optional[List[Attachment]] = []

{%- for field in spec.user_dset_fields(manual=False)|sort(attribute="name") %}
    @property
    def {{ field.name }}(self) -> Optional[{{ field.type_for("user") }}]:
        """{{ field.description }}"""
        return self._{{ field.name }}
{% if field.upload %}
    @{{ field.name }}.setter
    def {{ field.name }}(self, {{ field.name }}: {{ init_arg_type(field) }}) -> None:
        """{{ field.description }}"""
        {{ field_assignment(field) }}
{% endif %}
{%- endfor %}

    @property
    def meta(self) -> Dict[str, Any]:
        """Dict of scientific metadata."""
        return self._meta

    @meta.setter
    def meta(self, meta: Dict[str, Any]) -> None:
        """Dict of scientific metadata."""
        self._meta = meta

    @property
    def type(self) -> DatasetType:
        """Characterize type of dataset, either 'raw' or 'derived'. Autofilled when choosing the proper inherited models."""
        return self._type

    @staticmethod
    def _prepare_fields_from_download(
        download_model: DownloadDataset
    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
        init_args = {}
        read_only = {}
        for field in DatasetBase._FIELD_SPEC:
            if field.read_only:
                read_only["_"+field.name] = getattr(download_model, field.scicat_name)
            else:
                init_args[field.name] = getattr(download_model, field.scicat_name)

        init_args["meta"] = download_model.scientificMetadata
        _convert_download_fields_in_place(init_args, read_only)

        return init_args, read_only

    @staticmethod
    def _convert_readonly_fields_in_place(read_only: Dict[str, Any]) -> None:
    {% for field in spec.user_dset_fields(manual=False)|sort(attribute="name")|selectattr("conversion") %}
        {%- if not field.upload %}
        if ({{ field.name }} := read_only.get("_{{ field.name }}")) is not None:
            read_only["_{{ field.name }}"] = {{ field.conversion.func }}({{ field.name }})
        {%- endif -%}
    {%- endfor %}


def _convert_download_fields_in_place(
    init_args: Dict[str, Any], read_only: Dict[str, Any]
) -> None:
    for mod, key in ((Technique, "techniques"), (Relationship, "relationships")):
        init_args[key] = _list_field_from_download(mod, init_args.get(key))

    DatasetBase._convert_readonly_fields_in_place(read_only)
    if (lifecycle := read_only.get("_lifecycle")) is not None:
        read_only["_lifecycle"] = Lifecycle.from_download_model(
            _as_model(DownloadLifecycle, lifecycle)
        )


def _list_field_from_download(
    mod: Type[BaseUserModel], value: Optional[List[Any]]
) -> Optional[List[BaseUserModel]]:
    if value is None:
        return None
    return [
        mod.from_download_model(_as_model(mod.download_model_type(), item))
        for item in value
    ]


# If validation fails, sub models are not converted automatically by Pydantic.
def _as_model(
    mod: Type[M], value: Union[M, Dict[str, Any]]
) -> M:
    if isinstance(value, dict):
        return construct(mod, **value, _strict_validation=False)
    return value
