from __future__ import annotations
from os import PathLike

import jinja2
from typing import Any, Callable, Iterator, Mapping, Optional, Sequence, Union, cast
from jinja2.nativetypes import NativeEnvironment
from the1conf.attr_dict import AttrDict
from .utils import is_sequence, _undef
from pathlib import Path
from dataclasses import dataclass, field


class SegmentedPath(PathLike[str]):
    """
    A class that represents a path as a list of segments.

    Allow to perform variable substitution on each segment separately and then reassemble the path with the same structure as the input.

    Behaves like a path-like object and therefore can be used in a pathlib.Path context or with open()
    Stringified as a path string.
    """

    _raw_segments: list[str | Path]
    _resolved_path: Optional[Path]
    _resolved_stringification: Optional[str]

    @classmethod
    def create(
        cls, value: str | Path | SegmentedPath, *extra_segments: Any
    ) -> SegmentedPath:
        """Factory method to create a SegmentedPath from a string, a Path or another SegmentedPath and add optional extra segments."""
        if isinstance(value, SegmentedPath):
            base = value
        elif isinstance(value, (str, Path)):
            base = SegmentedPath(value)
        else:
            raise TypeError(f"P expects str | Path | SegmentedPath, got {type(value)}")

        for segment in extra_segments:
            if not isinstance(segment, (str, Path, SegmentedPath)):
                raise TypeError(f"Invalid segment type for P: {type(segment)}")
            base = base / segment

        return base

    def __init__(self, *args: str | Path | SegmentedPath) -> None:
        """Create a SegmentedPath from a list of segments, which can be strings, Paths or other SegmentedPath."""
        self._raw_segments = []
        for arg in args:
            self._add(arg)

        self._resolved_path = None
        self._resolved_stringification = None

    def _add(self, segment: str | Path | SegmentedPath) -> None:
        """add a segment to the SegmentedSegment
        We don't want to use the __add__ operator because the semantic of Path is to use '/' for path concatenation
        """
        if isinstance(segment, SegmentedPath):
            self._raw_segments.extend(segment._raw_segments)
        elif isinstance(segment, (str, Path)):
            self._raw_segments.append(segment)
        else:
            raise TypeError(f"Invalid argument type: {type(segment)}")

    def __truediv__(self, other: str | Path | SegmentedPath) -> SegmentedPath:
        """allow using operator / to concatenate SegmentedPath with strings, Paths or other SegmentedPath."""
        if isinstance(other, (str, Path, SegmentedPath)):
            return SegmentedPath(self, other)
        else:
            return NotImplemented

    def __rtruediv__(self, other: str | Path | SegmentedPath) -> SegmentedPath:
        """allow using operator / to concatenate string, Paths or other SegmentedPath with SegmentedPath ."""
        if isinstance(other, (str, Path, SegmentedPath)):
            return SegmentedPath(other, self)
        else:
            return NotImplemented

    def __str__(self) -> str:
        return self.__repr__()

    def __fspath__(self) -> str:
        """allow using SegmentedPath in a pathlib.Path context or with open() by returning the string representation
        of the path if the object has been resolved, otherwise raise en exception.
        Note the type of exception is important to allow fallback to __rtruediv__ in Path operations if the object is not resolved yet.
        """
        if self._resolved_stringification is None:
            raise TypeError(
                "SegmentedPath has not been resolved yet. Call resolve(context) before using it as a path like object."
            )
        return self._resolved_stringification

    def __repr__(self) -> str:
        """stringify the SegmentedPath as a path string. If the path has not been resolved yet, it will use the raw segments, otherwise it will use the solved segments."""
        if self._resolved_stringification is None:
            res = str(Path(*self._raw_segments))
        else:
            res = self._resolved_stringification
        return res

    def resolve(self, context: RenderContext) -> SegmentedPath:
        """Render each segment with Jinja using the provided context.

        returns self to allow chaining of resolve and other operations.
        """
        solved_segments = [render_template(s, context) for s in self._raw_segments]
        self._resolved_path = Path(*solved_segments)
        self._resolved_stringification = str(self._resolved_path)
        return self


@dataclass
class SubstitutionResult:
    """
    A class to hold the result of a substitution operation.

    includes the rendered values of all candidate templates and the final rendered value.
    """

    rendered_value: Any = field(default=_undef)
    rendered_candidates: list[Any] = field(default_factory=list)


# Native Jinja env: returns native Python objects when possible
_JINJA_ENV: NativeEnvironment = NativeEnvironment(
    autoescape=False,
    undefined=jinja2.StrictUndefined,
)

# Filter style: {{ "/tmp" | P / "base" }}
_JINJA_ENV.filters["p"] = SegmentedPath.create

# Global function style: {{ P("/tmp") / "base" }}
_JINJA_ENV.globals["P"] = SegmentedPath.create


def render_jinja_native(template_value: str, context: RenderContext) -> Any:
    """Render a template with NativeEnvironment and return native objects."""
    template = _JINJA_ENV.from_string(template_value)
    return template.render(context)


Sp = SegmentedPath

RenderContext = Mapping[str, Any]


def render_template(
    value: Any,
    context: RenderContext,
) -> Any:
    """
    Render a value rendered as a jinja2 template using the provided context.

    The value can be a jinja template or not. If it is not a jinja template, it is returned as is.
    If it is a jinja template, it is rendered with the provided context and the result is returned.

    The value can be of any type but only string and SegmentedPath and Path are processed directly, dict and sequence
    are processed recursively, while other types are returned as they are.

    - If the value is not a string, nor a Path, nor a SegmentedPath nor a dict or a sequence , it is returned as is.
    - If the value is a string it is rendered and returned as a string
    - If the value is a SegmentedPath it is resolved and returned as a SegmentedPath.
    - If the value is a Path it is stringified and rendered and returned as a Path.
    - If the value is a sequence (except strings), each item is rendered and the result is returned as a sequence .
    - If the value is a dict, each value is rendered and the result is returned as a dict with the same keys and the rendered values.
    """

    if value is None:
        return None
    if isinstance(value, SegmentedPath):
        return Path(value.resolve(context))
    if is_sequence(value):
        return type(value)([render_template(item, context) for item in value])
    if isinstance(value, dict):
        return type(value)({k: render_template(v, context) for k, v in value.items()})
    if isinstance(value, Path):
        val_to_render = str(
            value
        )  # convert to string for rendering, we will convert back to Path after rendering
    elif isinstance(value, str):
        val_to_render = value
    else:
        return value

    try:
        rendered = render_jinja_native(val_to_render, context)

        if isinstance(rendered, SegmentedPath):
            return str(rendered.resolve(context))
        if isinstance(value, Path):
            return str(rendered)
        return rendered

    except jinja2.exceptions.UndefinedError as e:
        raise ValueError(
            f"Undefined variable in template: {val_to_render} with context:\n {str(AttrDict(context))}"
        ) from e
    except jinja2.TemplateError as e:
        raise ValueError(
            f"Error rendering template: {val_to_render} with context:\n {str(AttrDict(context))}"
        ) from e


def resolve_candidate_list(
    *,
    candidates: str | Path | SegmentedPath | None | Sequence[str | Path | SegmentedPath],
    context: RenderContext,
    check_exists: Callable[[str], Any],
) -> SubstitutionResult:
    """
    Takes a list of candidates values (file paths,env name, value key or file_key) and resolves them in
    order, returning the first one for which check_exists returns a non-_undef value.

    Args:
        candidates: A single string, a Path or a SegmentedPath, or a list of these. Values can be
        potential jinja templates for string and SegmentedPath segments . Not for Path types.
        context: The context used for rendering templates.
        check_exists: An predicate to validate the rendered string.
                      returns _undef if not valid, or the value to return if valid.
    Returns:
        A SubstitutionResult containing the rendered candidates and the first rendered value that passed the
        check_exists validation. If no candidate is valid, rendered_value will be _undef.
    """
    res = SubstitutionResult()
    if candidates is None:
        return res

    candidate_list: list[str | Path | SegmentedPath] = []

    if is_sequence(candidates):
        candidate_list = cast(list[str | Path | SegmentedPath], candidates)
    else:
        candidate_list = [cast(str | Path | SegmentedPath, candidates)]

    for raw_candidate in candidate_list:
        rendered = render_template(raw_candidate, context)
        res.rendered_candidates.append(rendered)
        valid = check_exists(rendered)
        if valid != _undef:
            res.rendered_value = valid
            return res

    return res
