import abc
import re
import time
from typing import Any, Callable, Optional, Sequence, Tuple

from jinja2 import BaseLoader, TemplateNotFound, Undefined
from jinja2.environment import Environment
from jinja2.sandbox import SandboxedEnvironment

from .packages import BasePackageRegistry
from .apptypes import (
    PackageAlias,
    PackageReleaseId,
    Reference,
    RefRenderingResult,
    RenderRequest,
    RenderStats,
    RenderingResultSuccess,
)


class InvalidRefError(TypeError):
    """
    Error thrown when something invalid (e.g. an empty string or jinja2.Undefined) is
    passed as a parameter to a ref, ref_link, or similar call.
    """


def safe_literal(value: str, chars: Sequence[str] = ("'", '"')):
    # make this a raw docstring so it renders unambiguously as a docstring and as a code comment
    r"""Implement the "escape" filter. Chars is an optional parameter to the filter that can be used to change what is
    filtered:

    {{ '"foo"' | safe_literal }} -> \"foo\"

    {{ '`foo`' | safe_literal(["'", '"', '`'])}} -> \`foo\`
    """
    safe_value = value
    # Replace characters
    for char in chars:
        reg_pattern = r"(?<!\\)(?:\\{2})*" + char
        safe_value = re.sub(reg_pattern, f"\\{char}", safe_value)
    return safe_value


def postprocess_rendered_text(rendered_template: str):
    """
    Postprocess the result of rendering a template. Strip preceding and trailing spaces.
    """
    rendered_template = rendered_template.strip(" ")
    return rendered_template


class BaseRefCollector(abc.ABC):
    """The base definition of the "Ref collector" type that the environment expects."""

    @abc.abstractmethod
    def collect(self) -> RefRenderingResult:
        """Collect the results from rendering into lists of unique references by type of reference"""

    @abc.abstractmethod
    def get_exports(self, context: dict[str, Any]) -> dict[str, Any]:
        """Get a dict of things to export into the rendering context"""


class CoalesceJinjaEnvironment(SandboxedEnvironment):
    """Coalesce's jinja environment, which supplies some extra filters"""

    def __init__(self, registry: BasePackageRegistry, render_platform_kind: str) -> None:
        super().__init__()
        self.filters["e"] = safe_literal
        self.filters["escape"] = safe_literal
        # the shared common package registry
        self.package_registry: BasePackageRegistry = registry
        # the alias -> release_id mapping for this environment. This will be set during rendering,
        # but set it here for typing
        self.alias_ids: dict[PackageAlias, PackageReleaseId] = {}
        self.loader = Loader(self.package_registry, self.alias_ids)
        self.render_platform_kind = render_platform_kind

    def get_ref_collector(self):
        """Get the ref collector"""
        # databricks and bigquery use backticks to quote identifiers
        # fabric uses straight brackets as identifiers
        # snowflake uses double quotes to quote identifiers
        if self.render_platform_kind == "databricks" or self.render_platform_kind == "bigquery":
            opening_mark_for_identifiers = "`"
            closing_mark_for_identifiers = "`"
        elif self.render_platform_kind == "fabric":
            opening_mark_for_identifiers = "["
            closing_mark_for_identifiers = "]"
        else:
            opening_mark_for_identifiers = '"'
            closing_mark_for_identifiers = '"'
        return RefCollector(opening_mark_for_identifiers, closing_mark_for_identifiers, self.render_platform_kind)

    def join_path(self, template: str, parent: str) -> str:
        """Join a template with the parent.  By default all the lookups are relative to the loader root so this method
        returns the `template` parameter unchanged, but if the paths should be relative to the parent template, this
        function can be used to calculate the real template name.

        Subclasses may override this method and implement template path joining here.
        """
        if isinstance(template, Undefined):
            raise ValueError(UNDEFINED_TEMPLATE_MESSAGE)
        try:
            parent = self.alias_ids[PackageAlias(parent)]
        except KeyError:
            parent = PackageReleaseId(parent)
        try:
            package = self.package_registry.get_package(parent)
        except KeyError:
            raise TemplateNotFound(template, f'parent package "{parent}" not found in registry') from None
        try:
            return package.dependencies[PackageAlias(template)]
        except KeyError:
            raise TemplateNotFound(template, f'Could not find package "{template}" in "{parent}"') from None

    def render_coalesce_template_v2(self, request: RenderRequest) -> RenderingResultSuccess:
        """Render the coalesce template based upon the render request"""
        # the loader holds a reference to this, so on render requests we want to clear the alias_ids and update it
        # Due to the caller's behavior, we actually know that alias_ids is always empty, but this is more maintainable
        self.alias_ids.clear()
        self.alias_ids.update(self.package_registry.alias_ids_from_packages(request.dependencies))
        values = {**request.template_vars, **request.template_vars.get("nodeMetadata", {})}
        rendered, refs, stats = self.render_until(
            template=request.template,
            values=values,
            # it'd be nice to port hash() and friends to python as well, or something, but that's hard!
            macros=request.global_macros,
            render_depth=request.render_depth,
        )
        rendered = postprocess_rendered_text(rendered)
        return self.create_render_result(rendered, refs, stats)

    def render_until(
        self, template: str, values: dict[str, Any], macros: str, render_depth: int
    ) -> Tuple[str, RefRenderingResult, RenderStats]:
        """
        Repeatedly render templates until there are no more jinja control characters to strip out. This also tracks
        some relevant state on the environment object itself.
        """
        references = self.get_ref_collector()
        depth: int = 0

        render_start_time = time.time()
        curr = template

        # add ref, ref_link, and ref_no_link to the context. These are template-execution-specific.
        values.update(references.get_exports(values))

        while depth < render_depth:
            depth += 1
            full_template = macros + curr
            curr = self.from_string(full_template).render(values)
            # note: we don't check for `{#`/`#}`
            if not any(substring in curr for substring in ["{{", "}}", "{%", "%}"]):
                break

        render_time = int((time.time() - render_start_time) * 1000)
        depth_str = f"{depth}" if depth < render_depth else f"{depth}-MAX"

        stats: RenderStats = {"depth": depth_str, "duration_ms": render_time}

        return curr, references.collect(), stats

    def create_render_result(
        self, rendered_sql: str, refs: RefRenderingResult, stats: RenderStats
    ) -> RenderingResultSuccess:
        """Given the rendered sql, use other information stored on this object to generate the desired output."""
        return {
            "runDetails": {
                "compiledSQL": rendered_sql,
                "exported_refs": refs.exported_refs,
                "invalid_exported_refs": refs.invalid_exported_refs,
                "ref_usage": {
                    "ref": refs.ref,
                    "ref_no_link": refs.ref_no_link,
                    "ref_link": refs.ref_link,
                    "invalid_ref": refs.invalid_ref,
                    "invalid_ref_no_link": refs.invalid_ref_no_link,
                    "invalid_ref_link": refs.invalid_ref_link,
                },
            },
            "render_stats": stats,
        }


UNDEFINED_TEMPLATE_MESSAGE = (
    "Template name was undefined - did you quote it? Import sources must be of the form '{% import \"x\" as x %}' or "
    "'{% from \"x\" import thing %}'"
)


def _stripped_value(src: dict[str, Any], key: str) -> str:
    value = src.get(key)
    if isinstance(value, str):
        return value.strip()
    return ""


def _get_node_ref_name_from_storage(storage: dict[str, Any], node_name: str, open_mark: str, close_mark: str) -> str:
    database = _stripped_value(storage, "database")
    schema = _stripped_value(storage, "schema")
    return f"{open_mark}{database}{close_mark}.{open_mark}{schema}{close_mark}.{open_mark}{node_name}{close_mark}"


def _get_node_ref_schema_name_from_storage(
    storage: dict[str, Any], node_name: str, open_mark: str, close_mark: str
) -> str:
    schema = _stripped_value(storage, "schema")
    return f"{open_mark}{schema}{close_mark}.{open_mark}{node_name}{close_mark}"


def _find_storage_in_metadata(context: dict[str, Any], location_name: str) -> Optional[dict[str, Any]]:
    """Given the node metadata, return the storage matching the given location name"""
    node_metadata = context.get("nodeMetadata", {})
    storage_mapping = node_metadata.get("mapping", {})
    return next((s for s in storage_mapping.values() if s.get("locationName") == location_name), None)


class Loader(BaseLoader):
    """An implementation of the jinja Loader type that uses a package registry"""

    def __init__(self, package_registry: BasePackageRegistry, name_ids: dict[PackageAlias, PackageReleaseId]) -> None:
        self.package_registry: BasePackageRegistry = package_registry
        self.name_ids: dict[PackageAlias, PackageReleaseId] = name_ids

    def get_source(
        self, environment: Environment, template: str
    ) -> Tuple[str, Optional[str], Optional[Callable[[], bool]]]:
        """Get the source text for the loader"""
        if isinstance(template, Undefined):
            raise ValueError(UNDEFINED_TEMPLATE_MESSAGE)
        try:
            template_id = self.name_ids[PackageAlias(template)]
        except KeyError:
            # if it's not in name_ids, assume the argument was actually an ID, not a name
            template_id = PackageReleaseId(template)
        try:
            package = self.package_registry.get_package(template_id)
        except KeyError:
            raise TemplateNotFound(template, f"Template {template} was not registered") from None
        return package.macros, None, lambda: True

    def list_templates(self) -> list[str]:
        """Return a list of the templates this loader knows about"""
        return list(self.name_ids)


class RefCollector(BaseRefCollector):
    """
    Collects ref usage during template evaluation and puts them into local storage. Provides callables that can be used
    for export_ref and invalid_export_ref in the rendering environment
    """

    def __init__(
        self,
        opening_mark_for_identifiers: str,
        closing_mark_for_identifiers: str,
        render_platform_kind: str,
        context_prefix: str = "_coalesce_internal_",
    ) -> None:
        self.exported_t_ref: set[tuple[str, str]] = set()
        self.exported_t_ref_no_link: set[tuple[str, str]] = set()
        self.exported_t_ref_link: set[tuple[str, str]] = set()

        self.render_platform_kind: str = render_platform_kind
        self.invalid_exported_t_ref: set[tuple[str, str]] = set()
        self.invalid_exported_t_ref_no_link: set[tuple[str, str]] = set()
        self.invalid_exported_t_ref_link: set[tuple[str, str]] = set()
        self.context_prefix: str = context_prefix
        self.opening_mark_for_identifiers: str = opening_mark_for_identifiers
        self.closing_mark_for_identifiers: str = closing_mark_for_identifiers

    def _set_context_path(self, context: dict[str, Any], key: str, value: Any):
        context.setdefault(self.context_prefix, {})[key] = value

    def get_ref(self, context: dict[str, Any]) -> Callable[[str, str], str]:
        """Make a 'ref' implementation for this context"""

        def ref(location_name: str, node_name: str) -> str:
            storage = _find_storage_in_metadata(context, location_name)
            if storage:
                self.exported_t_ref.add((location_name, node_name))
                return _get_node_ref_name_from_storage(
                    storage, node_name, self.opening_mark_for_identifiers, self.closing_mark_for_identifiers
                )
            else:
                self.invalid_exported_t_ref.add((location_name, node_name))
                return ""

        return ref

    def get_ref_link(self, context: dict[str, Any]) -> Callable[[str, str], str]:
        """Make a 'ref_link' implementation for this context"""

        def ref_link(location_name: str, node_name: str) -> str:
            storage = _find_storage_in_metadata(context, location_name)
            if storage:
                self.exported_t_ref_link.add((location_name, node_name))
            else:
                self.invalid_exported_t_ref_link.add((location_name, node_name))
            return ""

        return ref_link

    def get_ref_no_link(self, context: dict[str, Any]) -> Callable[[str, str], str]:
        """Make a 'ref_no_link' implementation for this context."""

        def ref_no_link(location_name: str, node_name: str, *, include_db: bool = True) -> str:
            storage = _find_storage_in_metadata(context, location_name)
            if storage:
                self.exported_t_ref_no_link.add((location_name, node_name))
                if include_db:
                    return _get_node_ref_name_from_storage(
                        storage, node_name, self.opening_mark_for_identifiers, self.closing_mark_for_identifiers
                    )
                else:
                    return _get_node_ref_schema_name_from_storage(
                        storage, node_name, self.opening_mark_for_identifiers, self.closing_mark_for_identifiers
                    )
            else:
                self.invalid_exported_t_ref_no_link.add((location_name, node_name))
                return ""

        return ref_no_link

    def _validate_field(self, param_name: str, field_val: Any) -> None:
        if isinstance(field_val, Undefined):
            # pylint: disable=protected-access
            raise InvalidRefError(f"Invalid {param_name}: {field_val._undefined_name} was not defined.")

        if not (field_val and isinstance(field_val, str)):
            raise InvalidRefError(f"Invalid {param_name}: was not a non-empty string")

    def _create_reflist(self, refset: set[tuple[str, str]]) -> list[Reference]:
        """Convert a reference set to a list of Reference dicts"""
        reflist: list[Reference] = []
        for location_name, node_name in refset:
            # If we manage to get here with a ref that is either an empty string or a
            # reference to something that does not exist, that is certainly an error of
            # either ours, or more likely, the user.
            self._validate_field("location", location_name)
            self._validate_field("node name", node_name)

            reflist.append({"location_name": location_name, "node_name": node_name})

        return reflist

    def collect(self) -> RefRenderingResult:
        """Create a list of dicts for each unique reference and return it."""
        exported_refs: list[Reference] = self._create_reflist(self.exported_t_ref | self.exported_t_ref_link)

        exported_t_ref: list[Reference] = self._create_reflist(self.exported_t_ref)

        exported_t_ref_no_link: list[Reference] = self._create_reflist(self.exported_t_ref_no_link)

        exported_t_ref_link: list[Reference] = self._create_reflist(self.exported_t_ref_link)

        invalid_exported_refs: list[Reference] = self._create_reflist(
            self.invalid_exported_t_ref | self.invalid_exported_t_ref_link
        )

        invalid_exported_t_ref: list[Reference] = self._create_reflist(self.invalid_exported_t_ref)

        invalid_exported_t_ref_no_link: list[Reference] = self._create_reflist(self.invalid_exported_t_ref_no_link)

        invalid_exported_t_ref_link: list[Reference] = self._create_reflist(self.invalid_exported_t_ref_link)

        return RefRenderingResult(
            exported_refs,
            invalid_exported_refs,
            exported_t_ref,
            exported_t_ref_no_link,
            exported_t_ref_link,
            invalid_exported_t_ref,
            invalid_exported_t_ref_no_link,
            invalid_exported_t_ref_link,
        )

    def get_exports(self, context: dict[str, Any]) -> dict[str, Any]:
        """The v2 renderer exports the `ref` family of functions s into the rendering context"""
        return {
            "ref": self.get_ref(context),
            "ref_link": self.get_ref_link(context),
            "ref_no_link": self.get_ref_no_link(context),
        }
