import re
from pathlib import Path
import logging as lg
from importlib.metadata import version
from typing import Literal

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

from prismarine.prisma_common import get_cluster


type ExtraImport = list[tuple[str, str]]


HEADER_TEMPLATE = '''\
# This file is generated by prismarine {version}. Do not modify it.

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

{type_imports}
from types import EllipsisType

from {access_module} import get_dynamo_access
{helper_block}
'''

DYNAMO_ACCESS = '''\
# Create a DynamoDB resource
dynamo = get_dynamo_access()
'''

PYDANTIC_HELPERS = '''\

from typing import Any, Optional, get_args, get_origin
from pydantic import BaseModel, Field, create_model
from pydantic.fields import PydanticUndefined


def _is_optional_type(annotation: Any):
    origin = get_origin(annotation)
    if origin:
        return type(None) in get_args(annotation)

    return annotation is type(None)


def _make_optional(annotation: Any):
    if _is_optional_type(annotation):
        return annotation

    try:
        return annotation | None  # type: ignore[operator]
    except TypeError:
        return Optional[annotation]


def _build_pydantic_update_model(model_cls):
    field_definitions = {}
    for name, field in model_cls.model_fields.items():
        annotation = _make_optional(field.annotation or Any)
        field_is_required = field.is_required() if callable(field.is_required) else field.is_required
        alias = field.alias if field.alias and field.alias != name else None

        if field_is_required:
            default = None
        else:
            default = field.default if field.default is not PydanticUndefined else None

        default_factory = None if field_is_required else field.default_factory
        field_kwargs = {}
        if alias:
            field_kwargs['alias'] = alias

        if default_factory is not None:
            field_definitions[name] = (
                annotation,
                Field(default_factory=default_factory, **field_kwargs)
            )
        elif field_kwargs:
            field_definitions[name] = (
                annotation,
                Field(default=default, **field_kwargs)
            )
        else:
            field_definitions[name] = (annotation, default)

    return create_model(
        f'{model_cls.__name__}UpdateDTO',
        __base__=BaseModel,
        __module__=__name__,
        **field_definitions
    )


def _dump_model(instance):
    if isinstance(instance, BaseModel):
        return instance.model_dump(mode='python', by_alias=True)

    return instance


def _dump_update_model(instance):
    if isinstance(instance, BaseModel):
        return instance.model_dump(mode='python', by_alias=True, exclude_unset=True)

    return instance


def _load_model(model_cls, value):
    if isinstance(value, model_cls):
        return value

    return model_cls.model_validate(value)


def _load_models(model_cls, items):
    return [_load_model(model_cls, item) for item in items]
'''


def build_client(
    cluster,
    base_dir: Path,
    runtime: str | None,
    access_module: str | None,
    extra_imports: ExtraImport = [],
    *,
    model_library: Literal['typed-dict', 'pydantic'] = 'typed-dict'
):
    '''
    Build a Prismarine client for a given cluster package.

    Args:
        cluster: The Cluster object.
        base_dir: The base directory of the project.
        runtime: The runtime package where the models are defined.
        access_module: The database access module to use.
        extra_imports: Extra imports to add to the client in format:
            [('path.to.module', 'ClassName')]
        model_library: Controls the code that is generated for data models.
            Supported values: 'typed-dict' (default) or 'pydantic'.
    '''
    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(encoding='utf-8'))
    module_body = ''
    imports = set()
    use_pydantic = model_library == 'pydantic'

    if use_pydantic:
        try:
            from pydantic import BaseModel as _BaseModel  # type: ignore import-not-found
        except ImportError as exc:
            raise RuntimeError(
                'Pydantic support requires the optional dependency. '
                'Install it via "pip install prismarine[pydantic]".'
            ) from exc

    for model in sorted(cluster.models, key=lambda x: x['class_name']):
        name = model['name'] or model['class_name']
        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)

        if use_pydantic:
            if not issubclass(model['cls'], _BaseModel):
                raise TypeError(
                    f'Model {model["class_name"]} must inherit from pydantic.BaseModel '
                    'when model_library="pydantic"'
                )
            dto_lines: list[str] = []
        else:
            source = re.sub(
                r'[a-zA-Z\d]+\(TypedDict\)', 'UpdateDTO(TypedDict, total=False)', source
            )
            dto_lines = source.split('\n')

        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'],
                UsePydantic=use_pydantic,
                Indexes=[
                    {
                        'name': index_name,
                        'PartitionKey': index['PK'],
                        'SortKey': index['SK']
                    }
                    for index_name, index
                    in sorted(model['indexes'].items(), key=lambda x: x[0])
                ],
                TableName=model['table'],
                DtoLines=dto_lines
            )
        )

    # 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'
    if use_pydantic:
        type_imports = 'from typing import List, Literal  # noqa'
        helper_block = PYDANTIC_HELPERS
    else:
        type_imports = 'from typing import TypedDict, List, Literal, NotRequired  # noqa'
        helper_block = ''

    header = HEADER_TEMPLATE.format(
        access_module=access_module,
        version=version('prismarine'),
        type_imports=type_imports,
        helper_block=helper_block
    )

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

    header += '# Model classes imports\n'
    for i in sorted(imports, key=lambda x: x[1]):
        header += f'from {runtime_prefix}{i[0]} import {i[1]}  # noqa: F401, E402\n'

    if extra_imports:
        header += '# Extra imports\n'
        for i in sorted(extra_imports, key=lambda x: x[1]):
            header += f'from {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)


def generate_client(
    base_dir: Path,
    cluster_package: str,
    *,
    runtime: str | None,
    access_module: str | None,
    extra_imports: ExtraImport = [],
    model_library: Literal['typed-dict', 'pydantic'] = 'typed-dict'
):
    '''
    Generate a Prismarine client for a given cluster package.

    Args:
        base_dir: The base directory of the project.
        cluster_package: The name of the cluster package.
        runtime: The runtime package where the models are defined.
        access_module: The database access module to use.
        extra_imports: Extra imports to add to the client in format:
            [('path.to.module', 'ClassName')]
        model_library: Controls the code that is generated for data models.
            Supported values: 'typed-dict' (default) or 'pydantic'.
    '''
    cluster = get_cluster(base_dir, cluster_package)

    content = build_client(
        cluster, base_dir, runtime, access_module,
        extra_imports=extra_imports,
        model_library=model_library
    )

    write_client(content, base_dir, cluster_package)
