import re
from pathlib import Path
import logging as lg

import click
from mako.template import Template
from caseconverter import snakecase, pascalcase
import inspect
import gray_formatter

from prismarine.prisma_common import get_cluster


HEADER = '''\
# This file is generated by prismarine. Do not modify it.

from prismarine.runtime.dynamo_crud import (
    _query,
    _get_item,
    _update,
    _put_item,
    _delete,
    _scan,
    _save,
    Model,
    without,
    DbNotFound
)

from typing import TypedDict, List, Literal, NotRequired  # noqa
from types import EllipsisType

from {access_module} import get_dynamo_access
'''

DYNAMO_ACCESS = '''\
dynamo = get_dynamo_access()
'''


def build_client(cluster, base_dir: Path, runtime: str, access_module: str | None):
    r_base_dir = base_dir.resolve()
    this_dir = Path(__file__).resolve().parent
    template_file = Path(this_dir, 'model.mako')
    model_template = Template(template_file.read_text())
    module_body = ''
    imports = set()

    for name, model in cluster.models.items():
        file_path = Path(inspect.getfile(model['cls']))
        relative_arr = file_path.relative_to(r_base_dir).as_posix().split('/')
        relative_arr[-1] = relative_arr[-1].replace('.py', '')
        imports.add(('.'.join(relative_arr), model['class_name']))
        source = inspect.getsource(model['cls'])
        source = re.sub(r'@.+?\n', '', source)

        source = re.sub(
            r'[a-zA-Z\d]+\(TypedDict\)', 'UpdateDTO(TypedDict, total=False)', source
        )

        module_body += '\n\n\n' + str(
            model_template.render(
                snake=snakecase,
                pascal=pascalcase,
                ModelClass=f'{name}Model',
                Model=model['class_name'],
                PartitionKey=model['main']['PK'],
                SortKey=model['main']['SK'],
                Indexes=[
                    {
                        'name': index_name,
                        'PartitionKey': index['PK'],
                        'SortKey': index['SK']
                    }
                    for index_name, index in model['indexes'].items()
                ],
                TableName=model['table'],
                DtoLines=source.split('\n')
            )
        )

    # Additional exported classes (but not models)
    for export in cluster.exports.values():
        file_path = Path(inspect.getfile(export['cls']))
        relative_arr = file_path.relative_to(r_base_dir).as_posix().split('/')
        relative_arr[-1] = relative_arr[-1].replace('.py', '')
        imports.add(('.'.join(relative_arr), export['class_name']))

    access_module = access_module or 'prismarine.runtime.dynamo_default'
    header = HEADER.format(access_module=access_module)

    runtime_prefix = f'{runtime}.' if runtime else ''

    for i in imports:
        header += f'from {runtime_prefix}{i[0]} import {i[1]}\n'

    content = header + DYNAMO_ACCESS + module_body
    content = gray_formatter.fix_content(content)
    return content


def write_client(content, base_dir: Path, cluster_package: str):
    client_path = Path(base_dir, cluster_package, 'prismarine_client.py')
    lg.info('Writing ' + str(client_path))
    client_path.write_text(content)


@click.command()
@click.pass_obj
@click.argument('cluster_package')
def generate_client(obj, cluster_package):
    cluster = get_cluster(obj['BaseDir'], cluster_package)

    content = build_client(
        cluster, obj['BaseDir'], obj['Runtime'], obj['DynamoAccessModule']
    )

    write_client(content, obj['BaseDir'], cluster_package)
