"""Expression template module for parameterized, persistable filter expressions.

This module provides:

- :class:`TemplateParam` placeholder substitution via :func:`bind` and :func:`collect_params`
- :class:`ExprTemplate` — a fully serializable wrapper combining an expression with named
  parameter rules evaluated from a runtime context
- :class:`RuleSerializer` — parser/serializer for the transform pipeline DSL
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

from therismos.expr._expr import (
    AllExpr,
    AnyExpr,
    Between,
    Expr,
    In,
    NotExpr,
    SingleValuedFieldExpr,
    TemplateParam,
)

if TYPE_CHECKING:
    from therismos.expr.serializer import Serializer
    from therismos.expr.transforms import TransformRegistry

_MISSING: Any = object()


class MissingParamError(KeyError):
    """Raised when a required template parameter has no value during :func:`bind`.

    :param name: The name of the missing parameter.
    :type name: str
    """

    def __init__(self, name: str) -> None:
        super().__init__(name)
        self.name = name

    def __str__(self) -> str:
        return f"Missing template parameter: ${self.name}"


def collect_params(expr: Expr) -> dict[str, TemplateParam]:
    """Return all :class:`TemplateParam` nodes in the expression tree, keyed by name.

    :param expr: The expression to inspect.
    :type expr: Expr
    :return: Mapping of parameter name to :class:`TemplateParam` instance.
    :rtype: dict[str, TemplateParam]
    """
    result: dict[str, TemplateParam] = {}

    if isinstance(expr, SingleValuedFieldExpr):
        if isinstance(expr.value, TemplateParam):
            result[expr.value.name] = expr.value
    elif isinstance(expr, Between):
        for v in (expr.lower, expr.upper):
            if isinstance(v, TemplateParam):
                result[v.name] = v
    elif isinstance(expr, In):
        for v in expr.values:
            if isinstance(v, TemplateParam):
                result[v.name] = v
    elif isinstance(expr, (AllExpr, AnyExpr)):
        for child in expr.exprs:
            result.update(collect_params(child))
    elif isinstance(expr, NotExpr):
        result.update(collect_params(expr.expr))

    return result


def _resolve(param: TemplateParam, raw: Any) -> Any:
    """Apply the param's type cast to a raw value."""
    if param.type_ is None:
        return raw
    if isinstance(param.type_, type) and isinstance(raw, param.type_):
        return raw
    return param.type_(raw)


def bind(expr: Expr, params: dict[str, Any]) -> Expr:
    """Substitute all :class:`TemplateParam` nodes with concrete values.

    Mirrors the recursive tree-walk pattern of the optimizer.

    :param expr: The expression to bind.
    :type expr: Expr
    :param params: Mapping of parameter name to concrete value.
    :type params: dict[str, Any]
    :return: A new expression with all :class:`TemplateParam` replaced.
    :rtype: Expr
    :raises MissingParamError: If any :class:`TemplateParam` has no entry in ``params``.
    """
    if isinstance(expr, SingleValuedFieldExpr):
        if isinstance(expr.value, TemplateParam):
            param = expr.value
            if param.name not in params:
                raise MissingParamError(param.name)
            concrete = _resolve(param, params[param.name])
            return type(expr)(expr.field, concrete)
        return expr

    if isinstance(expr, Between):
        lower = expr.lower
        upper = expr.upper
        if isinstance(lower, TemplateParam):
            if lower.name not in params:
                raise MissingParamError(lower.name)
            lower = _resolve(lower, params[lower.name])
        if isinstance(upper, TemplateParam):
            if upper.name not in params:
                raise MissingParamError(upper.name)
            upper = _resolve(upper, params[upper.name])
        if lower is expr.lower and upper is expr.upper:
            return expr
        return Between(expr.field, lower, upper)

    if isinstance(expr, In):
        new_values = []
        changed = False
        for v in expr.values:
            if isinstance(v, TemplateParam):
                if v.name not in params:
                    raise MissingParamError(v.name)
                new_values.append(_resolve(v, params[v.name]))
                changed = True
            else:
                new_values.append(v)
        if not changed:
            return expr
        return In(expr.field, tuple(new_values))

    if isinstance(expr, AllExpr):
        children = tuple(bind(child, params) for child in expr.exprs)
        if children == expr.exprs:
            return expr
        return AllExpr(*children)

    if isinstance(expr, AnyExpr):
        children = tuple(bind(child, params) for child in expr.exprs)
        if children == expr.exprs:
            return expr
        return AnyExpr(*children)

    if isinstance(expr, NotExpr):
        child = bind(expr.expr, params)
        if child is expr.expr:
            return expr
        return NotExpr(child)

    return expr


@dataclass(frozen=True)
class TransformStep:
    """One step in a transform pipeline.

    :ivar name: Transform function name (e.g., ``"extract_date"``).
    :vartype name: str
    :ivar args: Parsed literal arguments (``str``, ``int``, ``float``, ``timedelta``).
    :vartype args: tuple[Any, ...]
    """

    name: str
    args: tuple[Any, ...] = field(default_factory=tuple)


@dataclass(frozen=True)
class ParamRule:
    """A computation pipeline that derives a param value from a context variable.

    :ivar source: Context variable name (without leading ``$``).
    :vartype source: str
    :ivar steps: Ordered pipeline steps applied to the source value.
    :vartype steps: tuple[TransformStep, ...]
    """

    source: str
    steps: tuple[TransformStep, ...] = field(default_factory=tuple)


@dataclass(frozen=True)
class TemplateParamSpec:
    """Optional metadata for a named template parameter.

    :ivar default: Value to use when the parameter is not in the context and has no rule.
                   Use the module-level ``_MISSING`` sentinel (default) to indicate no default.
    :ivar description: Human-readable description.
    :vartype description: str
    """

    default: Any = _MISSING
    description: str = ""


class RuleSerializer:
    """Parses and serializes :class:`ParamRule` objects to/from pipeline DSL strings.

    **Syntax:**  ``$source | step1 | step2(arg1, arg2)``

    Supported argument literals: quoted strings, integers, floats, duration literals
    (``7d``, ``1h``, ``30m``, ``90s``, ``500ms``).

    :param registry: Optional transform registry used to validate step names.
    :type registry: TransformRegistry | None
    """

    def __init__(self, registry: TransformRegistry | None = None) -> None:
        self._registry = registry

    def serialize(self, rule: ParamRule) -> str:
        """Serialize a :class:`ParamRule` to a pipeline DSL string.

        :param rule: The rule to serialize.
        :type rule: ParamRule
        :return: Pipeline DSL string, e.g. ``"$now | extract_date | sub_time(7d)"``.
        :rtype: str
        """
        parts = [f"${rule.source}"]
        for step in rule.steps:
            if step.args:
                args_str = ", ".join(self._serialize_arg(a) for a in step.args)
                parts.append(f"{step.name}({args_str})")
            else:
                parts.append(step.name)
        return " | ".join(parts)

    def deserialize(self, s: str) -> ParamRule:
        """Parse a pipeline DSL string into a :class:`ParamRule`.

        :param s: The pipeline string, e.g. ``"$now | extract_date | sub_time(7d)"``.
        :type s: str
        :return: The parsed :class:`ParamRule`.
        :rtype: ParamRule
        :raises ValueError: If the string is syntactically invalid.
        """
        tokens = [t.strip() for t in s.split("|")]
        if not tokens:
            raise ValueError(f"Empty rule: {s!r}")

        source_token = tokens[0]
        if not source_token.startswith("$"):
            raise ValueError(f"Rule source must start with '$': {source_token!r}")
        source = source_token[1:].strip()
        if not source:
            raise ValueError(f"Rule source name is empty: {s!r}")

        steps = tuple(self._parse_step(t) for t in tokens[1:])
        return ParamRule(source=source, steps=steps)

    def _parse_step(self, token: str) -> TransformStep:
        """Parse a single pipeline step token."""
        token = token.strip()
        paren_open = token.find("(")
        if paren_open == -1:
            return TransformStep(name=token, args=())

        name = token[:paren_open].strip()
        if not token.endswith(")"):
            raise ValueError(f"Malformed step (missing closing ')'): {token!r}")
        args_str = token[paren_open + 1 : -1]
        args = tuple(self._parse_arg(a.strip()) for a in args_str.split(",") if a.strip())
        return TransformStep(name=name, args=args)

    def _parse_arg(self, token: str) -> Any:
        """Parse a single argument token to its Python value."""
        if _DURATION_PARSE_RE.match(token):
            return _parse_duration(token)

        if token.startswith('"') and token.endswith('"'):
            return token[1:-1].replace('\\"', '"').replace("\\\\", "\\")

        if re.match(r"^-?\d+\.\d+$", token):
            return float(token)
        if re.match(r"^-?\d+$", token):
            return int(token)

        raise ValueError(f"Cannot parse argument: {token!r}")

    def _serialize_arg(self, value: Any) -> str:
        """Serialize a single argument value."""
        from datetime import timedelta

        if isinstance(value, timedelta):
            return _format_duration(value)
        if isinstance(value, str):
            escaped = value.replace("\\", "\\\\").replace('"', '\\"')
            return f'"{escaped}"'
        return str(value)


_DURATION_PARSE_RE = re.compile(r"^(\d+)(ms|d|h|m|s)$")


def _parse_duration(s: str) -> Any:
    """Parse a duration string (``7d``, ``1h``, ``30m``, ``90s``, ``500ms``) to :class:`timedelta`.

    :param s: Duration string.
    :type s: str
    :return: Corresponding :class:`datetime.timedelta`.
    :raises ValueError: If the format is not recognized.
    """
    from datetime import timedelta

    m = _DURATION_PARSE_RE.match(s)
    if not m:
        raise ValueError(f"Cannot parse duration: {s!r}. Expected format: <number>(ms|d|h|m|s)")
    n, unit = int(m.group(1)), m.group(2)
    if unit == "d":
        return timedelta(days=n)
    if unit == "h":
        return timedelta(hours=n)
    if unit == "m":
        return timedelta(minutes=n)
    if unit == "s":
        return timedelta(seconds=n)
    return timedelta(milliseconds=n)


def _format_duration(td: Any) -> str:
    """Serialize a :class:`timedelta` to the shortest unambiguous duration string.

    :param td: The timedelta to format.
    :type td: datetime.timedelta
    :return: Duration string like ``"7d"``, ``"1h"``, ``"30m"``, ``"90s"``, ``"500ms"``.
    :rtype: str
    """
    from datetime import timedelta

    total_ms = int(td / timedelta(milliseconds=1))
    units = [
        ("d", 86_400_000),
        ("h", 3_600_000),
        ("m", 60_000),
        ("s", 1_000),
        ("ms", 1),
    ]
    for suffix, ms_per_unit in units:
        if total_ms % ms_per_unit == 0:
            return f"{total_ms // ms_per_unit}{suffix}"
    return f"{total_ms}ms"


@dataclass(frozen=True)
class ExprTemplate:
    """A parameterized, persistable expression template.

    Wraps an :class:`Expr` that may contain :class:`TemplateParam` placeholders.
    Rules define how parameter values are computed from a runtime context.

    :ivar expr: Expression containing :class:`TemplateParam` placeholders.
    :vartype expr: Expr
    :ivar params: Per-parameter metadata (type, default, description). All keys optional.
    :vartype params: dict[str, TemplateParamSpec]
    :ivar rules: Pipeline rules mapping parameter names to their computation.
    :vartype rules: dict[str, ParamRule]
    """

    expr: Expr
    params: dict[str, TemplateParamSpec] = field(default_factory=dict)
    rules: dict[str, ParamRule] = field(default_factory=dict)

    def bind(
        self,
        context: dict[str, Any] | None = None,
        registry: TransformRegistry | None = None,
        **overrides: Any,
    ) -> Expr:
        """Evaluate rules against *context*, apply *overrides*, fall back to defaults, bind.

        Resolution priority (highest first):
        1. ``overrides``
        2. Rule-evaluated values (from ``context``)
        3. :attr:`TemplateParamSpec.default`

        :param context: Runtime context (e.g., ``{"now": datetime(...)}``) used by rules.
        :type context: dict[str, Any] | None
        :param registry: Transform registry for evaluating pipeline steps.
                         Defaults to :data:`~therismos.expr.transforms.DEFAULT_TRANSFORM_REGISTRY`.
        :type registry: TransformRegistry | None
        :param overrides: Explicit parameter values that take priority over rules.
        :return: Fully bound :class:`Expr` with no remaining :class:`TemplateParam` nodes.
        :rtype: Expr
        :raises MissingParamError: If any parameter cannot be resolved.
        :raises KeyError: If a rule references a context key that is absent.
        """
        if registry is None:
            from therismos.expr.transforms import DEFAULT_TRANSFORM_REGISTRY

            registry = DEFAULT_TRANSFORM_REGISTRY

        ctx = context or {}
        resolved: dict[str, Any] = {}

        for name, rule in self.rules.items():
            try:
                value = ctx[rule.source]
            except KeyError:
                raise KeyError(
                    f"Rule for parameter {name!r} references missing context key {rule.source!r}"
                ) from None
            for step in rule.steps:
                value = registry.apply(step.name, value, *step.args)
            resolved[name] = value

        resolved.update(overrides)

        for name, spec in self.params.items():
            if name not in resolved and spec.default is not _MISSING:
                resolved[name] = spec.default

        return bind(self.expr, resolved)

    def collect_missing(self, context: dict[str, Any] | None = None) -> set[str]:
        """Return parameter names that cannot be resolved from rules + context.

        Does not account for ``overrides``; those are supplied at :meth:`bind` time.

        :param context: Runtime context to check rules against.
        :type context: dict[str, Any] | None
        :return: Set of unresolvable parameter names.
        :rtype: set[str]
        """
        ctx = context or {}
        all_params = collect_params(self.expr)
        missing = set()
        for name in all_params:
            if name in self.rules and self.rules[name].source in ctx:
                continue
            spec = self.params.get(name)
            if spec is not None and spec.default is not _MISSING:
                continue
            missing.add(name)
        return missing

    def to_dict(self, serializer: Serializer | None = None) -> dict[str, Any]:
        """Serialize this template to a JSON-compatible dict.

        :param serializer: Serializer for the expression string.
                           Defaults to a plain :class:`~therismos.expr.serializer.Serializer`.
        :type serializer: Serializer | None
        :return: JSON-compatible dict with keys ``version``, ``expr``, ``params``, ``rules``.
        :rtype: dict[str, Any]
        """
        if serializer is None:
            from therismos.expr.serializer import Serializer

            serializer = Serializer()

        rule_ser = RuleSerializer()

        params_dict: dict[str, Any] = {}
        for name, spec in self.params.items():
            entry: dict[str, Any] = {}
            if spec.description:
                entry["description"] = spec.description
            if spec.default is not _MISSING:
                entry["default"] = spec.default
            params_dict[name] = entry

        rules_dict = {name: rule_ser.serialize(rule) for name, rule in self.rules.items()}

        return {
            "version": "1",
            "expr": serializer.serialize(self.expr),
            "params": params_dict,
            "rules": rules_dict,
        }

    @classmethod
    def from_dict(
        cls,
        data: dict[str, Any],
        serializer: Serializer | None = None,
    ) -> ExprTemplate:
        """Deserialize a template from a JSON-compatible dict.

        :param data: Dict produced by :meth:`to_dict`.
        :type data: dict[str, Any]
        :param serializer: Serializer for parsing the expression string.
                           Defaults to a plain :class:`~therismos.expr.serializer.Serializer`.
        :type serializer: Serializer | None
        :return: Reconstructed :class:`ExprTemplate`.
        :rtype: ExprTemplate
        :raises ValueError: If the version field is not ``"1"``.
        """
        version = data.get("version", "1")
        if version != "1":
            raise ValueError(f"Unsupported ExprTemplate version: {version!r}")

        if serializer is None:
            from therismos.expr.serializer import Serializer

            serializer = Serializer()

        expr = serializer.deserialize(data["expr"])

        rule_ser = RuleSerializer()
        rules = {name: rule_ser.deserialize(s) for name, s in data.get("rules", {}).items()}

        params: dict[str, TemplateParamSpec] = {}
        for name, entry in data.get("params", {}).items():
            default = entry.get("default", _MISSING)
            description = entry.get("description", "")
            params[name] = TemplateParamSpec(default=default, description=description)

        return cls(expr=expr, params=params, rules=rules)

    def to_json(self, serializer: Serializer | None = None) -> str:
        """Serialize to a JSON string.

        :param serializer: Optional serializer for the expression.
        :type serializer: Serializer | None
        :return: JSON string.
        :rtype: str
        """
        return json.dumps(self.to_dict(serializer), indent=2)

    @classmethod
    def from_json(
        cls,
        s: str,
        serializer: Serializer | None = None,
    ) -> ExprTemplate:
        """Deserialize from a JSON string.

        :param s: JSON string produced by :meth:`to_json`.
        :type s: str
        :param serializer: Optional serializer for parsing the expression.
        :type serializer: Serializer | None
        :return: Reconstructed :class:`ExprTemplate`.
        :rtype: ExprTemplate
        """
        return cls.from_dict(json.loads(s), serializer=serializer)
