{% include '_header.py.jinja' %}
{% from '_utils.py.jinja' import is_async, maybe_async_def, maybe_await, methods, operations with context %}
# -- template client.py.jinja --
from types import TracebackType
from collections import defaultdict
from pydantic import BaseModel, validate_arguments

from . import types, models, errors
from ._types import BaseModelT
from .engine import QueryEngine
from .builder import QueryBuilder


__all__ = (
    'Client',
    'load_env',
)

SCHEMA = '''
{{ datamodel }}
'''


def load_env(*, override: bool = True, **kwargs: Any) -> None:
    """Load environemntal variables from dotenv files

    Loads from the following files relative to the current
    working directory:

    - .env
    - prisma/.env
    """
    from dotenv import load_dotenv

    load_dotenv('.env', override=override, **kwargs)
    load_dotenv('prisma/.env', override=override, **kwargs)


class Client:
    {% for model in dmmf.datamodel.models %}
    {{ model.name.lower() }}: '{{ model.name }}Actions'
    {% endfor %}

    def __init__(
        self,
        *,
        use_dotenv: bool = True,
        log_queries: bool = False,
        # no idea why this type ignore is required, pyright works correctly, mypy does not
        datasource: Optional[types.DatasourceOverride] = None,  # type: ignore[name-defined]
    ) -> None:
        {% for model in dmmf.datamodel.models %}
        self.{{ model.name.lower() }} = {{ model.name }}Actions(self)
        {% endfor %}
        self.__engine: Optional[QueryEngine] = None
        self._active_provider = '{{ datasources[0].active_provider }}'
        self._log_queries = log_queries
        self._datasource = datasource

        if use_dotenv:
            load_env()

    def __del__(self) -> None:
        if self.__engine is not None:
            self.__engine.stop()
            self.__engine = None

    def is_connected(self) -> bool:
        """Returns True if the client is connected to the query engine, False otherwise."""
        return self.__engine is not None

    {{ maybe_async_def }}connect(self, timeout: int = 10) -> None:
        """Connect to the Prisma query engine.

        It is required to call this before accessing data.
        """
        if self.__engine is None:
            self.__engine = QueryEngine(dml=SCHEMA, log_queries=self._log_queries)

        datasources: Optional[List[types.DatasourceOverride]] = None
        if self._datasource is not None:
            ds = self._datasource.copy()
            ds.setdefault('name', '{{ datasources[0].name }}')
            datasources = [ds]

        {{ maybe_await }}self.__engine.connect(
            timeout=timeout,
            datasources=datasources,
        )

    {{ maybe_async_def }}disconnect(self) -> None:
        """Disconnect the Prisma query engine."""
        if self.__engine is not None:
            self.__engine.disconnect()
            {{ maybe_await }}self.__engine.close_session()
            self.__engine = None

    {{ maybe_async_def }}execute_raw(self, query: str, *args: Any) -> int:
        resp = {{ maybe_await }}self._execute(
            operation='{{ operations.execute_raw }}',
            method='{{ methods.execute_raw }}',
            arguments={
                'query': query,
                'parameters': args,
            }
        )
        return int(resp['data']['result'])

    @overload
    {{ maybe_async_def }}query_first(self, query: str, *args: Any) -> Any:
        ...

    @overload
    {{ maybe_async_def }}query_first(self, query: str, *args: Any, model: Type[BaseModelT]) -> Optional[BaseModelT]:
        ...

    {{ maybe_async_def }}query_first(
        self, query: str, *args: Any, model: Optional[Type[BaseModelT]] = None
    ) -> Union[Optional[BaseModelT], Any]:
        if model is not None:
            results = {{ maybe_await }}self.query_raw(query, *args, model=model)
        else:
            results = {{ maybe_await }}self.query_raw(query, *args)

        if not results:
            return None

        return results[0]

    @overload
    {{ maybe_async_def }}query_raw(self, query: str, *args: Any) -> Any:
        ...

    @overload
    {{ maybe_async_def }}query_raw(self, query: str, *args: Any, model: Type[BaseModelT]) -> List[BaseModelT]:
        ...

    {{ maybe_async_def }}query_raw(
        self, query: str, *args: Any, model: Optional[Type[BaseModelT]] = None
    ) -> Union[List[BaseModelT], Any]:
        resp = {{ maybe_await }}self._execute(
            operation='{{ operations.query_raw }}',
            method='{{ methods.query_raw }}',
            arguments={
                'query': query,
                'parameters': args,
            }
        )
        result = resp['data']['result']
        if model is not None:
            return [model.parse_obj(r) for r in result]
        return result

    def batch_(self) -> 'Batch':
        """Returns a context manager for grouping write queries into a single transaction."""
        return Batch(client=self)

    # TODO: don't return Any
    {{ maybe_async_def }}_execute(
        self,
        method: str,
        operation: str,
        arguments: Dict[str, Any],
        model: Optional[str] = None,
        root_selection: Optional[List[str]] = None
    ) -> Any:
        builder = QueryBuilder(
            operation=operation,
            method=method,
            model=model,
            arguments=arguments,
            root_selection=root_selection,
        )
        return {{ maybe_await }}self._engine.request('POST', '/', content=builder.build())

    @property
    def _engine(self) -> QueryEngine:
        engine = self.__engine
        if engine is None:
            raise errors.ClientNotConnectedError()
        return engine


# TODO: this should return the results as well
# TODO: don't require copy-pasting arguments between actions and batch actions
class Batch:
    {% for model in dmmf.datamodel.models %}
    {{ model.name.lower() }}: '{{ model.name }}BatchActions'
    {% endfor %}

    def __init__(self, client: Client) -> None:
        self.__client = client
        self.__queries: List[str] = []
        self._active_provider = client._active_provider
        {% for model in dmmf.datamodel.models %}
        self.{{ model.name.lower() }} = {{ model.name }}BatchActions(self)
        {% endfor %}

    def _add(self, **kwargs: Any) -> None:
        builder = QueryBuilder(**kwargs)
        self.__queries.append(builder.build_query())

    {{ maybe_async_def }}commit(self) -> None:
        """Execute the queries"""
        # TODO: normalise this, we should still call client._execute
        from .builder import dumps

        queries = self.__queries
        self.__queries = []

        payload = {
            'batch': [
                {
                    'query': query,
                    'variables': {},
                }
                for query in queries
            ],
            'transaction': True,
        }
        {{ maybe_await }}self.__client._engine.request('POST', '/', content=dumps(payload))

    {% if is_async %}
    async def __aenter__(self) -> 'Batch':
        return self

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        if exc is None:
            await self.commit()
    {% else %}
    def __enter__(self) -> 'Batch':
        return self

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        if exc is None:
            self.commit()
    {% endif %}

{% for model in dmmf.datamodel.models %}

class {{ model.name }}Actions:
    def __init__(self, client: Client):
        self._client = client

    {{ maybe_async_def }}create(
        self,
        data: types.{{ model.name }}CreateInput,
        include: Optional[types.{{ model.name}}Include] = None
    ) -> models.{{ model.name }}:
        resp = {{ maybe_await }}self._client._execute(
            operation='{{ operations.create }}',
            method='{{ methods.create }}',
            model='{{ model.name }}',
            arguments={
                'data': data,
                'include': include,
            },
        )
        return models.{{ model.name }}.parse_obj(resp['data']['result'])

    {{ maybe_async_def }}create_many(
        self,
        data: List[types.{{ model.name }}CreateWithoutRelationsInput],
        *,
        skip_duplicates: Optional[bool] = None,
    ) -> int:
        if self._client._active_provider == 'sqlite':
            raise errors.UnsupportedDatabaseError('sqlite', 'create_many()')

        resp = {{ maybe_await }}self._client._execute(
            operation='{{ operations.create_many }}',
            method='{{ methods.create_many }}',
            model='{{ model.name }}',
            arguments={
                'data': data,
                'skipDuplicates': skip_duplicates,
            },
            root_selection=['count'],
        )
        return int(resp['data']['result']['count'])

    {{ maybe_async_def }}delete(
        self,
        where: types.{{ model.name }}WhereUniqueInput,
        include: Optional[types.{{ model.name}}Include] = None
    ) -> Optional[models.{{ model.name }}]:
        try:
            resp = {{ maybe_await }}self._client._execute(
                operation='{{ operations.delete }}',
                method='{{ methods.delete }}',
                model='{{ model.name }}',
                arguments={
                    'where': where,
                    'include': include,
                },
            )
        except errors.RecordNotFoundError:
            return None

        return models.{{ model.name }}.parse_obj(resp['data']['result'])

    {{ maybe_async_def }}find_unique(
        self,
        where: types.{{ model.name }}WhereUniqueInput,
        include: Optional[types.{{ model.name}}Include] = None
    ) -> Optional[models.{{ model.name }}]:
        resp = {{ maybe_await }}self._client._execute(
            operation='{{ operations.find_unique }}',
            method='{{ methods.find_unique }}',
            model='{{ model.name }}',
            arguments={
                'where': where,
                'include': include,
            },
        )
        result = resp['data']['result']
        if result is None:
            return None
        return models.{{ model.name }}.parse_obj(result)

    {{ maybe_async_def }}find_many(
        self,
        take: Optional[int] = None,
        skip: Optional[int] = None,
        where: Optional[types.{{ model.name }}WhereInput] = None,
        cursor: Optional[types.{{ model.name }}WhereUniqueInput] = None,
        include: Optional[types.{{ model.name }}Include] = None,
        order: Optional[Union[types.{{ model.name }}OrderByInput, List[types.{{ model.name }}OrderByInput]]] = None,
    ) -> List[models.{{ model.name }}]:
        resp = {{ maybe_await }}self._client._execute(
            operation='{{ operations.find_many }}',
            method='{{ methods.find_many }}',
            model='{{ model.name }}',
            arguments={
                'take': take,
                'skip': skip,
                'where': where,
                'order_by': order,
                'cursor': cursor,
                'include': include,
            },
        )
        return [models.{{ model.name }}.parse_obj(r) for r in resp['data']['result']]

    {{ maybe_async_def }}find_first(
        self,
        skip: Optional[int] = None,
        where: Optional[types.{{ model.name }}WhereInput] = None,
        cursor: Optional[types.{{ model.name }}WhereUniqueInput] = None,
        include: Optional[types.{{ model.name }}Include] = None,
        order: Optional[Union[types.{{ model.name }}OrderByInput, List[types.{{ model.name }}OrderByInput]]] = None,
    ) -> Optional[models.{{ model.name }}]:
        resp = {{ maybe_await }}self._client._execute(
            operation='{{ operations.find_first }}',
            method='{{ methods.find_first }}',
            model='{{ model.name }}',
            arguments={
                'skip': skip,
                'where': where,
                'order_by': order,
                'cursor': cursor,
                'include': include
            },
        )
        result = resp['data']['result']
        if result is None:
            return None
        return models.{{ model.name }}.parse_obj(result)

    {{ maybe_async_def }}update(
        self,
        data: types.{{ model.name }}UpdateInput,
        where: types.{{ model.name }}WhereUniqueInput,
        include: Optional[types.{{ model.name}}Include] = None
    ) -> Optional[models.{{ model.name }}]:
        try:
            resp = {{ maybe_await }}self._client._execute(
                operation='{{ operations["update"] }}',
                method='{{ methods["update"] }}',
                model='{{ model.name }}',
                arguments={
                    'data': data,
                    'where': where,
                    'include': include,
                },
            )
        except errors.RecordNotFoundError:
            return None

        return models.{{ model.name }}.parse_obj(resp['data']['result'])

    {{ maybe_async_def }}upsert(
        self,
        where: types.{{ model.name }}WhereUniqueInput,
        data: types.{{ model.name }}UpsertInput,
        include: Optional[types.{{ model.name}}Include] = None,
    ) -> models.{{ model.name}}:
        resp = {{ maybe_await }}self._client._execute(
            operation='{{ operations.upsert }}',
            method='{{ methods.upsert }}',
            model='{{ model.name }}',
            arguments={
                'where': where,
                'include': include,
                'create': data.get('create'),
                'update': data.get('update'),
            },
        )
        return models.{{ model.name }}.parse_obj(resp['data']['result'])

    {{ maybe_async_def }}update_many(
        self,
        data: types.{{ model.name }}UpdateManyMutationInput,
        where: types.{{ model.name }}WhereInput,
    ) -> int:
        resp = {{ maybe_await }}self._client._execute(
            operation='{{ operations.update_many }}',
            method='{{ methods.update_many }}',
            model='{{ model.name }}',
            arguments={'data': data, 'where': where,},
            root_selection=['count'],
        )
        return int(resp['data']['result']['count'])

    {{ maybe_async_def }}count(
        self,
        take: Optional[int] = None,
        skip: Optional[int] = None,
        where: Optional[types.{{ model.name }}WhereInput] = None,
        cursor: Optional[types.{{ model.name }}WhereUniqueInput] = None,
        order: Optional[Union[types.{{ model.name }}OrderByInput, List[types.{{ model.name }}OrderByInput]]] = None,
    ) -> int:
        # TODO: support select
        resp = {{ maybe_await }}self._client._execute(
            operation='{{ operations.count }}',
            method='{{ methods.count }}',
            model='{{ model.name }}',
            arguments={
                'take': take,
                'skip': skip,
                'where': where,
                'order_by': order,
                'cursor': cursor,
            },
            root_selection=['_count { _all }'],
        )
        return cast(int, resp['data']['result']['_count']['_all'])

    {{ maybe_async_def }}delete_many(
        self,
        where: Optional[types.{{ model.name }}WhereInput] = None
    ) -> int:
        resp = {{ maybe_await }}self._client._execute(
            operation='{{ operations.delete_many }}',
            method='{{ methods.delete_many }}',
            model='{{ model.name }}',
            arguments={'where': where},
            root_selection=['count'],
        )
        return int(resp['data']['result']['count'])


# NOTE: some arguments are meaningless in this context but are included
# for completeness sake
class {{ model.name }}BatchActions:
    def __init__(self, batcher: Batch) -> None:
        self._batcher = batcher

    def create(
        self,
        data: types.{{ model.name }}CreateInput,
        include: Optional[types.{{ model.name}}Include] = None
    ) -> None:
        self._batcher._add(
            operation='{{ operations.create }}',
            method='{{ methods.create }}',
            model='{{ model.name }}',
            arguments={
                'data': data,
                'include': include,
            },
        )

    def create_many(
        self,
        data: List[types.{{ model.name }}CreateWithoutRelationsInput],
        *,
        skip_duplicates: Optional[bool] = None,
    ) -> None:
        if self._batcher._active_provider == 'sqlite':
            raise errors.UnsupportedDatabaseError('sqlite', 'create_many()')

        self._batcher._add(
            operation='{{ operations.create_many }}',
            method='{{ methods.create_many }}',
            model='{{ model.name }}',
            arguments={
                'data': data,
                'skipDuplicates': skip_duplicates,
            },
            root_selection=['count'],
        )

    def delete(
        self,
        where: types.{{ model.name }}WhereUniqueInput,
        include: Optional[types.{{ model.name}}Include] = None,
    ) -> None:
        self._batcher._add(
            operation='{{ operations.delete }}',
            method='{{ methods.delete }}',
            model='{{ model.name }}',
            arguments={
                'where': where,
                'include': include,
            },
        )

    def update(
        self,
        data: types.{{ model.name }}UpdateInput,
        where: types.{{ model.name }}WhereUniqueInput,
        include: Optional[types.{{ model.name}}Include] = None
    ) -> None:
        self._batcher._add(
            operation='{{ operations["update"] }}',
            method='{{ methods["update"] }}',
            model='{{ model.name }}',
            arguments={
                'data': data,
                'where': where,
                'include': include,
            },
        )

    def upsert(
        self,
        where: types.{{ model.name }}WhereUniqueInput,
        data: types.{{ model.name }}UpsertInput,
        include: Optional[types.{{ model.name}}Include] = None,
    ) -> None:
        self._batcher._add(
            operation='{{ operations.upsert }}',
            method='{{ methods.upsert }}',
            model='{{ model.name }}',
            arguments={
                'where': where,
                'include': include,
                'create': data.get('create'),
                'update': data.get('update'),
            },
        )

    def update_many(
        self,
        data: types.{{ model.name }}UpdateManyMutationInput,
        where: types.{{ model.name }}WhereInput,
    ) -> None:
        self._batcher._add(
            operation='{{ operations.update_many }}',
            method='{{ methods.update_many }}',
            model='{{ model.name }}',
            arguments={'data': data, 'where': where,},
            root_selection=['count'],
        )

    def delete_many(
        self,
        where: Optional[types.{{ model.name }}WhereInput] = None,
    ) -> None:
        self._batcher._add(
            operation='{{ operations.delete_many }}',
            method='{{ methods.delete_many }}',
            model='{{ model.name }}',
            arguments={'where': where},
            root_selection=['count'],
        )


{% endfor %}
