	try:
		return json.dumps(obj, ensure_ascii=False)
	except Exception:
		return str(obj)

def _extract_between(s: str, start_token: str, end_token: str) -> Optional[str]:
	"""Extract substring between tokens with better error handling."""
	try:
		start = s.find(start_token)
		if start == -1:
			return None
		end = s.find(end_token, start)
		if end == -1:
			end = len(s)
		return s[start:end][:256]
	except Exception:
		return None

def _extract_key(s: str, key: str) -> Optional[str]:
	"""Extract JSON key value with validation."""
	try:
		pattern = r"\b" + re.escape(key) + r"\"\s*:\s*\"([^\"]+)\""
		m = re.search(pattern, s)
		if not m:
			return None
		value = m.group(1)[:256]
		return value if value.strip() else None
	except Exception:
		return None

def _extract_argv0(s: str) -> Optional[str]:
	"""Extract first argument from argv array."""
	try:
		m = re.search(r"\[\s*\"([^\"]+)\"", s)
		if not m:
			return None
		cmd = m.group(1)[:128]
		return cmd if cmd.strip() else None
	except Exception:
		return None


		if end == -1:
			end = len(s)
		return s[start:end][:256]
	except Exception:
		return None

def _extract_key(s: str, key: str) -> Optional[str]:
	"""Extract JSON key value with validation."""
	try:
		pattern = r"\b" + re.escape(key) + r"\"\s*:\s*\"([^\"]+)\""
		m = re.search(pattern, s)
		if not m:
			return None
		value = m.group(1)[:256]
		return value if value.strip() else None
	except Exception:
		return None

def _extract_argv0(s: str) -> Optional[str]:
	"""Extract first argument from argv array."""
	try:
		m = re.search(r"\[\s*\"([^\"]+)\"", s)
		if not m:
			return None
		cmd = m.group(1)[:128]
		return cmd if cmd.strip() else None
	except Exception:
		return None


from __future__ import annotations
import logging
import json
import re
from typing import Dict, Any, Tuple, Optional

class NextStatePredictor:
	"""Next-state predictor for world model rollouts.

	Predicts observation changes and side effects from actions:
	- DOM: navigation, form states, click effects
	- CLI: commands, working directory, output types
	- Generic: action metadata and timestamps
	"""

	def predict_next(self, obs: Dict[str, Any], action: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
		try:
			obs_next: Dict[str, Any] = dict(obs or {})
			effects: Dict[str, Any] = {}

			# Convert action to string for pattern matching
			a_json = _safe_json(action)

			# DOM action handling
			if self._is_dom_action(a_json):
				obs_next, effects = self._handle_dom_action(obs_next, action, a_json)

			# CLI action handling
			elif self._is_cli_action(a_json):
				obs_next, effects = self._handle_cli_action(obs_next, action, a_json)

			# Record action metadata
			self._record_action_metadata(obs_next, action)

			return obs_next, effects
		except Exception as e:
			# Log error but continue with safe fallback
			effects = {"error": str(e)}
			return dict(obs or {}), effects

	def _is_dom_action(self, a_json: str) -> bool:
		"""Check if action is DOM-related."""
		dom_patterns = [
			"browser.read",
			"browser.click",
			"browser.fill",
			"url",
			"selector",
			"form",
			"button",
			"input"
		]
		return any(p in a_json for p in dom_patterns)

	def _is_cli_action(self, a_json: str) -> bool:
		"""Check if action is CLI-related."""
		cli_patterns = [
			"argv",
			"tool",


			# Convert action to string for pattern matching
			a_json = _safe_json(action)

			# DOM action handling
			if self._is_dom_action(a_json):
				obs_next, effects = self._handle_dom_action(obs_next, action, a_json)

			# CLI action handling
			elif self._is_cli_action(a_json):
				obs_next, effects = self._handle_cli_action(obs_next, action, a_json)

			# Record action metadata
			self._record_action_metadata(obs_next, action)

			return obs_next, effects
		except Exception as e:
			# Log error but continue with safe fallback
			effects = {"error": str(e)}
			return dict(obs or {}), effects

	def _is_dom_action(self, a_json: str) -> bool:
		"""Check if action is DOM-related."""
		dom_patterns = [
			"browser.read",
			"browser.click",
			"browser.fill",
			"url",
			"selector",
			"form",
			"button",
			"input"
		]
		return any(p in a_json for p in dom_patterns)

	def _is_cli_action(self, a_json: str) -> bool:
		"""Check if action is CLI-related."""
		cli_patterns = [
			"argv",
			"tool",
			"cmd",
			"shell",
			"terminal",
			"exec"
		]
		return any(p in a_json for p in cli_patterns)

	def _handle_dom_action(
		self, obs_next: Dict[str, Any], action: Dict[str, Any], a_json: str
	) -> Tuple[Dict[str, Any], Dict[str, Any]]:
		"""Handle DOM-related actions and predict state changes."""
		effects: Dict[str, Any] = {}

		# Extract basic DOM info


			return obs_next, effects
		except Exception as e:
			# Log error but continue with safe fallback
			effects = {"error": str(e)}
			return dict(obs or {}), effects

	def _is_dom_action(self, a_json: str) -> bool:
		"""Check if action is DOM-related."""
		dom_patterns = [
			"browser.read",
			"browser.click",
			"browser.fill",
			"url",
			"selector",
			"form",
			"button",
			"input"
		]
		return any(p in a_json for p in dom_patterns)

	def _is_cli_action(self, a_json: str) -> bool:
		"""Check if action is CLI-related."""
		cli_patterns = [
			"argv",
			"tool",
			"cmd",
			"shell",
			"terminal",
			"exec"
		]
		return any(p in a_json for p in cli_patterns)

	def _handle_dom_action(
		self, obs_next: Dict[str, Any], action: Dict[str, Any], a_json: str
	) -> Tuple[Dict[str, Any], Dict[str, Any]]:
		"""Handle DOM-related actions and predict state changes."""
		effects: Dict[str, Any] = {}

		# Extract basic DOM info
		url = _extract_between(a_json, "http", '"')
		selector = _extract_key(a_json, "selector")

		# Track navigation
		if url:
			obs_next["last_url"] = url
			effects["url"] = url
			# Clear form state on navigation
			obs_next.pop("form_state", None)

		# Track element interaction
		if selector:


			"url",
			"selector",
			"form",
			"button",
			"input"
		]
		return any(p in a_json for p in dom_patterns)

	def _is_cli_action(self, a_json: str) -> bool:
		"""Check if action is CLI-related."""
		cli_patterns = [
			"argv",
			"tool",
			"cmd",
			"shell",
			"terminal",
			"exec"
		]
		return any(p in a_json for p in cli_patterns)

	def _handle_dom_action(
		self, obs_next: Dict[str, Any], action: Dict[str, Any], a_json: str
	) -> Tuple[Dict[str, Any], Dict[str, Any]]:
		"""Handle DOM-related actions and predict state changes."""
		effects: Dict[str, Any] = {}

		# Extract basic DOM info
		url = _extract_between(a_json, "http", '"')
		selector = _extract_key(a_json, "selector")

		# Track navigation
		if url:
			obs_next["last_url"] = url
			effects["url"] = url
			# Clear form state on navigation
			obs_next.pop("form_state", None)

		# Track element interaction
		if selector:
			obs_next["last_selector"] = selector
			effects["selector"] = selector

			# Predict form state changes
			if "fill" in a_json:
				form_state = obs_next.get("form_state", {})
				value = _extract_key(a_json, "value")
				if value:
					form_state[selector] = value
					obs_next["form_state"] = form_state
					effects["form_update"] = {selector: value}

			# Predict click effects
			elif "click" in a_json:
				effects["clicked"] = selector
				# Clear form state if it looks like a submit
				if re.search(r"submit|login|signup|send|save", selector.lower()):
					obs_next.pop("form_state", None)
					effects["form_submit"] = True

		return obs_next, effects

	def _handle_cli_action(
		self, obs_next: Dict[str, Any], action: Dict[str, Any], a_json: str
	) -> Tuple[Dict[str, Any], Dict[str, Any]]:
		"""Handle CLI-related actions and predict command output."""
		effects: Dict[str, Any] = {}

		# Extract command info
		cmd = _extract_argv0(a_json)
		if cmd:
			obs_next["last_cmd"] = cmd
			effects["cmd"] = cmd

			# Predict common command patterns
			if cmd in ["ls", "dir"]:
				effects["output_type"] = "file_list"
			elif cmd in ["cat", "type"]:
				effects["output_type"] = "file_content"
			elif cmd in ["grep", "find"]:
				effects["output_type"] = "search_results"


			# Predict form state changes
			if "fill" in a_json:
				form_state = obs_next.get("form_state", {})
				value = _extract_key(a_json, "value")
				if value:
					form_state[selector] = value
					obs_next["form_state"] = form_state
					effects["form_update"] = {selector: value}

			# Predict click effects
			elif "click" in a_json:
				effects["clicked"] = selector
				# Clear form state if it looks like a submit
				if re.search(r"submit|login|signup|send|save", selector.lower()):
					obs_next.pop("form_state", None)
					effects["form_submit"] = True

		return obs_next, effects

	def _handle_cli_action(
		self, obs_next: Dict[str, Any], action: Dict[str, Any], a_json: str
	) -> Tuple[Dict[str, Any], Dict[str, Any]]:
		"""Handle CLI-related actions and predict command output."""
		effects: Dict[str, Any] = {}

		# Extract command info
		cmd = _extract_argv0(a_json)
		if cmd:
			obs_next["last_cmd"] = cmd
			effects["cmd"] = cmd

			# Predict common command patterns
			if cmd in ["ls", "dir"]:
				effects["output_type"] = "file_list"
			elif cmd in ["cat", "type"]:
				effects["output_type"] = "file_content"
			elif cmd in ["grep", "find"]:
				effects["output_type"] = "search_results"
			elif cmd in ["git"]:
				effects["output_type"] = "git_status"

			# Track working directory
			if cmd == "cd":
				try:
					args = json.loads(a_json).get("argv", [])
					if len(args) > 1:
						new_dir = args[1]
						obs_next["cwd"] = new_dir
						effects["cwd"] = new_dir
				except Exception:
					pass

		return obs_next, effects

	def _record_action_metadata(self, obs_next: Dict[str, Any], action: Dict[str, Any]) -> None:
		"""Record general action metadata in observation."""
		try:
			if isinstance(action, dict) and action:
				obs_next["last_action_kind"] = next(iter(action.keys()))
				if "timestamp" in action:
					obs_next["last_action_time"] = action["timestamp"]
		except Exception:
			pass

def _safe_json(obj: Any) -> str:
	"""Safely convert object to JSON string."""
	try:
		return json.dumps(obj, ensure_ascii=False)
	except Exception:
		return str(obj)


			elif cmd in ["cat", "type"]:
				effects["output_type"] = "file_content"
			elif cmd in ["grep", "find"]:
				effects["output_type"] = "search_results"
			elif cmd in ["git"]:
				effects["output_type"] = "git_status"

			# Track working directory
			if cmd == "cd":
				try:
					args = json.loads(a_json).get("argv", [])
					if len(args) > 1:
						new_dir = args[1]
						obs_next["cwd"] = new_dir
						effects["cwd"] = new_dir
				except Exception:
					pass

		return obs_next, effects

	def _record_action_metadata(self, obs_next: Dict[str, Any], action: Dict[str, Any]) -> None:
		"""Record general action metadata in observation."""
		try:
			if isinstance(action, dict) and action:
				obs_next["last_action_kind"] = next(iter(action.keys()))
				if "timestamp" in action:
					obs_next["last_action_time"] = action["timestamp"]
		except Exception:
			pass

def _safe_json(obj: Any) -> str:
	"""Safely convert object to JSON string."""
	try:
		return json.dumps(obj, ensure_ascii=False)
	except Exception:
		return str(obj)

def _extract_between(s: str, start_token: str, end_token: str) -> Optional[str]:
	"""Extract substring between tokens with better error handling."""
	try:
		start = s.find(start_token)
		if start == -1:
			return None
		end = s.find(end_token, start)
		if end == -1:
			end = len(s)
		return s[start:end][:256]


from __future__ import annotations
import logging

from typing import Any, Dict

# Minimal schema definition for WM dataset examples
# An example row looks like:
# {
#   "input": {"obs": str, "plan": str, "action": str, "effects": str},
#   "success": int (0/1),
#   "risk": float [0..1]
# }
SCHEMA: Dict[str, Any] = {
	"type": "object",
	"required": ["input", "success", "risk"],
	"properties": {
		"input": {
			"type": "object",
			"required": ["obs", "plan", "action"],
			"properties": {
				"obs": {"type": "string"},
				"plan": {"type": "string"},
				"action": {"type": "string"},
				"effects": {"type": "string"},
			},
		},
		"success": {"type": "integer", "minimum": 0, "maximum": 1},
		"risk": {"type": "number", "minimum": 0.0, "maximum": 1.0},
	},
}

def validate_row(row: Dict[str, Any]) -> bool:
	try:
		inp = row.get("input")
		if not isinstance(inp, dict):
			return False
		if not isinstance(inp.get("obs", ""), str):
			return False
		if not isinstance(inp.get("plan", ""), str):
			return False
		if not isinstance(inp.get("action", ""), str):
			return False
		# effects optional but if present must be string
		if "effects" in inp and not isinstance(inp.get("effects"), str):
			return False
		succ = int(row.get("success", 0))
		if succ not in (0, 1):
			return False
		risk = float(row.get("risk", 0.5))
		if not (0.0 <= risk <= 1.0):
			return False
		return True
	except Exception:
		return False


SCHEMA: Dict[str, Any] = {
	"type": "object",
	"required": ["input", "success", "risk"],
	"properties": {
		"input": {
			"type": "object",
			"required": ["obs", "plan", "action"],
			"properties": {
				"obs": {"type": "string"},
				"plan": {"type": "string"},
				"action": {"type": "string"},
				"effects": {"type": "string"},
			},
		},
		"success": {"type": "integer", "minimum": 0, "maximum": 1},
		"risk": {"type": "number", "minimum": 0.0, "maximum": 1.0},
	},
}

def validate_row(row: Dict[str, Any]) -> bool:
	try:
		inp = row.get("input")
		if not isinstance(inp, dict):
			return False
		if not isinstance(inp.get("obs", ""), str):
			return False
		if not isinstance(inp.get("plan", ""), str):
			return False
		if not isinstance(inp.get("action", ""), str):
			return False
		# effects optional but if present must be string
		if "effects" in inp and not isinstance(inp.get("effects"), str):
			return False
		succ = int(row.get("success", 0))
		if succ not in (0, 1):
			return False
		risk = float(row.get("risk", 0.5))
		if not (0.0 <= risk <= 1.0):
			return False
		return True
	except Exception:
		return False


from __future__ import annotations
import logging

from typing import Dict, Any, List, Optional, Tuple
from dataclasses import dataclass
from pathlib import Path

from .next_state import NextStatePredictor
from .api import WorldModelPrior

@dataclass
class RolloutConfig:
	"""Configuration for rollout simulation."""
	horizon: int = 3
	max_risk_threshold: float = 0.8
	min_success_prob: float = 0.2
	enable_early_stopping: bool = True
	track_state_changes: bool = True

@dataclass
class RolloutStep:
	"""Single step in a rollout simulation."""
	idx: int
	action: Dict[str, Any]
	observation: Dict[str, Any]
	effects: Dict[str, Any]
	metrics: Dict[str, float]
	state_changes: Dict[str, Any]

@dataclass
class RolloutResult:
	"""Result of a rollout simulation."""
	steps: List[RolloutStep]
	cumulative_risk: float
	avg_risk: float
	success_probability: float
	terminated_early: bool
	termination_reason: Optional[str]
	state_validation: Dict[str, Any]  # Validation results for state requirements
	validation_points: List[Dict[str, Any]]  # Results of validation point checks

class RolloutAPI:
	"""API for simulating action plans with configurable horizon.

	Features:
	- k-step lookahead with state prediction
	- Risk and success probability estimation
	- Early stopping on risk/success thresholds
	- State change tracking
	- Detailed metrics per step
	"""

	def __init__(
		self,
		world_model: Optional[WorldModelPrior] = None,
		next_state_predictor: Optional[NextStatePredictor] = None,
		config: Optional[RolloutConfig] = None
	) -> None:
		self.world_model = world_model
		self.next_state_predictor = next_state_predictor or NextStatePredictor()
		self.config = config or RolloutConfig()

	@classmethod
	def from_checkpoint(cls, checkpoint_path: str | Path) -> "RolloutAPI":
		"""Initialize from a saved world model checkpoint."""
		world_model = WorldModelPrior.load(checkpoint_path)
		return cls(world_model=world_model)

	def simulate_sequence(
		self,
		initial_obs: Dict[str, Any],
		plan: Dict[str, Any],
		sequence: Dict[str, Any]
	) -> RolloutResult:
		"""Simulate a sequence with state tracking and validation."""
		# Extract sequence steps
		actions = []
		state_requirements = {}
		validation_points = []

		# Parse sequence
		for step in sequence.get("subgoals", []):
			# Convert step to action
			action = self._step_to_action(step)
			actions.append(action)

			# Collect state requirements
			if step.get("requires"):
				state_requirements[len(actions) - 1] = step["requires"]

			# Collect validation points
			if step.get("validation"):
				validation_points.append({
					"step": len(actions) - 1,
					"validation": step["validation"]
				})

		return self.simulate_plan(
			initial_obs,
			plan,
			actions,
			state_requirements=state_requirements,
			validation_points=validation_points
		)

	def _step_to_action(
		self,
		step: Dict[str, Any]
	) -> Dict[str, Any]:
		"""Convert a sequence step to an action."""
		action = {
			"type": self._infer_action_type(step["description"]),
			"args": {}
		}

		# Extract action details from description
		if action["type"] == "navigate":
			action["args"]["url"] = self._extract_url(step["description"])
		elif action["type"] == "fill":
			action["args"].update(self._extract_form_args(step["description"]))
		elif action["type"] == "click":
			action["args"]["selector"] = self._extract_selector(step["description"])

		return action

	def _infer_action_type(self, description: str) -> str:
		"""Infer action type from step description."""
		description = description.lower()

		if any(w in description for w in ["navigate", "go to", "visit", "open"]):
			return "navigate"
		elif any(w in description for w in ["fill", "enter", "type", "input"]):
			return "fill"
		elif any(w in description for w in ["click", "submit", "press", "select"]):
			return "click"
		else:
			return "read"  # Default to read action

	def _extract_url(self, description: str) -> str:
		"""Extract URL from description."""
		import re

		# Look for URL patterns
		url_pattern = r'https?://[^\s<>"\']+|www\.[^\s<>"\']+\.[^\s<>"\']+'
		urls = re.findall(url_pattern, description)

		if urls:
			return urls[0]
		return ""

	def _extract_selector(self, description: str) -> str:
		"""Extract selector from description."""
		import re

		# Look for quoted selectors
		selector_pattern = r'"([^"]+)"|\'([^\']+)\''
		selectors = re.findall(selector_pattern, description)

		if selectors:
			# Take first non-empty group
			selector = next((s for group in selectors for s in group if s), "")
			if selector:
				return selector

		# Look for common selector patterns
		patterns = [
			r'#[\w-]+',  # ID selectors
			r'\.[\w-]+',  # Class selectors
			r'\[[\w-]+=[^\]]+\]',  # Attribute selectors
			r'button\[type=submit\]',  # Common button selector
			r'input\[type=\w+\]'  # Common input selector
		]

		for pattern in patterns:
			match = re.search(pattern, description)
			if match:
				return match.group(0)

		return ""

	def _extract_form_args(self, description: str) -> Dict[str, str]:
		"""Extract form field arguments from description."""
		import re

		args = {}

		# Look for field-value pairs
		# Pattern: field "value" or field 'value' or field=value
		patterns = [
			r'(\w+)\s*["\']([^"\']+)["\']',
			r'(\w+)\s*=\s*([^\s,]+)'
		]

		for pattern in patterns:
			matches = re.findall(pattern, description)
			for field, value in matches:
				args[field] = value

		# Look for selector and value separately
		if not args:
			selector = self._extract_selector(description)
			if selector:
				# Look for quoted value after selector
				value_match = re.search(r'["\']([^"\']+)["\']', description)
				if value_match:
					args["selector"] = selector
					args["value"] = value_match.group(1)

		return args

	def simulate_plan(
		self,
		initial_obs: Dict[str, Any],
		plan: Dict[str, Any],
		actions: List[Dict[str, Any]],
		state_requirements: Optional[Dict[int, List[str]]] = None,
		validation_points: Optional[List[Dict[str, Any]]] = None
	) -> RolloutResult:
		"""Simulate a sequence of actions from an initial observation."""
		steps: List[RolloutStep] = []
		cumulative_risk = 0.0
		cur_obs = dict(initial_obs)
		terminated_early = False
		termination_reason = None

		# Validate and truncate actions to horizon
		horizon = min(self.config.horizon, len(actions))
		if not actions:
			return RolloutResult(
				steps=[],
				cumulative_risk=0.0,
				avg_risk=0.0,
				success_probability=0.0,
				terminated_early=True,
				termination_reason="No actions provided"
			)

		for idx, action in enumerate(actions[:horizon]):
			# Get world model predictions
			wm_prior = None
			if self.world_model:
				try:
					wm_prior = self.world_model.predict_prior(cur_obs, plan, action)
				except Exception:
					pass

			# Default metrics if WM fails
			metrics = {
				"risk": wm_prior.get("risk", 0.5) if wm_prior else 0.5,
				"success_prob": wm_prior.get("success_prob", 0.5) if wm_prior else 0.5,
				"success_entropy": wm_prior.get("success_entropy", 0.0) if wm_prior else 0.0
			}

			# Predict next state
			next_obs, effects = self.next_state_predictor.predict_next(cur_obs, action)

			# Track state changes if enabled
			state_changes = {}
			if self.config.track_state_changes:
				state_changes = self._detect_state_changes(cur_obs, next_obs)

			# Create step record
			step = RolloutStep(
				idx=idx,
				action=action,
				observation=next_obs,
				effects=effects,
				metrics=metrics,
				state_changes=state_changes
			)
			steps.append(step)

			# Update running metrics
			cumulative_risk += metrics["risk"]
			avg_risk = cumulative_risk / (idx + 1)

			# Check early stopping conditions
			if self.config.enable_early_stopping:
				if avg_risk > self.config.max_risk_threshold:
					terminated_early = True
					termination_reason = f"Average risk {avg_risk:.2f} exceeded threshold {self.config.max_risk_threshold}"
					break
				if metrics["success_prob"] < self.config.min_success_prob:
					terminated_early = True
					termination_reason = f"Success probability {metrics['success_prob']:.2f} below threshold {self.config.min_success_prob}"
					break

			# Update current observation for next iteration
			cur_obs = next_obs

		# Calculate final metrics
		num_steps = len(steps)
		avg_risk = cumulative_risk / num_steps if num_steps > 0 else 0.0
		success_prob = steps[-1].metrics["success_prob"] if steps else 0.0

		return RolloutResult(
			steps=steps,
			cumulative_risk=cumulative_risk,
			avg_risk=avg_risk,
			success_probability=success_prob,
			terminated_early=terminated_early,
			termination_reason=termination_reason
		)

	def _detect_state_changes(
		self,
		prev_obs: Dict[str, Any],
		next_obs: Dict[str, Any]
	) -> Dict[str, Any]:
		"""Detect meaningful changes between observations."""
		changes = {}

		# Track URL changes
		if prev_obs.get("last_url") != next_obs.get("last_url"):
			changes["url_changed"] = {
				"from": prev_obs.get("last_url"),
				"to": next_obs.get("last_url")
			}

		# Track form state changes
		prev_form = prev_obs.get("form_state", {})
		next_form = next_obs.get("form_state", {})
		if prev_form != next_form:
			changes["form_changed"] = {
				"added": {k: v for k, v in next_form.items() if k not in prev_form},
				"removed": {k: v for k, v in prev_form.items() if k not in next_form},
				"modified": {k: {"from": prev_form[k], "to": v} for k, v in next_form.items()
						   if k in prev_form and prev_form[k] != v}
			}

		# Track working directory changes
		if prev_obs.get("cwd") != next_obs.get("cwd"):
			changes["cwd_changed"] = {
				"from": prev_obs.get("cwd"),
				"to": next_obs.get("cwd")
			}

		# Track action type changes
		if prev_obs.get("last_action_kind") != next_obs.get("last_action_kind"):
			changes["action_type_changed"] = {
				"from": prev_obs.get("last_action_kind"),
				"to": next_obs.get("last_action_kind")
			}

		return changes


RolloutConfig

from __future__ import annotations
import logging

from typing import Dict, Any, List, Optional, Tuple
from dataclasses import dataclass
from pathlib import Path

from .next_state import NextStatePredictor
from .api import WorldModelPrior

@dataclass
class RolloutConfig:
	"""Configuration for rollout simulation."""
	horizon: int = 3
	max_risk_threshold: float = 0.8
	min_success_prob: float = 0.2
	enable_early_stopping: bool = True
	track_state_changes: bool = True

@dataclass
class RolloutStep:
	"""Single step in a rollout simulation."""
	idx: int
	action: Dict[str, Any]
	observation: Dict[str, Any]
	effects: Dict[str, Any]
	metrics: Dict[str, float]
	state_changes: Dict[str, Any]

@dataclass
class RolloutResult:
	"""Result of a rollout simulation."""
	steps: List[RolloutStep]
	cumulative_risk: float
	avg_risk: float
	success_probability: float


RolloutStep

from typing import Dict, Any, List, Optional, Tuple
from dataclasses import dataclass
from pathlib import Path

from .next_state import NextStatePredictor
from .api import WorldModelPrior

@dataclass
class RolloutConfig:
	"""Configuration for rollout simulation."""
	horizon: int = 3
	max_risk_threshold: float = 0.8
	min_success_prob: float = 0.2
	enable_early_stopping: bool = True
	track_state_changes: bool = True

@dataclass
class RolloutStep:
	"""Single step in a rollout simulation."""
	idx: int
	action: Dict[str, Any]
	observation: Dict[str, Any]
	effects: Dict[str, Any]
	metrics: Dict[str, float]
	state_changes: Dict[str, Any]

@dataclass
class RolloutResult:
	"""Result of a rollout simulation."""
	steps: List[RolloutStep]
	cumulative_risk: float
	avg_risk: float
	success_probability: float
	terminated_early: bool
	termination_reason: Optional[str]
	state_validation: Dict[str, Any]  # Validation results for state requirements
	validation_points: List[Dict[str, Any]]  # Results of validation point checks

class RolloutAPI:
	"""API for simulating action plans with configurable horizon.

	Features:
	- k-step lookahead with state prediction


RolloutResult

	"""Configuration for rollout simulation."""
	horizon: int = 3
	max_risk_threshold: float = 0.8
	min_success_prob: float = 0.2
	enable_early_stopping: bool = True
	track_state_changes: bool = True

@dataclass
class RolloutStep:
	"""Single step in a rollout simulation."""
	idx: int
	action: Dict[str, Any]
	observation: Dict[str, Any]
	effects: Dict[str, Any]
	metrics: Dict[str, float]
	state_changes: Dict[str, Any]

@dataclass
class RolloutResult:
	"""Result of a rollout simulation."""
	steps: List[RolloutStep]
	cumulative_risk: float
	avg_risk: float
	success_probability: float
	terminated_early: bool
	termination_reason: Optional[str]
	state_validation: Dict[str, Any]  # Validation results for state requirements
	validation_points: List[Dict[str, Any]]  # Results of validation point checks

class RolloutAPI:
	"""API for simulating action plans with configurable horizon.

	Features:
	- k-step lookahead with state prediction
	- Risk and success probability estimation
	- Early stopping on risk/success thresholds
	- State change tracking
	- Detailed metrics per step
	"""

	def __init__(
		self,
		world_model: Optional[WorldModelPrior] = None,
		next_state_predictor: Optional[NextStatePredictor] = None,
		config: Optional[RolloutConfig] = None
	) -> None:
		self.world_model = world_model


RolloutAPI

	action: Dict[str, Any]
	observation: Dict[str, Any]
	effects: Dict[str, Any]
	metrics: Dict[str, float]
	state_changes: Dict[str, Any]

@dataclass
class RolloutResult:
	"""Result of a rollout simulation."""
	steps: List[RolloutStep]
	cumulative_risk: float
	avg_risk: float
	success_probability: float
	terminated_early: bool
	termination_reason: Optional[str]
	state_validation: Dict[str, Any]  # Validation results for state requirements
	validation_points: List[Dict[str, Any]]  # Results of validation point checks

class RolloutAPI:
	"""API for simulating action plans with configurable horizon.

	Features:
	- k-step lookahead with state prediction
	- Risk and success probability estimation
	- Early stopping on risk/success thresholds
	- State change tracking
	- Detailed metrics per step
	"""

	def __init__(
		self,
		world_model: Optional[WorldModelPrior] = None,
		next_state_predictor: Optional[NextStatePredictor] = None,
		config: Optional[RolloutConfig] = None
	) -> None:
		self.world_model = world_model
		self.next_state_predictor = next_state_predictor or NextStatePredictor()
		self.config = config or RolloutConfig()

	@classmethod
	def from_checkpoint(cls, checkpoint_path: str | Path) -> "RolloutAPI":
		"""Initialize from a saved world model checkpoint."""
		world_model = WorldModelPrior.load(checkpoint_path)
		return cls(world_model=world_model)

	def simulate_sequence(
		self,
		initial_obs: Dict[str, Any],
		plan: Dict[str, Any],
		sequence: Dict[str, Any]
	) -> RolloutResult:
		"""Simulate a sequence with state tracking and validation."""
		# Extract sequence steps
		actions = []
		state_requirements = {}
		validation_points = []

		# Parse sequence
		for step in sequence.get("subgoals", []):
			# Convert step to action
			action = self._step_to_action(step)
			actions.append(action)

			# Collect state requirements
			if step.get("requires"):
				state_requirements[len(actions) - 1] = step["requires"]

			# Collect validation points
			if step.get("validation"):
				validation_points.append({
					"step": len(actions) - 1,
					"validation": step["validation"]
				})

		return self.simulate_plan(
			initial_obs,
			plan,
			actions,
			state_requirements=state_requirements,
			validation_points=validation_points
		)

	def _step_to_action(
		self,
		step: Dict[str, Any]
	) -> Dict[str, Any]:
		"""Convert a sequence step to an action."""
		action = {
			"type": self._infer_action_type(step["description"]),
			"args": {}
		}

		# Extract action details from description
		if action["type"] == "navigate":
			action["args"]["url"] = self._extract_url(step["description"])
		elif action["type"] == "fill":
			action["args"].update(self._extract_form_args(step["description"]))
		elif action["type"] == "click":
			action["args"]["selector"] = self._extract_selector(step["description"])

		return action

	def _infer_action_type(self, description: str) -> str:
		"""Infer action type from step description."""
		description = description.lower()

		if any(w in description for w in ["navigate", "go to", "visit", "open"]):
			return "navigate"
		elif any(w in description for w in ["fill", "enter", "type", "input"]):
			return "fill"
		elif any(w in description for w in ["click", "submit", "press", "select"]):
			return "click"
		else:
			return "read"  # Default to read action

	def _extract_url(self, description: str) -> str:
		"""Extract URL from description."""
		import re

		# Look for URL patterns
		url_pattern = r'https?://[^\s<>"\']+|www\.[^\s<>"\']+\.[^\s<>"\']+'
		urls = re.findall(url_pattern, description)

		if urls:
			return urls[0]
		return ""

	def _extract_selector(self, description: str) -> str:
		"""Extract selector from description."""
		import re

		# Look for quoted selectors
		selector_pattern = r'"([^"]+)"|\'([^\']+)\''
		selectors = re.findall(selector_pattern, description)

		if selectors:
			# Take first non-empty group
			selector = next((s for group in selectors for s in group if s), "")
			if selector:
				return selector

		# Look for common selector patterns
		patterns = [
			r'#[\w-]+',  # ID selectors
			r'\.[\w-]+',  # Class selectors
			r'\[[\w-]+=[^\]]+\]',  # Attribute selectors
			r'button\[type=submit\]',  # Common button selector
			r'input\[type=\w+\]'  # Common input selector
		]

		for pattern in patterns:
			match = re.search(pattern, description)
			if match:
				return match.group(0)

		return ""

	def _extract_form_args(self, description: str) -> Dict[str, str]:
		"""Extract form field arguments from description."""
		import re

		args = {}

		# Look for field-value pairs
		# Pattern: field "value" or field 'value' or field=value
		patterns = [
			r'(\w+)\s*["\']([^"\']+)["\']',
			r'(\w+)\s*=\s*([^\s,]+)'
		]

		for pattern in patterns:
			matches = re.findall(pattern, description)
			for field, value in matches:
				args[field] = value

		# Look for selector and value separately
		if not args:
			selector = self._extract_selector(description)
			if selector:
				# Look for quoted value after selector
				value_match = re.search(r'["\']([^"\']+)["\']', description)
				if value_match:
					args["selector"] = selector
					args["value"] = value_match.group(1)

		return args

	def simulate_plan(
		self,
		initial_obs: Dict[str, Any],
		plan: Dict[str, Any],
		actions: List[Dict[str, Any]],
		state_requirements: Optional[Dict[int, List[str]]] = None,
		validation_points: Optional[List[Dict[str, Any]]] = None
	) -> RolloutResult:
		"""Simulate a sequence of actions from an initial observation."""
		steps: List[RolloutStep] = []
		cumulative_risk = 0.0
		cur_obs = dict(initial_obs)
		terminated_early = False
		termination_reason = None

		# Validate and truncate actions to horizon
		horizon = min(self.config.horizon, len(actions))
		if not actions:
			return RolloutResult(
				steps=[],
				cumulative_risk=0.0,
				avg_risk=0.0,
				success_probability=0.0,
				terminated_early=True,
				termination_reason="No actions provided"
			)

		for idx, action in enumerate(actions[:horizon]):
			# Get world model predictions
			wm_prior = None
			if self.world_model:
				try:
					wm_prior = self.world_model.predict_prior(cur_obs, plan, action)
				except Exception:
					pass

			# Default metrics if WM fails
			metrics = {
				"risk": wm_prior.get("risk", 0.5) if wm_prior else 0.5,
				"success_prob": wm_prior.get("success_prob", 0.5) if wm_prior else 0.5,
				"success_entropy": wm_prior.get("success_entropy", 0.0) if wm_prior else 0.0
			}

			# Predict next state
			next_obs, effects = self.next_state_predictor.predict_next(cur_obs, action)

			# Track state changes if enabled
			state_changes = {}
			if self.config.track_state_changes:
				state_changes = self._detect_state_changes(cur_obs, next_obs)

			# Create step record
			step = RolloutStep(
				idx=idx,
				action=action,
				observation=next_obs,
				effects=effects,
				metrics=metrics,
				state_changes=state_changes
			)
			steps.append(step)

			# Update running metrics
			cumulative_risk += metrics["risk"]
			avg_risk = cumulative_risk / (idx + 1)

			# Check early stopping conditions
			if self.config.enable_early_stopping:
				if avg_risk > self.config.max_risk_threshold:
					terminated_early = True
					termination_reason = f"Average risk {avg_risk:.2f} exceeded threshold {self.config.max_risk_threshold}"
					break
				if metrics["success_prob"] < self.config.min_success_prob:
					terminated_early = True
					termination_reason = f"Success probability {metrics['success_prob']:.2f} below threshold {self.config.min_success_prob}"
					break

			# Update current observation for next iteration
			cur_obs = next_obs

		# Calculate final metrics
		num_steps = len(steps)
		avg_risk = cumulative_risk / num_steps if num_steps > 0 else 0.0
		success_prob = steps[-1].metrics["success_prob"] if steps else 0.0

		return RolloutResult(
			steps=steps,
			cumulative_risk=cumulative_risk,
			avg_risk=avg_risk,
			success_probability=success_prob,
			terminated_early=terminated_early,
			termination_reason=termination_reason
		)

	def _detect_state_changes(
		self,
		prev_obs: Dict[str, Any],
		next_obs: Dict[str, Any]
	) -> Dict[str, Any]:
		"""Detect meaningful changes between observations."""
		changes = {}

		# Track URL changes
		if prev_obs.get("last_url") != next_obs.get("last_url"):
			changes["url_changed"] = {
				"from": prev_obs.get("last_url"),
				"to": next_obs.get("last_url")
			}

		# Track form state changes
		prev_form = prev_obs.get("form_state", {})
		next_form = next_obs.get("form_state", {})
		if prev_form != next_form:
			changes["form_changed"] = {
				"added": {k: v for k, v in next_form.items() if k not in prev_form},
				"removed": {k: v for k, v in prev_form.items() if k not in next_form},
				"modified": {k: {"from": prev_form[k], "to": v} for k, v in next_form.items()
						   if k in prev_form and prev_form[k] != v}
			}

		# Track working directory changes
		if prev_obs.get("cwd") != next_obs.get("cwd"):
			changes["cwd_changed"] = {
				"from": prev_obs.get("cwd"),
				"to": next_obs.get("cwd")
			}

		# Track action type changes
		if prev_obs.get("last_action_kind") != next_obs.get("last_action_kind"):
			changes["action_type_changed"] = {
				"from": prev_obs.get("last_action_kind"),
				"to": next_obs.get("last_action_kind")
			}

		return changes


	cumulative_risk: float
	avg_risk: float
	success_probability: float
	terminated_early: bool
	termination_reason: Optional[str]
	state_validation: Dict[str, Any]  # Validation results for state requirements
	validation_points: List[Dict[str, Any]]  # Results of validation point checks

class RolloutAPI:
	"""API for simulating action plans with configurable horizon.

	Features:
	- k-step lookahead with state prediction
	- Risk and success probability estimation
	- Early stopping on risk/success thresholds
	- State change tracking
	- Detailed metrics per step
	"""

	def __init__(
		self,
		world_model: Optional[WorldModelPrior] = None,
		next_state_predictor: Optional[NextStatePredictor] = None,
		config: Optional[RolloutConfig] = None
	) -> None:
		self.world_model = world_model
		self.next_state_predictor = next_state_predictor or NextStatePredictor()
		self.config = config or RolloutConfig()

	@classmethod
	def from_checkpoint(cls, checkpoint_path: str | Path) -> "RolloutAPI":
		"""Initialize from a saved world model checkpoint."""
		world_model = WorldModelPrior.load(checkpoint_path)
		return cls(world_model=world_model)

	def simulate_sequence(
		self,
		initial_obs: Dict[str, Any],
		plan: Dict[str, Any],
		sequence: Dict[str, Any]
	) -> RolloutResult:
		"""Simulate a sequence with state tracking and validation."""
		# Extract sequence steps
		actions = []
		state_requirements = {}
		validation_points = []

		# Parse sequence


	Features:
	- k-step lookahead with state prediction
	- Risk and success probability estimation
	- Early stopping on risk/success thresholds
	- State change tracking
	- Detailed metrics per step
	"""

	def __init__(
		self,
		world_model: Optional[WorldModelPrior] = None,
		next_state_predictor: Optional[NextStatePredictor] = None,
		config: Optional[RolloutConfig] = None
	) -> None:
		self.world_model = world_model
		self.next_state_predictor = next_state_predictor or NextStatePredictor()
		self.config = config or RolloutConfig()

	@classmethod
	def from_checkpoint(cls, checkpoint_path: str | Path) -> "RolloutAPI":
		"""Initialize from a saved world model checkpoint."""
		world_model = WorldModelPrior.load(checkpoint_path)
		return cls(world_model=world_model)

	def simulate_sequence(
		self,
		initial_obs: Dict[str, Any],
		plan: Dict[str, Any],
		sequence: Dict[str, Any]
	) -> RolloutResult:
		"""Simulate a sequence with state tracking and validation."""
		# Extract sequence steps
		actions = []
		state_requirements = {}
		validation_points = []

		# Parse sequence
		for step in sequence.get("subgoals", []):
			# Convert step to action
			action = self._step_to_action(step)
			actions.append(action)

			# Collect state requirements


	- State change tracking
	- Detailed metrics per step
	"""

	def __init__(
		self,
		world_model: Optional[WorldModelPrior] = None,
		next_state_predictor: Optional[NextStatePredictor] = None,
		config: Optional[RolloutConfig] = None
	) -> None:
		self.world_model = world_model
		self.next_state_predictor = next_state_predictor or NextStatePredictor()
		self.config = config or RolloutConfig()

	@classmethod
	def from_checkpoint(cls, checkpoint_path: str | Path) -> "RolloutAPI":
		"""Initialize from a saved world model checkpoint."""
		world_model = WorldModelPrior.load(checkpoint_path)
		return cls(world_model=world_model)

	def simulate_sequence(
		self,
		initial_obs: Dict[str, Any],
		plan: Dict[str, Any],
		sequence: Dict[str, Any]
	) -> RolloutResult:
		"""Simulate a sequence with state tracking and validation."""
		# Extract sequence steps
		actions = []
		state_requirements = {}
		validation_points = []

		# Parse sequence
		for step in sequence.get("subgoals", []):
			# Convert step to action
			action = self._step_to_action(step)
			actions.append(action)

			# Collect state requirements
			if step.get("requires"):
				state_requirements[len(actions) - 1] = step["requires"]

			# Collect validation points
			if step.get("validation"):
				validation_points.append({
					"step": len(actions) - 1,
					"validation": step["validation"]
				})

		return self.simulate_plan(
			initial_obs,
			plan,
			actions,
			state_requirements=state_requirements,
			validation_points=validation_points
		)

	def _step_to_action(
		self,
		step: Dict[str, Any]
	) -> Dict[str, Any]:
		"""Convert a sequence step to an action."""
		action = {
			"type": self._infer_action_type(step["description"]),
			"args": {}
		}

		# Extract action details from description
		if action["type"] == "navigate":
			action["args"]["url"] = self._extract_url(step["description"])
		elif action["type"] == "fill":
			action["args"].update(self._extract_form_args(step["description"]))
		elif action["type"] == "click":
			action["args"]["selector"] = self._extract_selector(step["description"])

		return action


			# Collect state requirements
			if step.get("requires"):
				state_requirements[len(actions) - 1] = step["requires"]

			# Collect validation points
			if step.get("validation"):
				validation_points.append({
					"step": len(actions) - 1,
					"validation": step["validation"]
				})

		return self.simulate_plan(
			initial_obs,
			plan,
			actions,
			state_requirements=state_requirements,
			validation_points=validation_points
		)

	def _step_to_action(
		self,
		step: Dict[str, Any]
	) -> Dict[str, Any]:
		"""Convert a sequence step to an action."""
		action = {
			"type": self._infer_action_type(step["description"]),
			"args": {}
		}

		# Extract action details from description
		if action["type"] == "navigate":
			action["args"]["url"] = self._extract_url(step["description"])
		elif action["type"] == "fill":
			action["args"].update(self._extract_form_args(step["description"]))
		elif action["type"] == "click":
			action["args"]["selector"] = self._extract_selector(step["description"])

		return action

	def _infer_action_type(self, description: str) -> str:
		"""Infer action type from step description."""
		description = description.lower()

		if any(w in description for w in ["navigate", "go to", "visit", "open"]):
			return "navigate"
		elif any(w in description for w in ["fill", "enter", "type", "input"]):
			return "fill"
		elif any(w in description for w in ["click", "submit", "press", "select"]):
			return "click"
		else:
			return "read"  # Default to read action

	def _extract_url(self, description: str) -> str:
		"""Extract URL from description."""
		import re

		# Look for URL patterns
		url_pattern = r'https?://[^\s<>"\']+|www\.[^\s<>"\']+\.[^\s<>"\']+'


	def _step_to_action(
		self,
		step: Dict[str, Any]
	) -> Dict[str, Any]:
		"""Convert a sequence step to an action."""
		action = {
			"type": self._infer_action_type(step["description"]),
			"args": {}
		}

		# Extract action details from description
		if action["type"] == "navigate":
			action["args"]["url"] = self._extract_url(step["description"])
		elif action["type"] == "fill":
			action["args"].update(self._extract_form_args(step["description"]))
		elif action["type"] == "click":
			action["args"]["selector"] = self._extract_selector(step["description"])

		return action

	def _infer_action_type(self, description: str) -> str:
		"""Infer action type from step description."""
		description = description.lower()

		if any(w in description for w in ["navigate", "go to", "visit", "open"]):
			return "navigate"
		elif any(w in description for w in ["fill", "enter", "type", "input"]):
			return "fill"
		elif any(w in description for w in ["click", "submit", "press", "select"]):
			return "click"
		else:
			return "read"  # Default to read action

	def _extract_url(self, description: str) -> str:
		"""Extract URL from description."""
		import re

		# Look for URL patterns
		url_pattern = r'https?://[^\s<>"\']+|www\.[^\s<>"\']+\.[^\s<>"\']+'
		urls = re.findall(url_pattern, description)

		if urls:
			return urls[0]
		return ""

	def _extract_selector(self, description: str) -> str:
		"""Extract selector from description."""
		import re

		# Look for quoted selectors
		selector_pattern = r'"([^"]+)"|\'([^\']+)\''
		selectors = re.findall(selector_pattern, description)


		elif action["type"] == "fill":
			action["args"].update(self._extract_form_args(step["description"]))
		elif action["type"] == "click":
			action["args"]["selector"] = self._extract_selector(step["description"])

		return action

	def _infer_action_type(self, description: str) -> str:
		"""Infer action type from step description."""
		description = description.lower()

		if any(w in description for w in ["navigate", "go to", "visit", "open"]):
			return "navigate"
		elif any(w in description for w in ["fill", "enter", "type", "input"]):
			return "fill"
		elif any(w in description for w in ["click", "submit", "press", "select"]):
			return "click"
		else:
			return "read"  # Default to read action

	def _extract_url(self, description: str) -> str:
		"""Extract URL from description."""
		import re

		# Look for URL patterns
		url_pattern = r'https?://[^\s<>"\']+|www\.[^\s<>"\']+\.[^\s<>"\']+'
		urls = re.findall(url_pattern, description)

		if urls:
			return urls[0]
		return ""

	def _extract_selector(self, description: str) -> str:
		"""Extract selector from description."""
		import re

		# Look for quoted selectors
		selector_pattern = r'"([^"]+)"|\'([^\']+)\''
		selectors = re.findall(selector_pattern, description)

		if selectors:
			# Take first non-empty group
			selector = next((s for group in selectors for s in group if s), "")
			if selector:
				return selector

		# Look for common selector patterns
		patterns = [
			r'#[\w-]+',  # ID selectors
			r'\.[\w-]+',  # Class selectors
			r'\[[\w-]+=[^\]]+\]',  # Attribute selectors


			return "navigate"
		elif any(w in description for w in ["fill", "enter", "type", "input"]):
			return "fill"
		elif any(w in description for w in ["click", "submit", "press", "select"]):
			return "click"
		else:
			return "read"  # Default to read action

	def _extract_url(self, description: str) -> str:
		"""Extract URL from description."""
		import re

		# Look for URL patterns
		url_pattern = r'https?://[^\s<>"\']+|www\.[^\s<>"\']+\.[^\s<>"\']+'
		urls = re.findall(url_pattern, description)

		if urls:
			return urls[0]
		return ""

	def _extract_selector(self, description: str) -> str:
		"""Extract selector from description."""
		import re

		# Look for quoted selectors
		selector_pattern = r'"([^"]+)"|\'([^\']+)\''
		selectors = re.findall(selector_pattern, description)

		if selectors:
			# Take first non-empty group
			selector = next((s for group in selectors for s in group if s), "")
			if selector:
				return selector

		# Look for common selector patterns
		patterns = [
			r'#[\w-]+',  # ID selectors
			r'\.[\w-]+',  # Class selectors
			r'\[[\w-]+=[^\]]+\]',  # Attribute selectors
			r'button\[type=submit\]',  # Common button selector
			r'input\[type=\w+\]'  # Common input selector
		]

		for pattern in patterns:
			match = re.search(pattern, description)
			if match:
				return match.group(0)

		return ""

	def _extract_form_args(self, description: str) -> Dict[str, str]:
		"""Extract form field arguments from description."""
		import re

		args = {}

		# Look for field-value pairs
		# Pattern: field "value" or field 'value' or field=value
		patterns = [
			r'(\w+)\s*["\']([^"\']+)["\']',
			r'(\w+)\s*=\s*([^\s,]+)'
		]

		for pattern in patterns:
			matches = re.findall(pattern, description)
			for field, value in matches:
				args[field] = value

		# Look for selector and value separately


			selector = next((s for group in selectors for s in group if s), "")
			if selector:
				return selector

		# Look for common selector patterns
		patterns = [
			r'#[\w-]+',  # ID selectors
			r'\.[\w-]+',  # Class selectors
			r'\[[\w-]+=[^\]]+\]',  # Attribute selectors
			r'button\[type=submit\]',  # Common button selector
			r'input\[type=\w+\]'  # Common input selector
		]

		for pattern in patterns:
			match = re.search(pattern, description)
			if match:
				return match.group(0)

		return ""

	def _extract_form_args(self, description: str) -> Dict[str, str]:
		"""Extract form field arguments from description."""
		import re

		args = {}

		# Look for field-value pairs
		# Pattern: field "value" or field 'value' or field=value
		patterns = [
			r'(\w+)\s*["\']([^"\']+)["\']',
			r'(\w+)\s*=\s*([^\s,]+)'
		]

		for pattern in patterns:
			matches = re.findall(pattern, description)
			for field, value in matches:
				args[field] = value

		# Look for selector and value separately
		if not args:
			selector = self._extract_selector(description)
			if selector:
				# Look for quoted value after selector
				value_match = re.search(r'["\']([^"\']+)["\']', description)
				if value_match:
					args["selector"] = selector
					args["value"] = value_match.group(1)

		return args

	def simulate_plan(
		self,
		initial_obs: Dict[str, Any],
		plan: Dict[str, Any],
		actions: List[Dict[str, Any]],
		state_requirements: Optional[Dict[int, List[str]]] = None,
		validation_points: Optional[List[Dict[str, Any]]] = None
	) -> RolloutResult:
		"""Simulate a sequence of actions from an initial observation."""
		steps: List[RolloutStep] = []
		cumulative_risk = 0.0
		cur_obs = dict(initial_obs)
		terminated_early = False
		termination_reason = None

		# Validate and truncate actions to horizon
		horizon = min(self.config.horizon, len(actions))
		if not actions:
			return RolloutResult(


			r'(\w+)\s*=\s*([^\s,]+)'
		]

		for pattern in patterns:
			matches = re.findall(pattern, description)
			for field, value in matches:
				args[field] = value

		# Look for selector and value separately
		if not args:
			selector = self._extract_selector(description)
			if selector:
				# Look for quoted value after selector
				value_match = re.search(r'["\']([^"\']+)["\']', description)
				if value_match:
					args["selector"] = selector
					args["value"] = value_match.group(1)

		return args

	def simulate_plan(
		self,
		initial_obs: Dict[str, Any],
		plan: Dict[str, Any],
		actions: List[Dict[str, Any]],
		state_requirements: Optional[Dict[int, List[str]]] = None,
		validation_points: Optional[List[Dict[str, Any]]] = None
	) -> RolloutResult:
		"""Simulate a sequence of actions from an initial observation."""
		steps: List[RolloutStep] = []
		cumulative_risk = 0.0
		cur_obs = dict(initial_obs)
		terminated_early = False
		termination_reason = None

		# Validate and truncate actions to horizon
		horizon = min(self.config.horizon, len(actions))
		if not actions:
			return RolloutResult(
				steps=[],
				cumulative_risk=0.0,
				avg_risk=0.0,
				success_probability=0.0,
				terminated_early=True,
				termination_reason="No actions provided"
			)

		for idx, action in enumerate(actions[:horizon]):
			# Get world model predictions
			wm_prior = None
			if self.world_model:
				try:
					wm_prior = self.world_model.predict_prior(cur_obs, plan, action)
				except Exception:
					pass

			# Default metrics if WM fails
			metrics = {
				"risk": wm_prior.get("risk", 0.5) if wm_prior else 0.5,
				"success_prob": wm_prior.get("success_prob", 0.5) if wm_prior else 0.5,
				"success_entropy": wm_prior.get("success_entropy", 0.0) if wm_prior else 0.0
			}

			# Predict next state
			next_obs, effects = self.next_state_predictor.predict_next(cur_obs, action)

			# Track state changes if enabled
			state_changes = {}
			if self.config.track_state_changes:
				state_changes = self._detect_state_changes(cur_obs, next_obs)

			# Create step record
			step = RolloutStep(
				idx=idx,
				action=action,
				observation=next_obs,
				effects=effects,
				metrics=metrics,
				state_changes=state_changes
			)
			steps.append(step)

			# Update running metrics
			cumulative_risk += metrics["risk"]
			avg_risk = cumulative_risk / (idx + 1)

			# Check early stopping conditions
			if self.config.enable_early_stopping:
				if avg_risk > self.config.max_risk_threshold:
					terminated_early = True
					termination_reason = f"Average risk {avg_risk:.2f} exceeded threshold {self.config.max_risk_threshold}"
					break
				if metrics["success_prob"] < self.config.min_success_prob:
					terminated_early = True
					termination_reason = f"Success probability {metrics['success_prob']:.2f} below threshold {self.config.min_success_prob}"
					break

			# Update current observation for next iteration
			cur_obs = next_obs

		# Calculate final metrics
		num_steps = len(steps)
		avg_risk = cumulative_risk / num_steps if num_steps > 0 else 0.0
		success_prob = steps[-1].metrics["success_prob"] if steps else 0.0

		return RolloutResult(
			steps=steps,
			cumulative_risk=cumulative_risk,
			avg_risk=avg_risk,
			success_probability=success_prob,
			terminated_early=terminated_early,
			termination_reason=termination_reason
		)

	def _detect_state_changes(
		self,
		prev_obs: Dict[str, Any],
		next_obs: Dict[str, Any]
	) -> Dict[str, Any]:
		"""Detect meaningful changes between observations."""
		changes = {}

		# Track URL changes
		if prev_obs.get("last_url") != next_obs.get("last_url"):
			changes["url_changed"] = {
				"from": prev_obs.get("last_url"),
				"to": next_obs.get("last_url")
			}

		# Track form state changes
		prev_form = prev_obs.get("form_state", {})
		next_form = next_obs.get("form_state", {})
		if prev_form != next_form:


					termination_reason = f"Success probability {metrics['success_prob']:.2f} below threshold {self.config.min_success_prob}"
					break

			# Update current observation for next iteration
			cur_obs = next_obs

		# Calculate final metrics
		num_steps = len(steps)
		avg_risk = cumulative_risk / num_steps if num_steps > 0 else 0.0
		success_prob = steps[-1].metrics["success_prob"] if steps else 0.0

		return RolloutResult(
			steps=steps,
			cumulative_risk=cumulative_risk,
			avg_risk=avg_risk,
			success_probability=success_prob,
			terminated_early=terminated_early,
			termination_reason=termination_reason
		)

	def _detect_state_changes(
		self,
		prev_obs: Dict[str, Any],
		next_obs: Dict[str, Any]
	) -> Dict[str, Any]:
		"""Detect meaningful changes between observations."""
		changes = {}

		# Track URL changes
		if prev_obs.get("last_url") != next_obs.get("last_url"):
			changes["url_changed"] = {
				"from": prev_obs.get("last_url"),
				"to": next_obs.get("last_url")
			}

		# Track form state changes
		prev_form = prev_obs.get("form_state", {})
		next_form = next_obs.get("form_state", {})
		if prev_form != next_form:
			changes["form_changed"] = {
				"added": {k: v for k, v in next_form.items() if k not in prev_form},
				"removed": {k: v for k, v in prev_form.items() if k not in next_form},
				"modified": {k: {"from": prev_form[k], "to": v} for k, v in next_form.items()
						   if k in prev_form and prev_form[k] != v}
			}

		# Track working directory changes
		if prev_obs.get("cwd") != next_obs.get("cwd"):
			changes["cwd_changed"] = {
				"from": prev_obs.get("cwd"),
				"to": next_obs.get("cwd")
			}

		# Track action type changes
		if prev_obs.get("last_action_kind") != next_obs.get("last_action_kind"):
			changes["action_type_changed"] = {
				"from": prev_obs.get("last_action_kind"),
				"to": next_obs.get("last_action_kind")
			}

		return changes


from __future__ import annotations
import logging

from dataclasses import dataclass
from typing import Dict, Any, Optional, List

@dataclass
class RewardConfig:
	"""Configuration for reward shaping."""
	# Base rewards
	success_reward: float = 1.0
	failure_penalty: float = -0.5

	# Risk-based adjustments
	risk_penalty_factor: float = 0.3
	high_risk_threshold: float = 0.7

	# Complexity penalties
	file_count_penalty: float = 0.1
	loc_penalty: float = 0.05
	max_files_threshold: int = 3
	max_loc_threshold: int = 200

	# Near-miss bonuses
	near_miss_bonus: float = 0.2
	partial_success_factor: float = 0.5

	# Test-based rewards
	test_pass_bonus: float = 0.3
	test_improvement_bonus: float = 0.2

	# Intent alignment rewards
	intent_alignment_bonus: float = 0.2
	primary_file_bonus: float = 0.1

class ValidatorRewardShaper:
	"""Shapes rewards for off-policy training using validation signals.

	Features:
	1. Test-based rewards (pass/fail, improvement)
	2. Risk-adjusted rewards
	3. Complexity penalties (files, LOC)
	4. Intent alignment bonuses
	5. Near-miss detection and rewards
	"""

	def __init__(self, config: Optional[RewardConfig] = None):
		self.config = config or RewardConfig()

	def compute_reward(
		self,
		action_result: Dict[str, Any],
		intent: Dict[str, Any],
		world_model_prior: Optional[Dict[str, Any]] = None
	) -> Dict[str, float]:
		"""Compute shaped reward for an action result."""
		rewards = {}
		total_reward = 0.0

		# Base reward from test results
		success = action_result.get("tests_passed", False)
		rewards["success"] = self.config.success_reward if success else self.config.failure_penalty
		total_reward += rewards["success"]

		# Risk-based adjustment
		if world_model_prior:
			risk = float(world_model_prior.get("risk", 0.5))
			if risk > self.config.high_risk_threshold:
				risk_penalty = -self.config.risk_penalty_factor * (risk - self.config.high_risk_threshold)
				rewards["risk"] = risk_penalty
				total_reward += risk_penalty

		# Complexity penalties
		n_files = int(action_result.get("files_changed", 0))
		if n_files > self.config.max_files_threshold:
			file_penalty = -self.config.file_count_penalty * (n_files - self.config.max_files_threshold)
			rewards["files"] = file_penalty
			total_reward += file_penalty

		loc_changed = int(action_result.get("lines_changed", 0))
		if loc_changed > self.config.max_loc_threshold:
			loc_penalty = -self.config.loc_penalty * (loc_changed - self.config.max_loc_threshold)
			rewards["loc"] = loc_penalty
			total_reward += loc_penalty

		# Near-miss detection and bonus
		if self._is_near_miss(action_result):
			rewards["near_miss"] = self.config.near_miss_bonus
			total_reward += self.config.near_miss_bonus

			# Additional bonus for partial success
			if self._is_partial_success(action_result):
				partial_bonus = self.config.near_miss_bonus * self.config.partial_success_factor
				rewards["partial_success"] = partial_bonus
				total_reward += partial_bonus

		# Test improvement bonus
		if self._has_test_improvement(action_result):
			rewards["test_improvement"] = self.config.test_improvement_bonus
			total_reward += self.config.test_improvement_bonus

		# Intent alignment bonus
		alignment_score = self._compute_intent_alignment(action_result, intent)
		if alignment_score > 0:
			intent_bonus = self.config.intent_alignment_bonus * alignment_score
			rewards["intent_alignment"] = intent_bonus
			total_reward += intent_bonus

			# Additional bonus for modifying primary file
			if self._touches_primary_file(action_result, intent):
				rewards["primary_file"] = self.config.primary_file_bonus
				total_reward += self.config.primary_file_bonus

		rewards["total"] = total_reward
		return rewards

	def _is_near_miss(self, result: Dict[str, Any]) -> bool:
		"""Detect if the result is a near-miss (close to success)."""
		if result.get("tests_passed", False):
			return False

		indicators = [
			result.get("partial_success", False),
			result.get("some_tests_passed", False),
			result.get("syntax_valid", True),
			result.get("compiles", True),
			not result.get("runtime_error", False)
		]
		return sum(indicators) >= 3

	def _is_partial_success(self, result: Dict[str, Any]) -> bool:
		"""Check if some success criteria were met."""
		return (
			result.get("partial_success", False) or
			result.get("some_tests_passed", False) or
			(result.get("tests_passed_ratio", 0.0) > 0.5)
		)

	def _has_test_improvement(self, result: Dict[str, Any]) -> bool:
		"""Check if there was improvement in test results."""
		prev_failing = result.get("prev_failing_tests", 0)
		curr_failing = result.get("curr_failing_tests", 0)
		return curr_failing < prev_failing

	def _compute_intent_alignment(self, result: Dict[str, Any], intent: Dict[str, Any]) -> float:
		"""Compute how well the action aligns with intent."""
		alignment_score = 0.0

		# Check if modified files match intent
		if intent.get("target_files"):
			target_files = set(intent["target_files"])
			modified_files = set(result.get("modified_files", []))
			if target_files & modified_files:
				alignment_score += 0.5
				if target_files == modified_files:
					alignment_score += 0.5

		# Check if changes match intent type
		if intent.get("change_type") == result.get("change_type"):
			alignment_score += 0.5

		# Check if size constraints were respected
		if (
			intent.get("max_files", float("inf")) >= result.get("files_changed", 0) and
			intent.get("max_lines", float("inf")) >= result.get("lines_changed", 0)
		):
			alignment_score += 0.5

		return min(1.0, alignment_score)

	def _touches_primary_file(self, result: Dict[str, Any], intent: Dict[str, Any]) -> bool:
		"""Check if the action modified the primary target file."""
		primary_file = intent.get("primary_file")
		if not primary_file:
			return False
		return primary_file in result.get("modified_files", [])


RewardConfig

from __future__ import annotations
import logging

from dataclasses import dataclass
from typing import Dict, Any, Optional, List

@dataclass
class RewardConfig:
	"""Configuration for reward shaping."""
	# Base rewards
	success_reward: float = 1.0
	failure_penalty: float = -0.5

	# Risk-based adjustments
	risk_penalty_factor: float = 0.3
	high_risk_threshold: float = 0.7

	# Complexity penalties
	file_count_penalty: float = 0.1
	loc_penalty: float = 0.05
	max_files_threshold: int = 3
	max_loc_threshold: int = 200

	# Near-miss bonuses
	near_miss_bonus: float = 0.2
	partial_success_factor: float = 0.5

	# Test-based rewards
	test_pass_bonus: float = 0.3
	test_improvement_bonus: float = 0.2

	# Intent alignment rewards
	intent_alignment_bonus: float = 0.2
	primary_file_bonus: float = 0.1

class ValidatorRewardShaper:
	"""Shapes rewards for off-policy training using validation signals.

	Features:
	1. Test-based rewards (pass/fail, improvement)
	2. Risk-adjusted rewards
	3. Complexity penalties (files, LOC)
	4. Intent alignment bonuses
	5. Near-miss detection and rewards
	"""

	def __init__(self, config: Optional[RewardConfig] = None):
		self.config = config or RewardConfig()

	def compute_reward(
		self,
		action_result: Dict[str, Any],
		intent: Dict[str, Any],


ValidatorRewardShaper

	# Complexity penalties
	file_count_penalty: float = 0.1
	loc_penalty: float = 0.05
	max_files_threshold: int = 3
	max_loc_threshold: int = 200

	# Near-miss bonuses
	near_miss_bonus: float = 0.2
	partial_success_factor: float = 0.5

	# Test-based rewards
	test_pass_bonus: float = 0.3
	test_improvement_bonus: float = 0.2

	# Intent alignment rewards
	intent_alignment_bonus: float = 0.2
	primary_file_bonus: float = 0.1

class ValidatorRewardShaper:
	"""Shapes rewards for off-policy training using validation signals.

	Features:
	1. Test-based rewards (pass/fail, improvement)
	2. Risk-adjusted rewards
	3. Complexity penalties (files, LOC)
	4. Intent alignment bonuses
	5. Near-miss detection and rewards
	"""

	def __init__(self, config: Optional[RewardConfig] = None):
		self.config = config or RewardConfig()

	def compute_reward(
		self,
		action_result: Dict[str, Any],
		intent: Dict[str, Any],
		world_model_prior: Optional[Dict[str, Any]] = None
	) -> Dict[str, float]:
		"""Compute shaped reward for an action result."""
		rewards = {}
		total_reward = 0.0

		# Base reward from test results
		success = action_result.get("tests_passed", False)
		rewards["success"] = self.config.success_reward if success else self.config.failure_penalty
		total_reward += rewards["success"]

		# Risk-based adjustment
		if world_model_prior:
			risk = float(world_model_prior.get("risk", 0.5))
			if risk > self.config.high_risk_threshold:
				risk_penalty = -self.config.risk_penalty_factor * (risk - self.config.high_risk_threshold)
				rewards["risk"] = risk_penalty
				total_reward += risk_penalty

		# Complexity penalties
		n_files = int(action_result.get("files_changed", 0))
		if n_files > self.config.max_files_threshold:
			file_penalty = -self.config.file_count_penalty * (n_files - self.config.max_files_threshold)
			rewards["files"] = file_penalty
			total_reward += file_penalty

		loc_changed = int(action_result.get("lines_changed", 0))
		if loc_changed > self.config.max_loc_threshold:
			loc_penalty = -self.config.loc_penalty * (loc_changed - self.config.max_loc_threshold)
			rewards["loc"] = loc_penalty
			total_reward += loc_penalty

		# Near-miss detection and bonus
		if self._is_near_miss(action_result):
			rewards["near_miss"] = self.config.near_miss_bonus
			total_reward += self.config.near_miss_bonus

			# Additional bonus for partial success
			if self._is_partial_success(action_result):
				partial_bonus = self.config.near_miss_bonus * self.config.partial_success_factor
				rewards["partial_success"] = partial_bonus
				total_reward += partial_bonus

		# Test improvement bonus
		if self._has_test_improvement(action_result):
			rewards["test_improvement"] = self.config.test_improvement_bonus
			total_reward += self.config.test_improvement_bonus

		# Intent alignment bonus
		alignment_score = self._compute_intent_alignment(action_result, intent)
		if alignment_score > 0:
			intent_bonus = self.config.intent_alignment_bonus * alignment_score
			rewards["intent_alignment"] = intent_bonus
			total_reward += intent_bonus

			# Additional bonus for modifying primary file
			if self._touches_primary_file(action_result, intent):
				rewards["primary_file"] = self.config.primary_file_bonus
				total_reward += self.config.primary_file_bonus

		rewards["total"] = total_reward
		return rewards

	def _is_near_miss(self, result: Dict[str, Any]) -> bool:
		"""Detect if the result is a near-miss (close to success)."""
		if result.get("tests_passed", False):
			return False

		indicators = [
			result.get("partial_success", False),
			result.get("some_tests_passed", False),
			result.get("syntax_valid", True),
			result.get("compiles", True),
			not result.get("runtime_error", False)
		]
		return sum(indicators) >= 3

	def _is_partial_success(self, result: Dict[str, Any]) -> bool:
		"""Check if some success criteria were met."""
		return (
			result.get("partial_success", False) or
			result.get("some_tests_passed", False) or
			(result.get("tests_passed_ratio", 0.0) > 0.5)
		)

	def _has_test_improvement(self, result: Dict[str, Any]) -> bool:
		"""Check if there was improvement in test results."""
		prev_failing = result.get("prev_failing_tests", 0)
		curr_failing = result.get("curr_failing_tests", 0)
		return curr_failing < prev_failing

	def _compute_intent_alignment(self, result: Dict[str, Any], intent: Dict[str, Any]) -> float:
		"""Compute how well the action aligns with intent."""
		alignment_score = 0.0

		# Check if modified files match intent
		if intent.get("target_files"):
			target_files = set(intent["target_files"])
			modified_files = set(result.get("modified_files", []))
			if target_files & modified_files:
				alignment_score += 0.5
				if target_files == modified_files:
					alignment_score += 0.5

		# Check if changes match intent type
		if intent.get("change_type") == result.get("change_type"):
			alignment_score += 0.5

		# Check if size constraints were respected
		if (
			intent.get("max_files", float("inf")) >= result.get("files_changed", 0) and
			intent.get("max_lines", float("inf")) >= result.get("lines_changed", 0)
		):
			alignment_score += 0.5

		return min(1.0, alignment_score)

	def _touches_primary_file(self, result: Dict[str, Any], intent: Dict[str, Any]) -> bool:
		"""Check if the action modified the primary target file."""
		primary_file = intent.get("primary_file")
		if not primary_file:
			return False
		return primary_file in result.get("modified_files", [])


	# Test-based rewards
	test_pass_bonus: float = 0.3
	test_improvement_bonus: float = 0.2

	# Intent alignment rewards
	intent_alignment_bonus: float = 0.2
	primary_file_bonus: float = 0.1

class ValidatorRewardShaper:
	"""Shapes rewards for off-policy training using validation signals.

	Features:
	1. Test-based rewards (pass/fail, improvement)
	2. Risk-adjusted rewards
	3. Complexity penalties (files, LOC)
	4. Intent alignment bonuses
	5. Near-miss detection and rewards
	"""

	def __init__(self, config: Optional[RewardConfig] = None):
		self.config = config or RewardConfig()

	def compute_reward(
		self,
		action_result: Dict[str, Any],
		intent: Dict[str, Any],
		world_model_prior: Optional[Dict[str, Any]] = None
	) -> Dict[str, float]:
		"""Compute shaped reward for an action result."""
		rewards = {}
		total_reward = 0.0

		# Base reward from test results
		success = action_result.get("tests_passed", False)
		rewards["success"] = self.config.success_reward if success else self.config.failure_penalty
		total_reward += rewards["success"]

		# Risk-based adjustment
		if world_model_prior:
			risk = float(world_model_prior.get("risk", 0.5))
			if risk > self.config.high_risk_threshold:


	# Intent alignment rewards
	intent_alignment_bonus: float = 0.2
	primary_file_bonus: float = 0.1

class ValidatorRewardShaper:
	"""Shapes rewards for off-policy training using validation signals.

	Features:
	1. Test-based rewards (pass/fail, improvement)
	2. Risk-adjusted rewards
	3. Complexity penalties (files, LOC)
	4. Intent alignment bonuses
	5. Near-miss detection and rewards
	"""

	def __init__(self, config: Optional[RewardConfig] = None):
		self.config = config or RewardConfig()

	def compute_reward(
		self,
		action_result: Dict[str, Any],
		intent: Dict[str, Any],
		world_model_prior: Optional[Dict[str, Any]] = None
	) -> Dict[str, float]:
		"""Compute shaped reward for an action result."""
		rewards = {}
		total_reward = 0.0

		# Base reward from test results
		success = action_result.get("tests_passed", False)
		rewards["success"] = self.config.success_reward if success else self.config.failure_penalty
		total_reward += rewards["success"]

		# Risk-based adjustment
		if world_model_prior:
			risk = float(world_model_prior.get("risk", 0.5))
			if risk > self.config.high_risk_threshold:
				risk_penalty = -self.config.risk_penalty_factor * (risk - self.config.high_risk_threshold)
				rewards["risk"] = risk_penalty
				total_reward += risk_penalty

		# Complexity penalties
		n_files = int(action_result.get("files_changed", 0))
		if n_files > self.config.max_files_threshold:
			file_penalty = -self.config.file_count_penalty * (n_files - self.config.max_files_threshold)
			rewards["files"] = file_penalty
			total_reward += file_penalty

		loc_changed = int(action_result.get("lines_changed", 0))
		if loc_changed > self.config.max_loc_threshold:
			loc_penalty = -self.config.loc_penalty * (loc_changed - self.config.max_loc_threshold)
			rewards["loc"] = loc_penalty
			total_reward += loc_penalty

		# Near-miss detection and bonus
		if self._is_near_miss(action_result):
			rewards["near_miss"] = self.config.near_miss_bonus
			total_reward += self.config.near_miss_bonus

			# Additional bonus for partial success
			if self._is_partial_success(action_result):
				partial_bonus = self.config.near_miss_bonus * self.config.partial_success_factor
				rewards["partial_success"] = partial_bonus
				total_reward += partial_bonus

		# Test improvement bonus
		if self._has_test_improvement(action_result):
			rewards["test_improvement"] = self.config.test_improvement_bonus
			total_reward += self.config.test_improvement_bonus

		# Intent alignment bonus
		alignment_score = self._compute_intent_alignment(action_result, intent)
		if alignment_score > 0:
			intent_bonus = self.config.intent_alignment_bonus * alignment_score
			rewards["intent_alignment"] = intent_bonus
			total_reward += intent_bonus

			# Additional bonus for modifying primary file
			if self._touches_primary_file(action_result, intent):
				rewards["primary_file"] = self.config.primary_file_bonus
				total_reward += self.config.primary_file_bonus

		rewards["total"] = total_reward
		return rewards

	def _is_near_miss(self, result: Dict[str, Any]) -> bool:
		"""Detect if the result is a near-miss (close to success)."""
		if result.get("tests_passed", False):
			return False

		indicators = [
			result.get("partial_success", False),
			result.get("some_tests_passed", False),
			result.get("syntax_valid", True),
			result.get("compiles", True),
			not result.get("runtime_error", False)
		]
		return sum(indicators) >= 3

	def _is_partial_success(self, result: Dict[str, Any]) -> bool:
		"""Check if some success criteria were met."""
		return (
			result.get("partial_success", False) or
			result.get("some_tests_passed", False) or


		# Test improvement bonus
		if self._has_test_improvement(action_result):
			rewards["test_improvement"] = self.config.test_improvement_bonus
			total_reward += self.config.test_improvement_bonus

		# Intent alignment bonus
		alignment_score = self._compute_intent_alignment(action_result, intent)
		if alignment_score > 0:
			intent_bonus = self.config.intent_alignment_bonus * alignment_score
			rewards["intent_alignment"] = intent_bonus
			total_reward += intent_bonus

			# Additional bonus for modifying primary file
			if self._touches_primary_file(action_result, intent):
				rewards["primary_file"] = self.config.primary_file_bonus
				total_reward += self.config.primary_file_bonus

		rewards["total"] = total_reward
		return rewards

	def _is_near_miss(self, result: Dict[str, Any]) -> bool:
		"""Detect if the result is a near-miss (close to success)."""
		if result.get("tests_passed", False):
			return False

		indicators = [
			result.get("partial_success", False),
			result.get("some_tests_passed", False),
			result.get("syntax_valid", True),
			result.get("compiles", True),
			not result.get("runtime_error", False)
		]
		return sum(indicators) >= 3

	def _is_partial_success(self, result: Dict[str, Any]) -> bool:
		"""Check if some success criteria were met."""
		return (
			result.get("partial_success", False) or
			result.get("some_tests_passed", False) or
			(result.get("tests_passed_ratio", 0.0) > 0.5)
		)

	def _has_test_improvement(self, result: Dict[str, Any]) -> bool:
		"""Check if there was improvement in test results."""
		prev_failing = result.get("prev_failing_tests", 0)
		curr_failing = result.get("curr_failing_tests", 0)
		return curr_failing < prev_failing

	def _compute_intent_alignment(self, result: Dict[str, Any], intent: Dict[str, Any]) -> float:
		"""Compute how well the action aligns with intent."""
		alignment_score = 0.0

		# Check if modified files match intent


				rewards["primary_file"] = self.config.primary_file_bonus
				total_reward += self.config.primary_file_bonus

		rewards["total"] = total_reward
		return rewards

	def _is_near_miss(self, result: Dict[str, Any]) -> bool:
		"""Detect if the result is a near-miss (close to success)."""
		if result.get("tests_passed", False):
			return False

		indicators = [
			result.get("partial_success", False),
			result.get("some_tests_passed", False),
			result.get("syntax_valid", True),
			result.get("compiles", True),
			not result.get("runtime_error", False)
		]
		return sum(indicators) >= 3

	def _is_partial_success(self, result: Dict[str, Any]) -> bool:
		"""Check if some success criteria were met."""
		return (
			result.get("partial_success", False) or
			result.get("some_tests_passed", False) or
			(result.get("tests_passed_ratio", 0.0) > 0.5)
		)

	def _has_test_improvement(self, result: Dict[str, Any]) -> bool:
		"""Check if there was improvement in test results."""
		prev_failing = result.get("prev_failing_tests", 0)
		curr_failing = result.get("curr_failing_tests", 0)
		return curr_failing < prev_failing

	def _compute_intent_alignment(self, result: Dict[str, Any], intent: Dict[str, Any]) -> float:
		"""Compute how well the action aligns with intent."""
		alignment_score = 0.0

		# Check if modified files match intent
		if intent.get("target_files"):
			target_files = set(intent["target_files"])
			modified_files = set(result.get("modified_files", []))
			if target_files & modified_files:
				alignment_score += 0.5
				if target_files == modified_files:
					alignment_score += 0.5


		if result.get("tests_passed", False):
			return False

		indicators = [
			result.get("partial_success", False),
			result.get("some_tests_passed", False),
			result.get("syntax_valid", True),
			result.get("compiles", True),
			not result.get("runtime_error", False)
		]
		return sum(indicators) >= 3

	def _is_partial_success(self, result: Dict[str, Any]) -> bool:
		"""Check if some success criteria were met."""
		return (
			result.get("partial_success", False) or
			result.get("some_tests_passed", False) or
			(result.get("tests_passed_ratio", 0.0) > 0.5)
		)

	def _has_test_improvement(self, result: Dict[str, Any]) -> bool:
		"""Check if there was improvement in test results."""
		prev_failing = result.get("prev_failing_tests", 0)
		curr_failing = result.get("curr_failing_tests", 0)
		return curr_failing < prev_failing

	def _compute_intent_alignment(self, result: Dict[str, Any], intent: Dict[str, Any]) -> float:
		"""Compute how well the action aligns with intent."""
		alignment_score = 0.0

		# Check if modified files match intent
		if intent.get("target_files"):
			target_files = set(intent["target_files"])
			modified_files = set(result.get("modified_files", []))
			if target_files & modified_files:
				alignment_score += 0.5
				if target_files == modified_files:
					alignment_score += 0.5

		# Check if changes match intent type
		if intent.get("change_type") == result.get("change_type"):
			alignment_score += 0.5

		# Check if size constraints were respected
		if (


			result.get("syntax_valid", True),
			result.get("compiles", True),
			not result.get("runtime_error", False)
		]
		return sum(indicators) >= 3

	def _is_partial_success(self, result: Dict[str, Any]) -> bool:
		"""Check if some success criteria were met."""
		return (
			result.get("partial_success", False) or
			result.get("some_tests_passed", False) or
			(result.get("tests_passed_ratio", 0.0) > 0.5)
		)

	def _has_test_improvement(self, result: Dict[str, Any]) -> bool:
		"""Check if there was improvement in test results."""
		prev_failing = result.get("prev_failing_tests", 0)
		curr_failing = result.get("curr_failing_tests", 0)
		return curr_failing < prev_failing

	def _compute_intent_alignment(self, result: Dict[str, Any], intent: Dict[str, Any]) -> float:
		"""Compute how well the action aligns with intent."""
		alignment_score = 0.0

		# Check if modified files match intent
		if intent.get("target_files"):
			target_files = set(intent["target_files"])
			modified_files = set(result.get("modified_files", []))
			if target_files & modified_files:
				alignment_score += 0.5
				if target_files == modified_files:
					alignment_score += 0.5

		# Check if changes match intent type
		if intent.get("change_type") == result.get("change_type"):
			alignment_score += 0.5

		# Check if size constraints were respected
		if (
			intent.get("max_files", float("inf")) >= result.get("files_changed", 0) and
			intent.get("max_lines", float("inf")) >= result.get("lines_changed", 0)
		):
			alignment_score += 0.5

		return min(1.0, alignment_score)

	def _touches_primary_file(self, result: Dict[str, Any], intent: Dict[str, Any]) -> bool:
		"""Check if the action modified the primary target file."""
		primary_file = intent.get("primary_file")
		if not primary_file:
			return False
		return primary_file in result.get("modified_files", [])


			target_files = set(intent["target_files"])
			modified_files = set(result.get("modified_files", []))
			if target_files & modified_files:
				alignment_score += 0.5
				if target_files == modified_files:
					alignment_score += 0.5

		# Check if changes match intent type
		if intent.get("change_type") == result.get("change_type"):
			alignment_score += 0.5

		# Check if size constraints were respected
		if (
			intent.get("max_files", float("inf")) >= result.get("files_changed", 0) and
			intent.get("max_lines", float("inf")) >= result.get("lines_changed", 0)
		):
			alignment_score += 0.5

		return min(1.0, alignment_score)

	def _touches_primary_file(self, result: Dict[str, Any], intent: Dict[str, Any]) -> bool:
		"""Check if the action modified the primary target file."""
		primary_file = intent.get("primary_file")
		if not primary_file:
			return False
		return primary_file in result.get("modified_files", [])


from __future__ import annotations
import logging

from pathlib import Path
from typing import Dict, Optional

def pick_from_bank(root: Path, bank_name: Optional[str]) -> Dict[str, str]:
	"""Return an adapter mapping for roles (e.g., {"planner": dir, "verifier": dir}) from a named bank.

	If bank_name is None or missing, returns an empty mapping.
	"""
	if not bank_name:
		return {}
	try:
		from agi_dw.core.memory.service import get_adapter_bank  # type: ignore
		return get_adapter_bank(root, bank_name)
	except Exception:
		return {}

def pick_skill_adapters(root: Path, domain: str, signature: Dict[str, str]) -> Dict[str, str]:
	"""Attempt to find a matching skill and return its adapters mapping if present.

	Signature is domain-specific; this is a best-effort helper and may return empty.
	"""
	try:
		from agi_dw.core.memory.service import pick_skill_adapters  # type: ignore
		return pick_skill_adapters(root, domain, signature)
	except Exception:
		pass
	return {}


import logging

from pathlib import Path
from typing import Dict, Optional

def pick_from_bank(root: Path, bank_name: Optional[str]) -> Dict[str, str]:
	"""Return an adapter mapping for roles (e.g., {"planner": dir, "verifier": dir}) from a named bank.

	If bank_name is None or missing, returns an empty mapping.
	"""
	if not bank_name:
		return {}
	try:
		from agi_dw.core.memory.service import get_adapter_bank  # type: ignore
		return get_adapter_bank(root, bank_name)
	except Exception:
		return {}

def pick_skill_adapters(root: Path, domain: str, signature: Dict[str, str]) -> Dict[str, str]:
	"""Attempt to find a matching skill and return its adapters mapping if present.

	Signature is domain-specific; this is a best-effort helper and may return empty.
	"""
	try:
		from agi_dw.core.memory.service import pick_skill_adapters  # type: ignore
		return pick_skill_adapters(root, domain, signature)
	except Exception:
		pass
	return {}


import logging
from typing import Any, Dict, List, Optional
import os

_MODEL_CACHE: Dict[str, "HFClient"] = {}

class HFClient:
	def __init__(self, model_id: str = "meta-llama/Llama-3.1-8B-Instruct", device_map: Optional[str] = "auto", torch_dtype: Optional[str] = None) -> None:
		try:
			import torch  # type: ignore
			from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore
		except Exception as e:
			raise RuntimeError("transformers/torch not installed: pip install -r requirements.txt") from e
		self.model_id = model_id
		self.tokenizer = AutoTokenizer.from_pretrained(model_id)
		# Allow env overrides and torchrun rank binding
		env_device_map = os.environ.get("HF_DEVICE_MAP") or device_map
		try:
			# If running under torchrun and device_map wasn't explicitly set, bind to this rank's GPU
			if (env_device_map is None or str(env_device_map).strip().lower() == "auto") and os.environ.get("LOCAL_RANK") is not None:
				lr = int(os.environ.get("LOCAL_RANK", "0") or 0)
				env_device_map = f"cuda:{lr}" if lr >= 0 else env_device_map
		except Exception:
			pass
		env_dtype = os.environ.get("HF_TORCH_DTYPE") or torch_dtype
		dtype = None
		if env_dtype == "float16":
			dtype = torch.float16
		elif env_dtype == "bfloat16":
			dtype = torch.bfloat16
		elif env_dtype is None and torch.cuda.is_available():
			# Prefer BF16 if supported, else FP16 for faster GPU inference
			try:
				if hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported():
					dtype = torch.bfloat16
				else:
					dtype = torch.float16
			except Exception:
				dtype = torch.float16
		try:
			self.model = AutoModelForCausalLM.from_pretrained(
				model_id,
				device_map=env_device_map,
				torch_dtype=dtype,
			)
		except Exception as e:
			# Graceful low-memory fallback: load on CPU if CUDA OOM or device map fails
			try:
				self.model = AutoModelForCausalLM.from_pretrained(
					model_id,
					device_map="cpu",
				)
			except Exception as e2:
				raise e
		# Save dtype for autocast
		self._dtype = dtype
		# Prefer faster attention kernels when available; disable sliding window warnings
		try:
			# Explicitly set attention implementation if configurable
			if hasattr(self.model, "config"):
				cfg = self.model.config
				# Disable sliding window settings that sdpa does not support
				if hasattr(cfg, "sliding_window"):
					try:
						setattr(cfg, "sliding_window", None)
					except Exception:
						pass
				if hasattr(cfg, "use_sliding_window"):
					try:
						setattr(cfg, "use_sliding_window", False)
					except Exception:
						pass
				if hasattr(cfg, "sliding_window_size"):
					try:
						setattr(cfg, "sliding_window_size", 0)
					except Exception:
						pass
				if hasattr(cfg, "use_cache"):
					setattr(cfg, "use_cache", True)
				# Prefer flash-attn2 if available and requested; else SDPA
				attn_impl = "sdpa"
				if os.environ.get("HF_FLASH_ATTN", "0") in ("1", "true", "True"):
					try:
						import flash_attn  # type: ignore  # noqa: F401
						attn_impl = "flash_attention_2"
					except Exception:
						attn_impl = "sdpa"
				if hasattr(cfg, "attn_implementation"):
					setattr(cfg, "attn_implementation", attn_impl)
				if hasattr(self.model, "generation_config"):
					if hasattr(self.model.generation_config, "use_cache"):
						self.model.generation_config.use_cache = True  # type: ignore[attr-defined]
					if hasattr(self.model.generation_config, "sliding_window"):
						try:
							self.model.generation_config.sliding_window = None  # type: ignore[attr-defined]
						except Exception:
							pass
		except Exception:
			pass
		# Ensure model is on GPU if available
		try:
			if torch.cuda.is_available():
				# Respect torchrun LOCAL_RANK when present
				try:
					if os.environ.get("LOCAL_RANK") is not None:
						lr = int(os.environ.get("LOCAL_RANK", "0") or 0)
						if lr >= 0 and hasattr(torch.cuda, "set_device"):
							torch.cuda.set_device(lr)
				except Exception:
					pass
				current = str(self.model.device)
				if "cuda" not in current:
					# Attempt to move to CUDA; if it fails due to OOM, continue on CPU
					try:
						self.model.to("cuda")
					except Exception:
						pass
			# Enable TF32 on Ampere+ for faster matmul where acceptable
			try:
				torch.backends.cuda.matmul.allow_tf32 = True  # type: ignore[attr-defined]
				torch.backends.cudnn.allow_tf32 = True  # type: ignore[attr-defined]
				if hasattr(torch, "set_float32_matmul_precision"):
					torch.set_float32_matmul_precision("high")  # type: ignore[attr-defined]
			except Exception:
				pass
		except Exception:
			pass

	@classmethod
	def get_cached(cls, model_id: str) -> "HFClient":
		client = _MODEL_CACHE.get(model_id)
		if client is None:
			client = HFClient(model_id=model_id)
			_MODEL_CACHE[model_id] = client
		return client

	def attach_adapter(self, adapter_dir: Optional[str]) -> None:
		"""Optionally load a PEFT adapter onto the base model if available."""
		if not adapter_dir:
			return
		try:
			from peft import PeftModel  # type: ignore
			self.model = PeftModel.from_pretrained(self.model, adapter_dir)
		except Exception:
			# Silently continue if PEFT not installed or adapter invalid
			pass

	def generate(self, prompt: str, max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None, stop: Optional[List[str]] = None, grammar: Optional[str] = None) -> str:
		import torch  # type: ignore
		inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
		do_sample = temperature > 0.0
		# Try grammar-constrained decoding via Outlines if requested
		if grammar:
			try:
				import outlines  # type: ignore
				from outlines.models.transformers import Transformers  # type: ignore
				# Map known grammar shorthands to a conservative regex that discourages
				# docstrings and line comments in Python function bodies.
				def _python_body_minimal_regex() -> str:
					# Matches up to max_new_tokens lines consisting of optional indentation
					# and any non-newline characters, but forbids lines starting with
					# triple quotes or '#'. Negative lookaheads apply per line.
					# Note: This is a best-effort regex; Outlines enforces it token by token.
					return r"(?:[ \t]*(?!#)(?!\"\"\"|''')[^\n]*\n?){1," + str(max(1, int(max_new_tokens))) + r"}"
				pattern = None
				if grammar == "python_function_body_minimal":
					pattern = _python_body_minimal_regex()
				else:
					# Allow passing a raw regex as grammar string
					pattern = str(grammar)
				omodel = Transformers(self.model, self.tokenizer)
				gen = outlines.generate.regex(omodel, pattern)
				# Outlines handles sampling internally; pass temperature when sampling
				kwargs: Dict[str, Any] = {}
				if do_sample:
					kwargs["temperature"] = float(temperature)
					if top_p is not None:
						kwargs["top_p"] = float(top_p)
					if top_k is not None:
						kwargs["top_k"] = int(top_k)
				# Execute constrained generation. If anything fails, fall back below.
				text = str(gen(prompt, max_tokens=int(max_new_tokens), **kwargs))
				# Heuristic to drop echoed prompt
				if text.startswith(prompt):
					text = text[len(prompt):]
				text = text.lstrip("\n").rstrip()
				# Optional post-hoc stop handling for safety
				try:
					if stop:
						cut = len(text)
						for s in stop:
							if not s:
								continue
							idx = text.find(s)
							if idx != -1:
								cut = min(cut, idx)
						text = text[:cut]
				except Exception:
					pass
				return text
			except Exception:
				# Fall back to vanilla generation below
				pass
		with torch.inference_mode():
			gen_kwargs: Dict[str, Any] = {
				"max_new_tokens": max_new_tokens,
				"do_sample": do_sample,
				"pad_token_id": self.tokenizer.eos_token_id,
			}
			if do_sample:
				gen_kwargs["temperature"] = temperature
				if top_p is not None:
					gen_kwargs["top_p"] = float(top_p)
				if top_k is not None:
					gen_kwargs["top_k"] = int(top_k)
			# Remove sampler-only params to avoid warnings
			if not do_sample:
				if hasattr(self.model.generation_config, "top_p"):
					try:
						self.model.generation_config.top_p = None  # type: ignore[attr-defined]
					except Exception:
						pass
				if hasattr(self.model.generation_config, "top_k"):
					try:
						self.model.generation_config.top_k = None  # type: ignore[attr-defined]
					except Exception:
						pass
				if hasattr(self.model.generation_config, "temperature"):
					try:
						self.model.generation_config.temperature = None  # type: ignore[attr-defined]
					except Exception:
						pass
			# Autocast to configured dtype
			try:
				if torch.cuda.is_available() and self._dtype in (torch.float16, torch.bfloat16):
					with torch.amp.autocast("cuda", dtype=self._dtype):  # type: ignore[attr-defined]
						output = self.model.generate(**inputs, **gen_kwargs)
				else:
					output = self.model.generate(**inputs, **gen_kwargs)
			except Exception:
				output = self.model.generate(**inputs, **gen_kwargs)
		text = self.tokenizer.decode(output[0], skip_special_tokens=True)
		# Heuristic to return only the completion while preserving leading indentation
		if text.startswith(prompt):
			text = text[len(prompt) :]
		else:
			text = text
		# Optional post-hoc stop truncation by strings
		try:
			if stop:
				cut = len(text)
				for s in stop:
					if not s:
						continue
					idx = text.find(s)
					if idx != -1:
						cut = min(cut, idx)
				text = text[:cut]
		except Exception:
			pass
		return text.lstrip("\n").rstrip()

	def generate_batch(self, prompts: List[str], max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[str]:
		import torch  # type: ignore
		do_sample = temperature > 0.0
		enc = self.tokenizer(prompts, return_tensors="pt", padding=True, truncation=True)
		enc = {k: v.to(self.model.device) for k, v in enc.items()}
		gen_kwargs: Dict[str, Any] = {
			"max_new_tokens": max_new_tokens,
			"do_sample": do_sample,
			"pad_token_id": self.tokenizer.eos_token_id,
		}
		if do_sample:
			gen_kwargs["temperature"] = temperature
			if top_p is not None:
				gen_kwargs["top_p"] = float(top_p)
			if top_k is not None:
				gen_kwargs["top_k"] = int(top_k)
		with torch.inference_mode():
			if torch.cuda.is_available() and self._dtype in (torch.float16, torch.bfloat16):
				with torch.amp.autocast("cuda", dtype=self._dtype):  # type: ignore[attr-defined]
					out = self.model.generate(**enc, **gen_kwargs)
			else:
				out = self.model.generate(**enc, **gen_kwargs)
		texts = self.tokenizer.batch_decode(out, skip_special_tokens=True)
		# Trim echoes where possible, preserving leading indentation
		cleaned: List[str] = []
		for p, t in zip(prompts, texts):
			if t.startswith(p):
				cleaned.append(t[len(p):].lstrip("\n").rstrip())
			else:
				cleaned.append(t.lstrip("\n").rstrip())
		return cleaned

	def chat(self, messages: List[Dict[str, str]], max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None, stop: Optional[List[str]] = None, grammar: Optional[str] = None) -> str:
		# Try chat template if available
		prompt = None
		if hasattr(self.tokenizer, "apply_chat_template"):
			try:
				prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
			except Exception:
				prompt = None
		if prompt is None:
			# Fallback to simple concat
			parts: List[str] = []
			for m in messages:
				role = m.get("role", "user")
				content = m.get("content", "")
				parts.append(f"{role}: {content}")
			prompt = "\n".join(parts) + "\nassistant:"
		# If a grammar is requested, route through generate() with grammar to enable
		# constrained decoding for chat-style prompts as well.
		return self.generate(prompt, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, stop=stop, grammar=grammar)