"""
A1111 Prompt Parser (Recursive)

Uses Lark grammar parser like A1111 for proper nested syntax support.

Supported Syntax:
- Emphasis: (text:1.2), (text), [text]
- Scheduling: [from:to:when], [add:when], [remove::when]
- Alternation: [A|B|C]
- BREAK: Splits prompt into isolated chunks

Step vs Percentage:
- If 'when' has a decimal point (e.g., 0.5), it's a percentage
- If 'when' is an integer (e.g., 10), it's a step count (requires steps input)
"""

import re
from dataclasses import dataclass, field
from typing import List, Optional, Tuple, Dict

# Try to import lark, fall back to simple parsing if not available
try:
    import lark

    HAS_LARK = True
except ImportError:
    HAS_LARK = False


@dataclass
class WeightedPart:
    """Represents a parsed prompt segment with weight and timing info."""

    text: str
    weight: float = 1.0
    start_percent: float = 0.0
    end_percent: float = 1.0
    is_alternating: bool = False
    alternating_options: List[str] = field(default_factory=list)

    def __str__(self) -> str:
        """Return the text content for string conversion."""
        return self.text


# A1111-style Lark grammar
GRAMMAR = r"""
!start: (prompt | /[][():]/+)*
prompt: (emphasized | scheduled | alternate | plain | WHITESPACE)*
!emphasized: "(" prompt ")"
        | "(" prompt ":" prompt ")"
        | "[" prompt "]"
scheduled: "[" [prompt ":"] prompt ":" [WHITESPACE] NUMBER [WHITESPACE] "]"
alternate: "[" prompt ("|" [prompt])+ "]"
WHITESPACE: /\s+/
plain: /([^\\\\[\]():|]|\\.)+/
%import common.SIGNED_NUMBER -> NUMBER
"""


def parse_a1111_prompt(
    prompt: str, steps: Optional[int] = None
) -> List[List[WeightedPart]]:
    """
    Main entry point.

    Args:
        prompt: The prompt string to parse
        steps: Total sampling steps (needed for step-based scheduling)

    Returns: List[List[WeightedPart]]
        Outer list = BREAK chunks
        Inner list = Parts with timing/weight info
    """
    # Split by BREAK first
    break_chunks = re.split(r"\s*\bBREAK\b\s*", prompt)

    if HAS_LARK:
        return [_parse_chunk_lark(chunk, steps) for chunk in break_chunks]
    else:
        return [_parse_chunk_simple(chunk, steps) for chunk in break_chunks]


def get_prompt_schedule(prompt: str, steps: int) -> List[Tuple[int, str]]:
    """
    A1111-style prompt schedule generation.

    Returns list of (end_at_step, prompt_text) tuples representing
    when each prompt variant should be used.

    This properly handles nested alternation and scheduling.
    """
    if not HAS_LARK:
        # Fallback for simple prompts
        return [(steps, prompt)]

    # Split by BREAK first, process each chunk, then rejoin
    break_chunks = re.split(r"\s*\bBREAK\b\s*", prompt)

    try:
        parser = lark.Lark(GRAMMAR)
        trees = [parser.parse(chunk) for chunk in break_chunks]
    except lark.exceptions.LarkError:
        return [(steps, prompt)]

    # Collect all relevant steps from all chunks
    all_steps = _collect_steps(steps, trees)

    # Generate prompt text for each step
    schedule = []
    for step in all_steps:
        chunk_texts = []
        for tree in trees:
            text = _at_step(step, steps, tree)
            if text.strip():
                chunk_texts.append(text)
        full_prompt = " BREAK ".join(chunk_texts)
        schedule.append((step, full_prompt))

    return schedule


def _collect_steps(steps: int, trees) -> List[int]:
    """Collect all steps where the prompt changes."""
    result = {steps}  # Always include final step

    class CollectSteps(lark.Visitor):
        def scheduled(self, tree):
            s = str(tree.children[-2])
            v = float(s)
            if "." in s:
                # Percentage
                step_num = int(v * steps)
            else:
                # Step count
                step_num = int(v)
            step_num = max(1, min(steps, step_num))
            if step_num >= 1:
                result.add(step_num)

        def alternate(self, tree):
            # Alternation changes every step
            result.update(range(1, steps + 1))

    for tree in trees:
        CollectSteps().visit(tree)

    return sorted(result)


def _at_step(step: int, total_steps: int, tree) -> str:
    """
    Generate the prompt text for a specific step.

    This follows A1111's approach of walking the tree and picking
    the correct option at each node based on the current step.
    """

    class AtStep(lark.Transformer):
        def scheduled(self, args):
            # A1111 pattern: before, after, _, when, _ = args
            # Use positional access like A1111
            before = args[0] if len(args) > 0 else None
            after = args[1] if len(args) > 1 else None
            when_token = args[3] if len(args) > 3 else None

            # Fallback: find NUMBER token if not at expected position
            if when_token is None or not (
                isinstance(when_token, lark.Token) and when_token.type == "NUMBER"
            ):
                for arg in args:
                    if isinstance(arg, lark.Token) and arg.type == "NUMBER":
                        when_token = arg
                        break

            if when_token is None:
                for a in args:
                    if a is not None:
                        yield a
                return

            when_str = str(when_token)
            if "." in when_str:
                when_step = int(float(when_str) * total_steps)
            else:
                when_step = int(float(when_str))

            # A1111: yield before if step <= when, else yield after
            if step <= when_step:
                if before is not None:
                    yield before
            else:
                if after is not None:
                    yield after

        def alternate(self, args):
            # Filter out empty/None
            options = [a for a in args if a]
            if not options:
                yield ""
                return
            idx = (step - 1) % len(options)
            yield options[idx]

        def emphasized(self, args):
            # Keep emphasis syntax in output
            if len(args) == 1:
                for a in args:
                    yield a
            elif len(args) >= 2:
                # (text:weight) - preserve weight
                yield "("
                yield from self._flatten_arg(args[0])
                yield ":"
                yield from self._flatten_arg(args[1])
                yield ")"
            else:
                for a in args:
                    yield a

        def _flatten_arg(self, arg):
            if isinstance(arg, str):
                yield arg
            elif isinstance(arg, lark.Token):
                yield str(arg)
            elif hasattr(arg, "__iter__"):
                for item in arg:
                    yield from self._flatten_arg(item)
            elif arg is not None:
                yield str(arg)

        def plain(self, args):
            yield args[0].value if args else ""

        def start(self, args):
            def flatten(x):
                if isinstance(x, str):
                    yield x
                elif isinstance(x, lark.Token):
                    yield str(x)
                elif x is None or x == ():
                    pass
                elif hasattr(x, "__iter__"):
                    for item in x:
                        yield from flatten(item)
                else:
                    yield str(x)

            return "".join(flatten(args))

        def __default__(self, data, children, meta):
            for child in children:
                yield child

    return AtStep().transform(tree)


def _parse_chunk_lark(chunk: str, steps: Optional[int]) -> List[WeightedPart]:
    """Parse a single chunk using Lark grammar."""
    try:
        parser = lark.Lark(GRAMMAR)
        tree = parser.parse(chunk)
        return _extract_parts_from_tree(tree, steps)
    except lark.exceptions.LarkError:
        # Fallback to simple parsing on error
        return _parse_chunk_simple(chunk, steps)


def _convert_when_to_percent(when_str: str, steps: Optional[int]) -> float:
    """
    Convert a 'when' value to a percentage.
    - If has decimal point: treat as percentage (0.5 = 50%)
    - If integer >= 1: treat as step count, convert using steps param
    """
    val = float(when_str)

    if "." in when_str:
        # Percentage
        return max(0.0, min(1.0, val))
    else:
        # Step count
        if steps is not None and steps > 0:
            return max(0.0, min(1.0, val / steps))
        else:
            # No steps provided, treat large values as percentage anyway
            if val >= 1.0:
                return 1.0  # Can't convert without steps
            return val


def _extract_text(item) -> str:
    """Recursively extract text from WeightedPart or nested structures."""
    if isinstance(item, WeightedPart):
        if item.is_alternating and item.alternating_options:
            # For alternating, return options joined (will be handled per-step elsewhere)
            return "|".join(item.alternating_options)
        return item.text
    elif isinstance(item, list):
        return "".join(_extract_text(x) for x in item)
    elif HAS_LARK and isinstance(item, lark.Token):
        return str(item)
    elif item is None:
        return ""
    else:
        return str(item)


def _extract_parts_from_tree(tree, steps: Optional[int]) -> List[WeightedPart]:
    """Extract WeightedParts from parsed Lark tree."""
    parts = []

    class PartExtractor(lark.Transformer):
        def plain(self, args):
            text = args[0].value if args else ""
            return WeightedPart(text=text, weight=1.0)

        def emphasized(self, args):
            if len(args) == 1:
                # Just (text) or [text]
                inner = args[0]
                if isinstance(inner, WeightedPart):
                    inner.weight *= 1.1
                    return inner
                return WeightedPart(text=_extract_text(inner), weight=1.1)
            elif len(args) == 2:
                # (text:weight)
                inner, weight_part = args
                try:
                    w = float(_extract_text(weight_part))
                except:
                    w = 1.0
                if isinstance(inner, WeightedPart):
                    inner.weight *= w
                    return inner
                return WeightedPart(text=_extract_text(inner), weight=w)
            return args[0] if args else WeightedPart(text="")

        def scheduled(self, args):
            # [before:after:when] or [:after:when] or [before::when]
            # args structure varies
            when_str = _extract_text(args[-2]) if len(args) >= 2 else "0.5"
            when_pct = _convert_when_to_percent(when_str, steps)

            if len(args) >= 3:
                before = args[0] or ""
                after = args[1] or ""
            else:
                before = ""
                after = args[0] or ""

            # Create two parts: before and after
            result = []
            if before:
                result.append(
                    WeightedPart(
                        text=_extract_text(before),
                        start_percent=0.0,
                        end_percent=when_pct,
                    )
                )
            if after:
                result.append(
                    WeightedPart(
                        text=_extract_text(after),
                        start_percent=when_pct,
                        end_percent=1.0,
                    )
                )
            return result

        def alternate(self, args):
            options = [_extract_text(a) if a else "" for a in args]
            return WeightedPart(
                text="", is_alternating=True, alternating_options=options
            )

        def prompt(self, args):
            return args

        def start(self, args):
            return args

    try:
        result = PartExtractor().transform(tree)

        # Flatten result
        def flatten(items):
            for item in items:
                if isinstance(item, list):
                    yield from flatten(item)
                elif isinstance(item, WeightedPart):
                    yield item
                elif item is not None:
                    yield WeightedPart(text=str(item))

        parts = list(flatten(result if isinstance(result, list) else [result]))
    except Exception:
        parts = [WeightedPart(text=tree)]

    return parts if parts else [WeightedPart(text="")]


# ============ SIMPLE FALLBACK PARSER (no Lark) ============

# Regex patterns for simple parser
re_schedule_switch = re.compile(r"^\[([^:\[\]]*):([^:\[\]]*):([0-9.]+)\]$")
re_schedule_remove = re.compile(r"^\[([^:\[\]]*)::([0-9.]+)\]$")
re_schedule_add = re.compile(r"^\[([^:\[\]]*):([0-9.]+)\]$")
re_alternation = re.compile(r"^\[([^|\]]+(?:\|[^|\]]+)+)\]$")

re_attention = re.compile(
    r"""
\\\\|
\\\)|
\\\[|
\\]|
\\\\|
\\|
\[[^\]|]*\|[^\]|]*\]|
\[[^\]:]*:[^\]]*\]|
\(|
\[|
:\s*([+-]?[.\d]+)\s*\)|
\)|
]|
[^\\()\[\]:]+|
:
""",
    re.X,
)


def _parse_chunk_simple(chunk: str, steps: Optional[int]) -> List[WeightedPart]:
    """Simple regex-based parsing (fallback when Lark not available)."""
    parsed_weights = _parse_weight_simple(chunk)
    parts = []

    for text, weight in parsed_weights:
        # Check for scheduling patterns
        m_switch = re_schedule_switch.search(text)
        if m_switch:
            p_from, p_to, p_when = m_switch.groups()
            val = _convert_when_to_percent(p_when, steps)
            parts.append(
                WeightedPart(
                    text=p_from, weight=weight, start_percent=0.0, end_percent=val
                )
            )
            parts.append(
                WeightedPart(
                    text=p_to, weight=weight, start_percent=val, end_percent=1.0
                )
            )
            continue

        m_remove = re_schedule_remove.search(text)
        if m_remove:
            p_from, p_when = m_remove.groups()
            val = _convert_when_to_percent(p_when, steps)
            parts.append(
                WeightedPart(
                    text=p_from, weight=weight, start_percent=0.0, end_percent=val
                )
            )
            continue

        m_add = re_schedule_add.search(text)
        if m_add:
            p_to, p_when = m_add.groups()
            val = _convert_when_to_percent(p_when, steps)
            parts.append(
                WeightedPart(
                    text=p_to, weight=weight, start_percent=val, end_percent=1.0
                )
            )
            continue

        m_alt = re_alternation.search(text)
        if m_alt:
            options = m_alt.group(1).split("|")
            parts.append(
                WeightedPart(
                    text=text,
                    weight=weight,
                    is_alternating=True,
                    alternating_options=options,
                )
            )
            continue

        parts.append(WeightedPart(text=text, weight=weight))

    return parts if parts else [WeightedPart(text="")]


def _parse_weight_simple(text: str) -> List[tuple]:
    """Parse emphasis weights using regex."""
    res = []
    round_brackets = []
    square_brackets = []

    round_bracket_multiplier = 1.1
    square_bracket_multiplier = 1 / 1.1

    def multiply_range(start_position, multiplier):
        for p in range(start_position, len(res)):
            res[p][1] *= multiplier

    for m in re_attention.finditer(text):
        token = m.group(0)
        weight = m.group(1)

        if token.startswith("\\"):
            res.append([token[1:], 1.0])
        elif (token.startswith("[") and "|" in token and token.endswith("]")) or (
            token.startswith("[") and ":" in token and token.endswith("]")
        ):
            res.append([token, 1.0])
        elif token == "(":
            round_brackets.append(len(res))
        elif token == "[":
            square_brackets.append(len(res))
        elif weight is not None and round_brackets:
            multiply_range(round_brackets.pop(), float(weight))
        elif token == ")" and round_brackets:
            multiply_range(round_brackets.pop(), round_bracket_multiplier)
        elif token == "]" and square_brackets:
            multiply_range(square_brackets.pop(), square_bracket_multiplier)
        else:
            res.append([token, 1.0])

    for pos in round_brackets:
        multiply_range(pos, round_bracket_multiplier)
    for pos in square_brackets:
        multiply_range(pos, square_bracket_multiplier)

    if not res:
        res = [["", 1.0]]

    # Merge consecutive identical weights
    i = 0
    while i + 1 < len(res):
        if res[i][1] == res[i + 1][1]:
            res[i][0] += res[i + 1][0]
            res.pop(i + 1)
        else:
            i += 1

    return [(t, w) for t, w in res]