"""Transform Lark parse trees into typed LarkElement rule trees."""

# pylint: disable=missing-function-docstring,unused-argument
from lark import Token, Tree, v_args, Transformer, Discard
from lark.tree import Meta

from hcl2.rules.base import (
    StartRule,
    BodyRule,
    BlockRule,
    AttributeRule,
)
from hcl2.rules.containers import (
    ObjectRule,
    ObjectElemRule,
    ObjectElemKeyRule,
    TupleRule,
    ObjectElemKeyExpressionRule,
)
from hcl2.rules.expressions import (
    BinaryTermRule,
    UnaryOpRule,
    BinaryOpRule,
    ExprTermRule,
    ConditionalRule,
)
from hcl2.rules.for_expressions import (
    ForTupleExprRule,
    ForObjectExprRule,
    ForIntroRule,
    ForCondRule,
)
from hcl2.rules.functions import ArgumentsRule, FunctionCallRule
from hcl2.rules.indexing import (
    IndexExprTermRule,
    SqbIndexRule,
    ShortIndexRule,
    GetAttrRule,
    GetAttrExprTermRule,
    AttrSplatExprTermRule,
    AttrSplatRule,
    FullSplatRule,
    FullSplatExprTermRule,
)
from hcl2.rules.literal_rules import (
    FloatLitRule,
    IntLitRule,
    IdentifierRule,
    BinaryOperatorRule,
    KeywordRule,
    LiteralValueRule,
)
from hcl2.rules.strings import (
    InterpolationRule,
    StringRule,
    StringPartRule,
    HeredocTemplateRule,
    HeredocTrimTemplateRule,
    TemplateStringRule,
)
from hcl2.rules.directives import (
    TemplateIfRule,
    TemplateForRule,
    TemplateIfStartRule,
    TemplateElseRule,
    TemplateEndifRule,
    TemplateForStartRule,
    TemplateEndforRule,
)
from hcl2.rules.tokens import (
    NAME,
    IntLiteral,
    FloatLiteral,
    StringToken,
    StaticStringToken,
)
from hcl2.rules.whitespace import NewLineOrCommentRule


class RuleTransformer(Transformer):
    """Takes a syntax tree generated by the parser and
    transforms it to a tree of LarkRule instances
    """

    with_meta: bool

    def transform(self, tree: Tree) -> StartRule:
        return super().transform(tree)

    def __init__(self, discard_new_line_or_comments: bool = False):
        super().__init__()
        self.discard_new_line_or_comments = discard_new_line_or_comments

    def __default_token__(self, token: Token) -> StringToken:
        # TODO make this return StaticStringToken where applicable
        value = token.value
        # The EQ terminal /[ \t]*=(?!=|>)/ captures leading whitespace.
        # Preserve it so the direct pipeline (to_lark → reconstruct) retains
        # original alignment.  The reconstructor skips its own space insertion
        # when the EQ token already carries leading whitespace.

        # Don't convert STRING_CHARS or ESCAPED_* tokens to static tokens.
        # E.g., STRING_CHARS("=") must stay STRING_CHARS, not become EQ.
        if token.type in ("STRING_CHARS", "ESCAPED_INTERPOLATION", "ESCAPED_DIRECTIVE"):
            return StringToken[token.type](value)  # type: ignore[misc]

        if value in StaticStringToken.classes_by_value:
            return StaticStringToken.classes_by_value[value]()
        return StringToken[token.type](value)  # type: ignore[misc]

    # pylint: disable=C0103
    def FLOAT_LITERAL(self, token: Token) -> FloatLiteral:
        return FloatLiteral(token.value)

    # pylint: disable=C0103
    def NAME(self, token: Token) -> NAME:
        return NAME(token.value)

    # pylint: disable=C0103
    def INT_LITERAL(self, token: Token) -> IntLiteral:
        return IntLiteral(token.value)

    @v_args(meta=True)
    def start(self, meta: Meta, args) -> StartRule:
        return StartRule(args, meta)

    @v_args(meta=True)
    def body(self, meta: Meta, args) -> BodyRule:
        return BodyRule(args, meta)

    @v_args(meta=True)
    def block(self, meta: Meta, args) -> BlockRule:
        return BlockRule(args, meta)

    @v_args(meta=True)
    def attribute(self, meta: Meta, args) -> AttributeRule:
        # _attribute_name is flattened, so args[0] may be KeywordRule,
        # LiteralValueRule, or IdentifierRule
        if isinstance(args[0], (KeywordRule, LiteralValueRule)):
            args[0] = IdentifierRule([NAME(args[0].token.value)], meta)
        return AttributeRule(args, meta)

    @v_args(meta=True)
    def new_line_or_comment(
        self, meta: Meta, args
    ):  # -> NewLineOrCommentRule | Discard
        if self.discard_new_line_or_comments:
            return Discard
        return NewLineOrCommentRule(args, meta)

    @v_args(meta=True)
    def identifier(self, meta: Meta, args) -> IdentifierRule:
        return IdentifierRule(args, meta)

    @v_args(meta=True)
    def keyword(self, meta: Meta, args) -> KeywordRule:
        return KeywordRule(args, meta)

    @v_args(meta=True)
    def literal_value(self, meta: Meta, args) -> LiteralValueRule:
        return LiteralValueRule(args, meta)

    @v_args(meta=True)
    def int_lit(self, meta: Meta, args) -> IntLitRule:
        return IntLitRule(args, meta)

    @v_args(meta=True)
    def float_lit(self, meta: Meta, args) -> FloatLitRule:
        return FloatLitRule(args, meta)

    @v_args(meta=True)
    def string(self, meta: Meta, args) -> StringRule:
        # Assemble flat directive parts into nested TemplateIfRule/TemplateForRule
        args = self._assemble_directives(list(args), meta)
        return StringRule(args, meta)

    def _assemble_directives(self, parts, meta: Meta):
        """Assemble flat directive string_parts into nested template rules.

        Scans for TemplateIfStartRule/TemplateForStartRule within StringPartRules
        and collects children up to matching endif/endfor, handling nesting.
        """
        result = []
        i = 0
        while i < len(parts):
            assembled, i = self._try_assemble_nested(parts, i, meta)
            if assembled is not None:
                result.append(StringPartRule([assembled], meta))
            else:
                result.append(parts[i])
                i += 1
        return result

    def _try_assemble_nested(self, parts, idx, meta):
        """If parts[idx] starts a directive, assemble and return (rule, next_idx).

        Returns (None, idx) if parts[idx] is not a directive opener.
        """
        part = parts[idx]
        if isinstance(part, StringPartRule):
            content = part.content
            if isinstance(content, TemplateIfStartRule):
                return self._assemble_template_if(parts, idx, meta)
            if isinstance(content, TemplateForStartRule):
                return self._assemble_template_for(parts, idx, meta)
        return None, idx

    def _collect_body(self, parts, start, end_types, meta):
        """Collect body parts from start until a StringPartRule with end_types content.

        Recursively assembles nested directives. Returns (body_list, end_content, next_idx).
        """
        body: list = []
        i = start
        while i < len(parts):
            part = parts[i]
            if isinstance(part, StringPartRule) and isinstance(part.content, end_types):
                return body, part.content, i + 1
            assembled, i = self._try_assemble_nested(parts, i, meta)
            if assembled is not None:
                body.append(StringPartRule([assembled], meta))
            else:
                body.append(parts[i])
                i += 1
        return body, None, i

    def _assemble_template_if(self, parts, start_idx, meta: Meta):
        """Assemble a TemplateIfRule from flat parts starting at start_idx."""
        if_start = parts[start_idx].content
        # Collect if-body until else or endif
        if_body, end, i = self._collect_body(
            parts, start_idx + 1, (TemplateElseRule, TemplateEndifRule), meta
        )
        else_rule = None
        else_body = None
        if isinstance(end, TemplateElseRule):
            else_rule = end
            else_body, end, i = self._collect_body(parts, i, (TemplateEndifRule,), meta)
        if not isinstance(end, TemplateEndifRule):
            raise RuntimeError("Unterminated template if directive")
        return TemplateIfRule(if_start, if_body, else_rule, else_body, end, meta), i

    def _assemble_template_for(self, parts, start_idx, meta: Meta):
        """Assemble a TemplateForRule from flat parts starting at start_idx."""
        for_start = parts[start_idx].content
        body, end, i = self._collect_body(
            parts, start_idx + 1, (TemplateEndforRule,), meta
        )
        if not isinstance(end, TemplateEndforRule):
            raise RuntimeError("Unterminated template for directive")
        return TemplateForRule(for_start, body, end, meta), i

    @v_args(meta=True)
    def string_part(self, meta: Meta, args) -> StringPartRule:
        return StringPartRule(args, meta)

    @v_args(meta=True)
    def interpolation(self, meta: Meta, args) -> InterpolationRule:
        return InterpolationRule(args, meta)

    @v_args(meta=True)
    def heredoc_template(self, meta: Meta, args) -> HeredocTemplateRule:
        return HeredocTemplateRule(args, meta)

    @v_args(meta=True)
    def heredoc_template_trim(self, meta: Meta, args) -> HeredocTrimTemplateRule:
        return HeredocTrimTemplateRule(args, meta)

    @v_args(meta=True)
    def expr_term(self, meta: Meta, args) -> ExprTermRule:
        return ExprTermRule(args, meta)

    def _extract_nl_prefix(self, token):
        """Strip leading newlines from a token value.

        If the token contains a newline prefix (from the postlexer merging a
        line-continuation newline into the operator token), strip it and
        return a NewLineOrCommentRule.  Otherwise return None.
        """
        value = str(token.value)
        stripped = value.lstrip("\n \t")
        if len(stripped) == len(value):
            return None
        nl_text = value[: len(value) - len(stripped)]
        token.set_value(stripped)
        if self.discard_new_line_or_comments:
            return None
        return NewLineOrCommentRule.from_string(nl_text)

    @v_args(meta=True)
    def conditional(self, meta: Meta, args) -> ConditionalRule:
        # args: [condition, QMARK, NL?, if_true, NL?, COLON, NL?, if_false]
        # QMARK is at index 1 — check for NL prefix from the postlexer
        qmark_token = args[1]
        nl_rule = self._extract_nl_prefix(qmark_token)
        if nl_rule is not None:
            args = list(args)
            args.insert(1, nl_rule)
        return ConditionalRule(args, meta)

    @v_args(meta=True)
    def binary_operator(self, meta: Meta, args) -> BinaryOperatorRule:
        return BinaryOperatorRule(args, meta)

    @v_args(meta=True)
    def binary_term(self, meta: Meta, args) -> BinaryTermRule:
        # args: [BinaryOperatorRule, NL?, ExprTermRule]
        # The operator's token may contain a NL prefix from the postlexer
        op_rule = args[0]
        nl_rule = self._extract_nl_prefix(op_rule.token)
        if nl_rule is not None:
            args = [nl_rule] + list(args)
        return BinaryTermRule(args, meta)

    @v_args(meta=True)
    def unary_op(self, meta: Meta, args) -> UnaryOpRule:
        return UnaryOpRule(args, meta)

    @v_args(meta=True)
    def binary_op(self, meta: Meta, args) -> BinaryOpRule:
        return BinaryOpRule(args, meta)

    @v_args(meta=True)
    def tuple(self, meta: Meta, args) -> TupleRule:
        return TupleRule(args, meta)

    @v_args(meta=True)
    def object(self, meta: Meta, args) -> ObjectRule:
        return ObjectRule(args, meta)

    @v_args(meta=True)
    def object_elem(self, meta: Meta, args) -> ObjectElemRule:
        return ObjectElemRule(args, meta)

    @v_args(meta=True)
    def object_elem_key(self, meta: Meta, args):
        expr = args[0]
        # Simple literals (identifier, string, int, float) wrapped in ExprTermRule
        if isinstance(expr, ExprTermRule) and len(expr.children) == 5:
            inner = expr.children[2]  # position 2 in [None, None, inner, None, None]
            if isinstance(
                inner,
                (
                    IdentifierRule,
                    StringRule,
                    IntLitRule,
                    FloatLitRule,
                    LiteralValueRule,
                ),
            ):
                # Convert literal_value to identifier for dict key compatibility
                if isinstance(inner, LiteralValueRule):
                    inner = IdentifierRule([NAME(inner.token.value)], meta)
                return ObjectElemKeyRule([inner], meta)
        # Any other expression (parenthesized or bare)
        return ObjectElemKeyExpressionRule([expr], meta)

    @v_args(meta=True)
    def arguments(self, meta: Meta, args) -> ArgumentsRule:
        return ArgumentsRule(args, meta)

    @v_args(meta=True)
    def function_call(self, meta: Meta, args) -> FunctionCallRule:
        return FunctionCallRule(args, meta)

    @v_args(meta=True)
    def index_expr_term(self, meta: Meta, args) -> IndexExprTermRule:
        return IndexExprTermRule(args, meta)

    @v_args(meta=True)
    def braces_index(self, meta: Meta, args) -> SqbIndexRule:
        return SqbIndexRule(args, meta)

    @v_args(meta=True)
    def short_index(self, meta: Meta, args) -> ShortIndexRule:
        return ShortIndexRule(args, meta)

    @v_args(meta=True)
    def get_attr(self, meta: Meta, args) -> GetAttrRule:
        # Convert literal_value (true/false/null) to identifier in attr access
        if len(args) >= 2 and isinstance(args[1], LiteralValueRule):
            args = list(args)
            args[1] = IdentifierRule([NAME(args[1].token.value)], meta)
        return GetAttrRule(args, meta)

    @v_args(meta=True)
    def get_attr_expr_term(self, meta: Meta, args) -> GetAttrExprTermRule:
        return GetAttrExprTermRule(args, meta)

    @v_args(meta=True)
    def attr_splat(self, meta: Meta, args) -> AttrSplatRule:
        return AttrSplatRule(args, meta)

    @v_args(meta=True)
    def attr_splat_expr_term(self, meta: Meta, args) -> AttrSplatExprTermRule:
        return AttrSplatExprTermRule(args, meta)

    @v_args(meta=True)
    def full_splat(self, meta: Meta, args) -> FullSplatRule:
        return FullSplatRule(args, meta)

    @v_args(meta=True)
    def full_splat_expr_term(self, meta: Meta, args) -> FullSplatExprTermRule:
        return FullSplatExprTermRule(args, meta)

    @v_args(meta=True)
    def for_tuple_expr(self, meta: Meta, args) -> ForTupleExprRule:
        return ForTupleExprRule(args, meta)

    @v_args(meta=True)
    def for_object_expr(self, meta: Meta, args) -> ForObjectExprRule:
        return ForObjectExprRule(args, meta)

    @v_args(meta=True)
    def for_intro(self, meta: Meta, args) -> ForIntroRule:
        return ForIntroRule(args, meta)

    @v_args(meta=True)
    def for_cond(self, meta: Meta, args) -> ForCondRule:
        return ForCondRule(args, meta)

    @v_args(meta=True)
    def template_if_start(self, meta: Meta, args) -> TemplateIfStartRule:
        return TemplateIfStartRule(args, meta)

    @v_args(meta=True)
    def template_else(self, meta: Meta, args) -> TemplateElseRule:
        return TemplateElseRule(args, meta)

    @v_args(meta=True)
    def template_endif(self, meta: Meta, args) -> TemplateEndifRule:
        return TemplateEndifRule(args, meta)

    @v_args(meta=True)
    def template_for_start(self, meta: Meta, args) -> TemplateForStartRule:
        return TemplateForStartRule(args, meta)

    @v_args(meta=True)
    def template_endfor(self, meta: Meta, args) -> TemplateEndforRule:
        return TemplateEndforRule(args, meta)

    @v_args(meta=True)
    def template_string(self, meta: Meta, args) -> TemplateStringRule:
        return TemplateStringRule(args, meta)
