{% include '_header.py.jinja' %}
# -- template models.py.jinja --
import os
from pydantic import BaseConfig, BaseModel, Field

from . import types, enums, errors, fields
from .generator import partial_models_ctx, PartialModelField


class Config(BaseConfig):
    use_enum_values: bool = True
    allow_population_by_field_name: bool = True


_created_partial_types: Set[str] = set()


{% for model in dmmf.datamodel.models %}
class {{ model.name }}(BaseModel):
    {% for field in model.all_fields %}
    {{ field.name }}:
            {%- if field.is_required and not field.relation_name -%}
                {{ ' ' }}{{ field.python_type }}
            {% else -%}
                {{ ' ' }}Optional[{{ field.python_type }}]
            {% endif %}
    {% endfor %}

    Config = Config

    @staticmethod
    def create_partial(
        name: str,
        include: Optional[Iterable['types.{{ model.name }}Keys']] = None,
        exclude: Optional[Iterable['types.{{ model.name }}Keys']] = None,
        required: Optional[Iterable['types.{{ model.name }}Keys']] = None,
        optional: Optional[Iterable['types.{{ model.name }}Keys']] = None,
        relations: Optional[Mapping['types.{{ model.name }}RelationalFieldKeys', str]] = None,
    ) -> None:
        if not os.environ.get('PRISMA_GENERATOR_INVOCATION'):
            raise RuntimeError(
                'Attempted to create a partial type outside of client generation.'
            )

        if name in _created_partial_types:
            raise ValueError(f'Partial type "{name}" has already been created.')

        if include is not None and exclude is not None:
            raise TypeError(f'Exclude and include are mutually exclusive.')

        if required and optional:
            shared = set(required) & set(optional)
            if shared:
                raise ValueError(f'Cannot make the same field(s) required and optional {shared}')

        fields: Dict['types.{{ model.name }}Keys', PartialModelField] = {}

        try:
            if include:
                for field in include:
                    fields[field] = _{{ model.name }}_fields[field]
            elif exclude:
                for field in exclude:
                    if field not in _{{ model.name }}_fields:
                        raise KeyError(field)

                fields = {
                    key: data
                    for key, data in _{{ model.name }}_fields.items()
                    if key not in exclude
                }
            else:
                fields = _{{ model.name }}_fields.copy()

            if required:
                for field in required:
                    fields[field] = fields[field].copy()
                    fields[field]['optional'] = False

            if optional:
                for field in optional:
                    fields[field] = fields[field].copy()
                    fields[field]['optional'] = True

            if relations:
                {% if model.has_relational_fields %}
                for field, type_ in relations.items():
                    if field not in _{{ model.name }}_relational_fields:
                        raise errors.UnknownRelationalFieldError('{{ model.name }}', field)

                    # TODO: this method of validating types is not ideal
                    # as it means we cannot two create partial types that
                    # reference each other
                    if type_ not in _created_partial_types:
                        raise ValueError(
                            f'Unknown partial type: "{type_}". '
                            f'Did you remember to generate the {type_} type before this one?'
                        )

                    # TODO: support non prisma.partials models
                    info = fields[field]
                    if info['is_list']:
                        info['type'] = f'List[\'partials.{type_}\']'
                    else:
                        info['type'] = f'\'partials.{type_}\''
                {% else %}
                raise ValueError('Model: "{{ model.name }}" has no relational fields.')
                {% endif %}
        except KeyError as exc:
            raise ValueError(
                f'{exc.args[0]} is not a valid {{ model.name }} / {name} field.'
            ) from None

        models = partial_models_ctx.get()

        # mypy does not like this as we are assigning a
        # Dict[Literal[str]] to a Dict[str] but this is fine
        models[name] = fields  # type: ignore[assignment]
        partial_models_ctx.set(models)
        _created_partial_types.add(name)


{% endfor %}

{% for model in dmmf.datamodel.models %}
{% if model.has_relational_fields -%}
    _{{ model.name }}_relational_fields: Set[str] = {
        {% for field in model.relational_fields %}
        '{{ field.name }}',
        {% endfor %}
    }
{% else -%}
    _{{ model.name }}_relational_fields: Set[str] = set()  # pyright: reportUnusedVariable=false
{% endif %}
_{{ model.name }}_fields: Dict['types.{{ model.name }}Keys', PartialModelField] = {
    {% for field in model.all_fields %}
    '{{ field.name }}': {
        'name': '{{ field.name }}',
        'is_list': {{ field.is_list }},
        'optional': {{ field.is_optional }},
        'type': {{ field.python_type_as_string }},
    },
    {% endfor %}
}

{% endfor %}


# we have to import ourselves as relation types are namespaced to models
# e.g. models.Post
from . import models

# required to support relationships between models
{% for model in dmmf.datamodel.models %}
{{ model.name }}.update_forward_refs()
{% endfor %}
