import json
from pathlib import Path
from typing import Any, Dict, List, Optional, TextIO, Tuple, Type, Union

import requests
from jinja2 import (
    BaseLoader,
    ChoiceLoader,
    DictLoader,
    Environment,
    FileSystemLoader,
    PackageLoader,
    StrictUndefined,
    Template,
    TemplateNotFound,
    select_autoescape,
)
from pydantic import BaseModel

from fused_batch._formatter.formatter_templates import (
    fused_template_repr,
    fused_templates_repr,
)
from fused_batch._options import options as OPTIONS
from fused_batch._request import raise_for_status
from fused_batch.models import (
    AnyJobStepConfig,
    JobConfig,
    JobStepConfig,
    RootAnyJobStepConfig,
)
from fused_batch.models.coerce_dataset import (
    CoerceableToDatasetInput,
    CoerceableToDatasetOutput,
    _object_to_dataset_input,
    _object_to_dataset_output,
)
from fused_batch.models.udf import AnyBaseUdf, RootAnyBaseUdf

from .loaders import RegisteredLoader, UrlLoader


class FusedTemplate:
    name: str
    """Name of this template."""
    source: str
    """The Jinja template source for the JSON representation of the object."""
    template: Template
    """The Jinja template."""
    defaults: Dict[str, Any]
    """Default keyword arguments for the template."""

    def __init__(
        self, name: str, source: str, template: Template, defaults: Dict[str, Any]
    ):
        self.name = name
        self.source = source
        self.template = template
        self.defaults = defaults

    def _repr_html_(self):
        return fused_template_repr(self)

    def print(self, **kwargs):
        """Print the code generated by the template.

        Keyword arguments are passed to the template."""
        raise NotImplementedError()

    def code(self, **kwargs):
        """Set the code generated by the template as the next IPython cell.

        Keyword arguments are passed to the template."""
        raise NotImplementedError()

    def render(self, **kwargs):
        """Render the template to an object.

        Keyword arguments are passed to the template."""
        raise NotImplementedError()

    async def render_async(self, **kwargs):
        """Render the template to an object.

        Keyword arguments are passed to the template."""
        raise NotImplementedError()

    def write(self, out: Union[Path, str, TextIO]):
        """Write the template source out.

        Args:
            out: A file name, or a file-like object.
        """
        if isinstance(out, (str, Path)):
            with open(out, "w") as f:
                f.write(self.source)
        else:
            out.write(self.source)


class JobStepConfigTemplate(FusedTemplate):
    """Template that generates a JobStepConfig.

    The extension for this type of template is `job_tmpl`.
    """

    def print(
        self,
        left: CoerceableToDatasetInput,
        right: CoerceableToDatasetInput,
        output: Optional[CoerceableToDatasetOutput] = None,
        **kwargs,
    ) -> None:
        print(
            self.render(
                left=left,
                right=right,
                output=output,
                **kwargs,
            ).print_code()
        )

    def code(
        self,
        left: CoerceableToDatasetInput,
        right: CoerceableToDatasetInput,
        output: Optional[CoerceableToDatasetOutput] = None,
        **kwargs,
    ) -> None:
        import IPython

        IPython.get_ipython().set_next_input(
            self.render(
                left=left,
                right=right,
                output=output,
                **kwargs,
            ).print_code(),
            replace=True,
        )

    def render(
        self,
        left: CoerceableToDatasetInput,
        right: CoerceableToDatasetInput,
        output: Optional[CoerceableToDatasetOutput] = None,
        **kwargs,
    ) -> AnyJobStepConfig:
        resolved_kwargs = {
            **self.defaults,
            "left": _object_to_dataset_input(left),
            "right": _object_to_dataset_input(right),
            "output": _object_to_dataset_output(output),
            **kwargs,
        }
        json = self.template.render(**resolved_kwargs)
        config = RootAnyJobStepConfig.model_validate_json(json).root
        return config

    async def render_async(
        self,
        left: CoerceableToDatasetInput,
        right: CoerceableToDatasetInput,
        output: Optional[CoerceableToDatasetOutput] = None,
        **kwargs,
    ) -> JobStepConfig:
        resolved_kwargs = {
            **self.defaults,
            "left": _object_to_dataset_input(left),
            "right": _object_to_dataset_input(right),
            "output": _object_to_dataset_output(output),
            **kwargs,
        }
        json = await self.template.render_async(**resolved_kwargs)
        config = RootAnyJobStepConfig.model_validate_json(json).root
        return config


class JobConfigTemplate(FusedTemplate):
    """Template that generates a JobConfig.

    The extension for this type of template is `job_tmpl`.
    """

    # TODO: print
    # TODO: code

    def render(
        self,
        input: Optional[CoerceableToDatasetInput] = None,
        left: Optional[CoerceableToDatasetInput] = None,
        right: Optional[CoerceableToDatasetInput] = None,
        output: Optional[CoerceableToDatasetOutput] = None,
        **kwargs,
    ) -> JobConfig:
        resolved_kwargs = {
            **self.defaults,
            "input": _object_to_dataset_input(input) if input else None,
            "left": _object_to_dataset_input(left),
            "right": _object_to_dataset_input(right),
            "output": _object_to_dataset_output(output),
            **kwargs,
        }
        json = self.template.render(**resolved_kwargs)
        config = JobConfig.model_validate_json(json)
        return config

    async def render_async(
        self,
        input: Optional[CoerceableToDatasetInput] = None,
        left: Optional[CoerceableToDatasetInput] = None,
        right: Optional[CoerceableToDatasetInput] = None,
        output: Optional[CoerceableToDatasetOutput] = None,
        **kwargs,
    ) -> JobConfig:
        resolved_kwargs = {
            **self.defaults,
            "input": _object_to_dataset_input(input) if input else None,
            "left": _object_to_dataset_input(left),
            "right": _object_to_dataset_input(right),
            "output": _object_to_dataset_output(output),
            **kwargs,
        }
        json = await self.template.render_async(**resolved_kwargs)
        config = JobConfig.model_validate_json(json)
        return config


class UdfTemplate(FusedTemplate):
    """Template that generates a UDF.

    The extension for this type of template is `udf_tmpl`.
    """

    def print(self, **kwargs) -> None:
        print(self.render(**kwargs).code)

    def code(self, **kwargs) -> None:
        import IPython

        IPython.get_ipython().set_next_input(self.render(**kwargs).code, replace=True)

    def render(self, **kwargs) -> AnyBaseUdf:
        resolved_kwargs = {
            **self.defaults,
            **kwargs,
        }
        json = self.template.render(**resolved_kwargs)
        return RootAnyBaseUdf.model_validate_json(json).root

    async def render_async(self, **kwargs) -> AnyBaseUdf:
        resolved_kwargs = {
            **self.defaults,
            **kwargs,
        }
        json = await self.template.render_async(**resolved_kwargs)
        return RootAnyBaseUdf.model_validate_json(json).root


class Templates:
    """A collection of templates.

    Each template can be accessed with subscript operator or dot operator, e.g.
    `templates["template_name"]` or `templates.template_name`. Templates may be
    registered on this object using the subscript operator.

    Example:

        ```py
        fused.templates.udfs['my_udf_template'] = my_udf
        ```

    Registered templates always take priority over templates loaded from files,
    and over built-in templates.
    """

    _env: Environment
    _choice_loader: ChoiceLoader
    _name: str
    _template_class: Type[FusedTemplate]
    _defaults: Dict[str, Any]
    _extension: str
    _registered: Dict[str, str]

    def __init__(
        self,
        name: str,
        template_class: Type,
        extension: str,
        defaults: Optional[Dict[str, Any]] = None,
    ):
        self._name = name
        self._template_class = template_class
        self._defaults = defaults if defaults is not None else {}
        self._extension = extension
        self._registered = {}
        self._choice_loader = ChoiceLoader([])
        self._reset_loaders()
        self._env = _create_env(self._choice_loader)
        self._env.filters["embedpy"] = _embedpy_filter

    def _repr_html_(self):
        return fused_templates_repr(self._name, self._all_templates_by_loader())

    def __getitem__(self, key: str) -> FusedTemplate:
        template_key = self._name_to_template_key(key)
        source, _filename, _uptodate = self._env.loader.get_source(
            self._env, template_key
        )
        template = self._env.get_template(template_key)
        return self._template_class(
            template=template,
            name=key,
            source=source,
            defaults=self._defaults,
        )

    def __setitem__(self, key: str, value: Union[BaseModel, str]):
        if isinstance(value, BaseModel):
            value = value.json(indent=4)
        if not isinstance(value, str):
            raise TypeError(f"Template source has unexpected type {type(value)}")
        template_key = self._name_to_template_key(key)
        self._registered[template_key] = value

    def __delitem__(self, key: str):
        template_key = self._name_to_template_key(key)
        del self._registered[template_key]

    def __getattribute__(self, key):
        try:
            return super().__getattribute__(key)
        except AttributeError:
            try:
                return self[key]
            # Note that we need to raise an AttributeError, **not a TemplateNotFound** so that
            # IPython's _repr_html_ works here
            except TemplateNotFound:
                raise AttributeError(
                    f"object of type {type(self).__name__} has no attribute {key}"
                ) from None

    def list_templates(self) -> List[str]:
        """Returns a list of all template names in this collection."""
        return [
            name.rsplit(f".{self._extension}", maxsplit=1)[0]
            for name in self._env.list_templates(extensions=[self._extension])
        ]

    def load_files(self, where: Union[Path, str], **kwargs) -> "Templates":
        """Add a directory to the search path for templates.

        Loaded templates may override built-in templates, but will not take priority
        over templates registered on this object.

        Arguments:
            where: Path to a directory to load. The directory should have files with the extension for this registry, e.g. `job_tmpl` for `fused.templates.jobs`.
        """
        # Index 1 puts the loaded templates *after* the DictLoader for the registered templates
        # TODO: Maybe add an option for loading a Python source file directly as a header
        self._choice_loader.loaders.insert(1, FileSystemLoader(where, **kwargs))
        return self

    def load_url(
        self,
        url: str,
        name: str,
        **request_kwargs,
    ) -> "Templates":
        """Add a URL to the search path for templates.

        Loaded templates may override built-in templates, but will not take priority
        over templates registered on this object.

        Arguments:
            url: URL to load from. A single template should be returned by the URL.
            name: Name to save the loaded template as.
        """
        # TODO: Support loading directly from gist here
        r = requests.get(
            url=url,
            timeout=OPTIONS.request_timeout,
            **request_kwargs,
        )
        raise_for_status(r)

        key = self._name_to_template_key(name)
        source = r.text

        self._choice_loader.loaders.insert(1, UrlLoader({key: source}, url=url))
        return self

    def from_object(
        self, obj: BaseModel, name: str, register: bool = False
    ) -> FusedTemplate:
        """Create a template from an object.

        Args:
            obj: The object to create a template from. It should be of the appropriate type for this registry, e.g. a `AnyBaseUdf` for the UDFs registry.
            name: The name of the template to create.
            register: If True, store the resulting template in this registry. Defaults to False.
        """
        if register:
            self[name] = obj
            return self[name]

        # Otherwise, we need to specifically create the template
        source = obj.json(indent=4)
        template = _create_env(DictLoader({name: source})).get_template(name)
        return self._template_class(
            template=template,
            name=name,
            source=source,
            defaults=self._defaults,
        )

    def _all_templates(self) -> Dict[str, FusedTemplate]:
        return {name: self[name] for name in self.list_templates()}

    def _all_templates_by_loader(self) -> Dict[str, Tuple[FusedTemplate, BaseLoader]]:
        # This supports finding all the templates and returning them with their loader
        # for use in the repr
        results: Dict[str, Tuple[FusedTemplate, BaseLoader]] = {}
        for loader in self._choice_loader.loaders:
            all_possible_names = loader.list_templates()
            # We have to apply the extension filter our selves here
            names_with_extensions = [
                name.rsplit(".", maxsplit=1)
                for name in all_possible_names
                if "." in name
            ]
            names = [
                name_and_extension[0]
                for name_and_extension in names_with_extensions
                if name_and_extension[1] == self._extension
            ]
            for name in names:
                if name not in results:
                    # Note that we don't necessarily know that the template is loaded from
                    # this loader. Instead we are replicating the "choose first" behavior from
                    # Jinja.
                    results[name] = (self[name], loader)
        return results

    def _name_to_template_key(self, key: str) -> str:
        return f"{key}.{self._extension}"

    def _reset_loaders(self):
        self._choice_loader.loaders = [
            RegisteredLoader(self._registered),
            PackageLoader("fused_batch._static"),
        ]

    def __dir__(self) -> List[str]:
        # Provide method name lookup and completion. Only provide 'public' methods.
        # This enables autocompletion
        extra_attrs = set(self.list_templates())
        normal_dir = {name for name in dir(type(self)) if not name.startswith("_")}

        return sorted(normal_dir | extra_attrs)


def _create_env(loader: BaseLoader) -> Environment:
    return Environment(
        loader=loader,
        autoescape=select_autoescape(),
        undefined=StrictUndefined,
    )


def _embedpy_filter(code: str) -> str:
    assert code is not None, "embedpy filter called on None"
    assert isinstance(
        code, str
    ), f"embedpy filter called on a non-string (type was actually {type(code).__name__})"
    return json.dumps(code)[1:-1]


jobs = Templates(
    name="jobs", template_class=JobStepConfigTemplate, extension="job_tmpl"
)
"""Templates for jobs (joins, maps, etc)"""
job = Templates(name="job", template_class=JobConfigTemplate, extension="job_tmpl")
"""Templates for jobs with one or more steps"""
udfs = Templates(name="udfs", template_class=UdfTemplate, extension="udf_tmpl")
"""Templates for UDFs"""
