"""
Flow Reasoning Engine with Fixed-Point Forcing (FPF).

Synthesized from:
"Flow Reasoning Models: Turning Flows Into Efficient Recurrent Reasoners"
(Georgia Tech / MIT / MIT-IBM / IBM Research, Sep 2026).

Treats reasoning as iterative solution refinement toward a stable attractor.
Rather than committing irrevocably to one-shot greedy decisions, candidate
browser actions and spatial grounding coordinates undergo recurrent flow
refinement over held states:
    s^{(k+1)} = Refine(s^{(k)}, Observation)
until converging to a globally consistent fixed point.
"""

from __future__ import annotations

import math
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple
from miniswardbower.core.schemas import BrowserAction, BrowserActionType, InteractiveElement, PrunedAXTree


@dataclass
class FlowRefinementState:
    """Represents the recurrent held-state evolving toward a stable attractor."""
    step_index: int
    candidate_action: BrowserAction
    spatial_coords: Tuple[float, float]
    confidence: float
    residual_norm: float = 1.0
    converged: bool = False
    iterations_run: int = 0
    refinement_trace: List[Dict[str, Any]] = field(default_factory=list)


class FlowReasoningRefiner:
    """
    Recurrent Flow Reasoner with Fixed-Point Iteration for Browser Grounding & Actions.
    Iteratively refines candidate coordinates and parameter schemas to prevent cascade failures.
    """

    def __init__(
        self,
        max_refinement_depth: int = 4,
        convergence_threshold: float = 0.05,
        damping_factor: float = 0.65,
    ):
        self.max_depth = max_refinement_depth
        self.threshold = convergence_threshold
        self.damping = damping_factor

    @staticmethod
    def _find_element(tree: Optional[PrunedAXTree], target_id: Optional[str]) -> Optional[InteractiveElement]:
        if not tree or not target_id:
            return None
        for el in tree.elements:
            if el.id == target_id or el.selector == target_id:
                return el
        return None

    def refine_action(
        self,
        initial_action: BrowserAction,
        tree: Optional[PrunedAXTree] = None,
        viewport_dims: Tuple[int, int] = (1280, 900),
    ) -> FlowRefinementState:
        """
        Runs recurrent fixed-point refinement on candidate action and coordinates.
        Converges candidate toward the optimal, collision-free attractor.
        """
        t0 = time.perf_counter()
        vw, vh = viewport_dims

        # Extract initial spatial coordinates
        cur_x, cur_y = 0.0, 0.0
        if getattr(initial_action, "coords", None):
            cur_x, cur_y = float(initial_action.coords[0]), float(initial_action.coords[1])
        elif tree and initial_action.target:
            target_el = self._find_element(tree, initial_action.target)
            if target_el and target_el.bbox:
                bx, by, bw, bh = target_el.bbox
                cur_x, cur_y = bx + bw / 2.0, by + bh / 2.0
            else:
                cur_x, cur_y = float(vw / 2.0), float(vh / 2.0)
        else:
            cur_x, cur_y = float(vw / 2.0), float(vh / 2.0)

        cur_action = initial_action.model_copy(deep=True)
        cur_confidence = 0.50

        state = FlowRefinementState(
            step_index=0,
            candidate_action=cur_action,
            spatial_coords=(cur_x, cur_y),
            confidence=cur_confidence,
        )

        trace = []

        # Recurrent refinement loop
        for depth in range(1, self.max_depth + 1):
            prev_x, prev_y = cur_x, cur_y
            prev_conf = cur_confidence

            # 1. Attractor gradient from DOM layout
            target_el: Optional[InteractiveElement] = None
            if tree and cur_action.target:
                target_el = self._find_element(tree, cur_action.target)

            attractor_x, attractor_y = prev_x, prev_y
            target_valid = False

            if target_el:
                target_valid = True
                if target_el.bbox:
                    bx, by, bw, bh = target_el.bbox
                    attractor_x = bx + bw / 2.0
                    attractor_y = by + bh / 2.0

                # Type mismatch refinement (e.g. TYPE requested on non-editable tag)
                if cur_action.op == BrowserActionType.TYPE:
                    if target_el.tag not in ("input", "textarea") and target_el.element_type not in ("textbox", "search"):
                        # Search for adjacent editable child/sibling in tree
                        for alt_el in tree.elements:
                            if alt_el.tag in ("input", "textarea") and alt_el.bbox:
                                cur_action.target = alt_el.id
                                cur_action.selector = alt_el.selector
                                abx, aby, abw, abh = alt_el.bbox
                                attractor_x = abx + abw / 2.0
                                attractor_y = aby + abh / 2.0
                                break

            # 2. Viewport boundary attraction (keep coordinates within visible viewport)
            attractor_x = max(10.0, min(float(vw - 10.0), attractor_x))
            attractor_y = max(10.0, min(float(vh - 10.0), attractor_y))

            # 3. Fixed-point update with damping
            dx = (attractor_x - cur_x) * self.damping
            dy = (attractor_y - cur_y) * self.damping

            cur_x += dx
            cur_y += dy

            # 4. Confidence refinement
            if target_valid:
                cur_confidence = min(0.99, cur_confidence + 0.15)
            else:
                cur_confidence = max(0.20, cur_confidence - 0.10)

            # Residual norm
            delta_dist = math.sqrt(dx * dx + dy * dy)
            residual = delta_dist / max(1.0, math.sqrt(float(vw * vw + vh * vh)))

            trace.append({
                "depth": depth,
                "coords": (round(cur_x, 2), round(cur_y, 2)),
                "residual": round(residual, 5),
                "confidence": round(cur_confidence, 3),
            })

            if residual < self.threshold:
                state.converged = True
                state.iterations_run = depth
                state.residual_norm = residual
                break

        if not state.converged:
            state.iterations_run = self.max_depth
            state.residual_norm = residual

        cur_action.coords = (cur_x, cur_y)
        state.candidate_action = cur_action
        state.spatial_coords = (cur_x, cur_y)
        state.confidence = cur_confidence
        state.refinement_trace = trace

        return state