"""{{module}}

This module has been generated with SqlPyGen.
"""

from pprint import pprint
from typing import NewType, Iterable, cast

import sqlite3
ConnectionType = sqlite3.Connection

{% macro comma_sep(loop) -%}
{% if not loop.last %}, {% endif %}
{%- endmacro %}

{% for stmt in import_stmts %}
from {{ stmt.module }} import {% for name in stmt.names %}{{ name }}{{ comma_sep(loop) }}{% endfor %}

{% endfor %}

{% for type in types %}
{{ type.name }} = NewType("{{ type.name }}", {{ type.old_type }})
{% endfor %}

{% macro expand_params(params) -%}
{% for param in params -%}
{{ param.name }}: {{ param.type }}{{ comma_sep(loop) }}
{%- endfor %}
{%- endmacro %}

{% macro do_expand_return(return_) -%}
list[tuple[{% for etype in return_ %}{{ etype }}{{ comma_sep(loop) }}{% endfor %}]]
{%- endmacro %}

{% macro expand_return(return_) -%}
{% if return_ %}{{ do_expand_return(return_) }}{% else %}None{% endif %}
{%- endmacro %}

{% macro do_expand_return_nt(return_) -%}
Iterable[tuple[{% for etype in return_ %}{{ etype }}{{ comma_sep(loop) }}{% endfor %}]]
{%- endmacro %}

{% macro expand_return_nt(return_) -%}
{% if return_ %}{{ do_expand_return_nt(return_) }}{% else %}None{% endif %}
{%- endmacro %}

SCHEMA = {}
{% for schema in schemas %}
SCHEMA["{{ schema.name }}"] = """
{{ schema.sql }}
"""

{% endfor %}

QUERY = {}
{% for query in queries %}
QUERY["{{ query.name }}"] = """
{{ query.sql }}
"""

{% endfor %}

def create_schema(connection: ConnectionType) -> None:
    """Create the table schema."""
    with connection:
        cursor = connection.cursor()

        {% for schema in schemas %}
        try:
            sql = SCHEMA["{{ schema.name }}"]

            cursor.execute(sql)
        except Exception as e:
            raise RuntimeError("An unexpected exception occurred when creating schema: {{ schema.name }}") from e
        {% endfor %}

{% for query in queries %}
def {{ query.name }}(connection: ConnectionType, {{ expand_params(query.params) }}) -> {{ expand_return(query.return_) }}:
    """Query {{query.name}} with transaction."""
    {% if query.params %}
    query_args = {
        {% for param in query.params %}
        "{{ param.name }}": {{ param.name }}{{ comma_sep(loop) }}
        {% endfor %}
    }
    {% endif %}

    with connection:
        cursor = connection.cursor()
        try:
            sql = QUERY["{{ query.name }}"]

            {% if query.params %}
            cursor.execute(sql, query_args)
            {% else %}
            cursor.execute(sql)
            {% endif %}

            {% if query.return_ %}
            return cursor.fetchall()
            {% endif %}
        except Exception as e:
            raise RuntimeError("An unexpected exception occurred while executing query: {{ query.name }}") from e

def {{ query.name }}_nt(connection: ConnectionType, {{ expand_params(query.params) }}) -> {{ expand_return_nt(query.return_) }}:
    """Query {{query.name}} no implied transaction."""
    {% if query.params %}
    query_args = {
        {% for param in query.params %}
        "{{ param.name }}": {{ param.name }}{{ comma_sep(loop) }}
        {% endfor %}
    }
    {% endif %}

    cursor = connection.cursor()
    try:
        sql = QUERY["{{ query.name }}"]

        {% if query.params %}
        cursor.execute(sql, query_args)
        {% else %}
        cursor.execute(sql)
        {% endif %}

        {% if query.return_ %}
        return cast({{ expand_return_nt(query.return_) }}, cursor)
        {% endif %}
    except Exception as e:
        raise RuntimeError("An unexpected exception occurred while executing query: {{ query.name }}") from e

{% endfor %}

def explain_queries() -> None:
    connection = sqlite3.connect(":memory:")
    create_schema(connection)

    with connection:
        cursor = connection.cursor()

        {% for query in queries %}
        try:
            sql = QUERY["{{ query.name }}"]
            sql = "EXPLAIN " + sql

            {% if query.params %}
            query_args = {
                {% for param in query.params %}
                "{{ param.name }}": None{{ comma_sep(loop) }}
                {% endfor %}
            }
            {% endif %}

            {% if query.params %}
            cursor.execute(sql, query_args)
            {% else %}
            cursor.execute(sql)
            {% endif %}

            print("Query explanation for {{ query.name }}")
            print("-" * 80)
            pprint(cursor.fetchall())
        except Exception as e:
            raise RuntimeError("An unexpected exception occurred while executing query plan for: {{ query.name }}") from e

        {% endfor %}

if __name__ == "__main__":
    explain_queries()
