						except Exception:
							continue
			except Exception:
				# bs4 not available or parse failed; leave text empty
				text = ""
			res = {"url": url, "selector": css_selector, "text": text, "error_type": None if text else "not_found", "attempts": 0}
			try:
				if use_cache:
					cache_set("dom_http", [url, css_selector], text)
			except Exception:
				pass
			try:
				meter_cost("dom_http", 1.0)
			except Exception:
				pass
			return res
		except Exception:
			return {"url": url, "selector": css_selector, "text": "", "error_type": "http_error", "attempts": 0}

def fetch_text(url: str, css_selector: str = "body", wait_visible_secs: int | None = None) -> Dict[str, Any]:
	# Egress allowlist enforcement
	if not _is_host_allowed(url):
		return {"url": url, "selector": css_selector, "text": "", "blocked": True}
	with trace_span("dom_fetch", {"url": url, "selector": css_selector}):
		try:
			import undetected_chromedriver as uc  # type: ignore
			from selenium.common.exceptions import SessionNotCreatedException  # type: ignore
			exists = True
		except Exception:
			exists = False
		if not exists:
			# Graceful fallback to HTTP fetch when ChromeDriver/UC is unavailable
			return _http_fetch(url, css_selector)
		from selenium.webdriver.common.by import By  # type: ignore
		from selenium.webdriver.support.ui import WebDriverWait  # type: ignore
		from selenium.webdriver.support import expected_conditions as EC  # type: ignore
		from selenium.common.exceptions import StaleElementReferenceException, TimeoutException  # type: ignore
		import re

		def _make_opts():
			# Create a FRESH options object each time; uc forbids reusing the same instance
			o = uc.ChromeOptions()
			o.add_argument("--headless=new")
			o.add_argument("--no-sandbox")
			o.add_argument("--disable-dev-shm-usage")
			# Reduce blocking on full page load
			try:
				# 'eager' returns after DOMContentLoaded
				o.page_load_strategy = 'eager'  # type: ignore[attr-defined]
			except Exception:
				pass
			return o
		browser = None
		try:
			# Attempt multiple cold-driver attempts to mitigate flakiness
			retries = max(1, _get_retries())
			for attempt in range(retries):
				try:
					try:
						# Fresh options per attempt to avoid reuse error
						opts = _make_opts()
						_apply_proxy_to_options(opts)
						chrome_binary = _get_chrome_binary()
						if chrome_binary:
							try:
								opts.binary_location = chrome_binary  # type: ignore[attr-defined]
							except Exception:
								pass
						cd_path = _get_chromedriver_path()
						if cd_path:
							browser = uc.Chrome(driver_executable_path=cd_path, options=opts)
						else:
							browser = uc.Chrome(options=opts)
						# Apply viewport and network optimizations
						try:
							w, h = _get_viewport_size()
							browser.set_window_size(w, h)
							try:
								browser.execute_cdp_cmd("Network.enable", {})
								# Set Accept-Language early
								try:
									browser.execute_cdp_cmd("Network.setExtraHTTPHeaders", {"headers": {"Accept-Language": _get_accept_language()}})
								except Exception:
									pass
								if _disable_images():
									try:
										browser.execute_cdp_cmd("Network.setBlockedURLs", {"urls": ["*.png", "*.jpg", "*.jpeg", "*.gif", "*.webp", "*.svg", "*.ico", "*.woff", "*.woff2", "*.ttf"]})
									except Exception:
										pass
							except Exception:
								pass
						except Exception:
							pass
					except SessionNotCreatedException as e:
						m = re.search(r"Current browser version is (\d+)", str(e))
						version_main = _get_chrome_version_main()
						if not version_main and not m:
							raise
						if not version_main:
							version_main = int(m.group(1))
						# Create a FRESH options object; uc forbids reusing the same instance
						opts2 = _make_opts()
						cd_path = _get_chromedriver_path()
						if cd_path:
							browser = uc.Chrome(driver_executable_path=cd_path, options=opts2, version_main=version_main)
						else:
							browser = uc.Chrome(options=opts2, version_main=version_main)
					# Bound page load time; prevent indefinite blocks inside driver.get
					try:
						browser.set_page_load_timeout(_get_timeout_secs())
					except Exception:
						pass
					# Install JS guards to block POST by default
					_install_net_guards(browser)
					try:
						browser.get(url)
					except TimeoutException:
						# Stop loading and proceed; we'll attempt to locate the element anyway
						try:
							browser.execute_script("window.stop();")
						except Exception:
							pass
					# Light readiness, set UA, and scroll to prompt lazy content
					try:
						# Set navigator.userAgent for some sites that check it (best-effort)
						try:
							browser.execute_cdp_cmd("Network.setUserAgentOverride", {"userAgent": _get_user_agent()})
						except Exception:
							pass
						# Disable animations/transitions for stability
						try:
							browser.execute_script("var s=document.createElement('style');s.innerHTML='*{animation:none!important;transition:none!important;scroll-behavior:auto!important}';document.head.appendChild(s);")
						except Exception:
							pass
						# Ensure target element is in view before reading
						browser.execute_script("return document.readyState")
						browser.execute_script("window.scrollTo(0, 200)")
						browser.execute_script("window.scrollTo(0, 0)")
					except Exception:
						pass
					# Optional extra wait for dynamic content stability
					try:
						extra = _get_extra_wait_secs()
						if extra > 0:
							import time
							time.sleep(min(extra, 3.0))
					except Exception:
						pass
					# Attempt to dismiss common cookie banners to unblock content
					try:
						from selenium.webdriver.common.by import By  # type: ignore
						from selenium.webdriver.support.ui import WebDriverWait  # type: ignore
						from selenium.webdriver.support import expected_conditions as EC  # type: ignore
						cookie_selectors = [
							"button#onetrust-accept-btn-handler",
							"button[aria-label*='Accept']",
							"button[aria-label*='accept']",
							"#onetrust-accept-btn-handler",
							"button.cookie-accept",
						]
						for sel in cookie_selectors:
							try:
								WebDriverWait(browser, 2).until(EC.element_to_be_clickable((By.CSS_SELECTOR, sel)))
								el = browser.find_element(By.CSS_SELECTOR, sel)
								el.click()
								break
							except Exception:
								continue
						# Fallback: try a generic XPath match for buttons containing 'Accept'
						try:
							elx = browser.find_element(By.XPATH, "//button[contains(translate(normalize-space(.), 'ABCDEFGHIJKLMNOPQRSTUVWXYZ', 'abcdefghijklmnopqrstuvwxyz'), 'accept')]")
							elx.click()
						except Exception:
							pass
					except Exception:
						pass
					# Layered waits: try multi-selector visible, then present, then small sleep
					elem = None
					for s in _split_selectors(css_selector or "body"):
						try:
							if _is_xpath(s):
								WebDriverWait(browser, int(wait_visible_secs or _get_timeout_secs())).until(EC.visibility_of_element_located((By.XPATH, s[6:] if s.lower().startswith("xpath=") else s)))
								elem = browser.find_element(By.XPATH, s[6:] if s.lower().startswith("xpath=") else s)
							else:
								WebDriverWait(browser, int(wait_visible_secs or _get_timeout_secs())).until(EC.visibility_of_element_located((By.CSS_SELECTOR, s)))
								elem = browser.find_element(By.CSS_SELECTOR, s)
							css_selector = s
							break
						except Exception:
							try:
								if _is_xpath(s):
									WebDriverWait(browser, max(2, int((wait_visible_secs or _get_timeout_secs()) // 4))).until(EC.presence_of_element_located((By.XPATH, s[6:] if s.lower().startswith("xpath=") else s)))
									elem = browser.find_element(By.XPATH, s[6:] if s.lower().startswith("xpath=") else s)
								else:
									WebDriverWait(browser, max(2, int((wait_visible_secs or _get_timeout_secs()) // 4))).until(EC.presence_of_element_located((By.CSS_SELECTOR, s)))
									elem = browser.find_element(By.CSS_SELECTOR, s)
								css_selector = s
								break
							except Exception:
								continue
					text = ""
					if elem is not None:
						try:
							# Scroll element into view to encourage layout/text updates
							try:
								browser.execute_script("arguments[0].scrollIntoView({block: 'center', inline: 'center'})", elem)
							except Exception:
								pass
							text = elem.text or elem.get_attribute("textContent") or ""
							text = text.strip()
						except StaleElementReferenceException:
							# Re-acquire element once
							try:
								WebDriverWait(browser, 3).until(EC.visibility_of_element_located((By.CSS_SELECTOR, css_selector)))
								elem = browser.find_element(By.CSS_SELECTOR, css_selector)
								text = elem.text or elem.get_attribute("textContent") or ""
								text = text.strip()
							except Exception:
								text = ""
					# Fallback: if still empty, try JS querySelector to read innerText/textContent (handles hidden/inert nodes)
					if not text:
						try:
							js_text = browser.execute_script(
								"var el=document.querySelector(arguments[0]); return el ? (el.innerText || el.textContent || el.value || '') : '';",
								css_selector,
							)
							if isinstance(js_text, str) and js_text.strip():
								text = js_text.strip()
						except Exception:
							pass
					# Optional debug: save screenshot for empty results
					if not text and _should_save_screenshot():
						try:
							import base64  # type: ignore
							png_b64 = browser.get_screenshot_as_base64()
							# write to sandbox tmp under repo
							tmp_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "sandbox", "tmp")
							try:
								os.makedirs(tmp_dir, exist_ok=True)
							except Exception:
								pass
							with open(os.path.join(tmp_dir, "dom_empty_screenshot.png"), "wb") as f:
								f.write(base64.b64decode(png_b64))
						except Exception:
							pass
					# If empty and first attempt, retry once with a new driver
					if text or attempt == 1:
						err_type = None
						if not text:
							err_type = "empty_text" if elem is not None else "not_found"
						try:
							meter_cost("dom_driver", 1.0)
						except Exception:
							pass
						return {"url": url, "selector": css_selector, "text": text, "error_type": err_type, "attempts": int(attempt + 1)}
				finally:
					if browser is not None:
						try:
							browser.quit()
						except Exception:
							pass
			# If all attempts failed
			text = ""
			return {"url": url, "selector": css_selector, "text": text, "error_type": "not_found", "attempts": int(retries)}
		except Exception:
			# Any driver-level error → graceful HTTP fallback
			return _http_fetch(url, css_selector)
		finally:
			# Already closed per attempt; nothing to do
			pass

def click_then_fetch(start_url: str, click_selector: str, target_selector: str, wait_visible_secs: int = 8) -> Dict[str, Any]:
	"""
	Open start_url, click element matching click_selector, then fetch text at target_selector.
	Returns {url, selector, text} with text possibly empty on failure.
	"""
	if not _is_host_allowed(start_url):
		return {"url": start_url, "selector": target_selector, "text": "", "blocked": True}
	try:
		import undetected_chromedriver as uc  # type: ignore
		from selenium.webdriver.common.by import By  # type: ignore
		from selenium.webdriver.support.ui import WebDriverWait  # type: ignore
		from selenium.webdriver.support import expected_conditions as EC  # type: ignore
		from selenium.common.exceptions import TimeoutException  # type: ignore
	except Exception:
		raise RuntimeError("undetected-chromedriver/selenium not installed")

	opts = uc.ChromeOptions()
	opts.add_argument("--headless=new")


						except Exception:
							pass
						return {"url": url, "selector": css_selector, "text": text, "error_type": err_type, "attempts": int(attempt + 1)}
				finally:
					if browser is not None:
						try:
							browser.quit()
						except Exception:
							pass
			# If all attempts failed
			text = ""
			return {"url": url, "selector": css_selector, "text": text, "error_type": "not_found", "attempts": int(retries)}
		except Exception:
			# Any driver-level error → graceful HTTP fallback
			return _http_fetch(url, css_selector)
		finally:
			# Already closed per attempt; nothing to do
			pass

def click_then_fetch(start_url: str, click_selector: str, target_selector: str, wait_visible_secs: int = 8) -> Dict[str, Any]:
	"""
	Open start_url, click element matching click_selector, then fetch text at target_selector.
	Returns {url, selector, text} with text possibly empty on failure.
	"""
	if not _is_host_allowed(start_url):
		return {"url": start_url, "selector": target_selector, "text": "", "blocked": True}
	try:
		import undetected_chromedriver as uc  # type: ignore
		from selenium.webdriver.common.by import By  # type: ignore
		from selenium.webdriver.support.ui import WebDriverWait  # type: ignore
		from selenium.webdriver.support import expected_conditions as EC  # type: ignore
		from selenium.common.exceptions import TimeoutException  # type: ignore
	except Exception:
		raise RuntimeError("undetected-chromedriver/selenium not installed")

	opts = uc.ChromeOptions()
	opts.add_argument("--headless=new")
	opts.add_argument("--no-sandbox")
	opts.add_argument("--disable-dev-shm-usage")
	try:
		opts.page_load_strategy = 'eager'  # type: ignore[attr-defined]
	except Exception:
		pass

	browser = None
	try:
		browser = uc.Chrome(options=opts)
		try:
			browser.set_page_load_timeout(_get_timeout_secs())
		except Exception:
			pass
		_install_net_guards(browser)
		try:
			browser.get(start_url)
		except TimeoutException:
			try:
				browser.execute_script("window.stop();")
			except Exception:
				pass
		# Try to click if present (multi-selector and XPath supported)
		clicked = False
		for s in _split_selectors(click_selector or ""):
			try:
				if _is_xpath(s):
					xp = s[6:] if s.lower().startswith("xpath=") else s
					WebDriverWait(browser, int(max(1, wait_visible_secs))).until(EC.element_to_be_clickable((By.XPATH, xp)))
					el = browser.find_element(By.XPATH, xp)
				else:
					WebDriverWait(browser, int(max(1, wait_visible_secs))).until(EC.element_to_be_clickable((By.CSS_SELECTOR, s)))
					el = browser.find_element(By.CSS_SELECTOR, s)
				try:
					browser.execute_script("arguments[0].scrollIntoView({block: 'center'});", el)
				except Exception:
					pass
				el.click()
				clicked = True
				break
			except Exception:
				continue
		# After click, fetch target (multi-selector, XPath, and JS fallback)
		picked = target_selector
		elem = None
		for s in _split_selectors(target_selector or "body"):
			try:
				if _is_xpath(s):
					xp = s[6:] if s.lower().startswith("xpath=") else s
					WebDriverWait(browser, int(max(1, wait_visible_secs))).until(EC.visibility_of_element_located((By.XPATH, xp)))
					elem = browser.find_element(By.XPATH, xp)
				else:
					WebDriverWait(browser, int(max(1, wait_visible_secs))).until(EC.visibility_of_element_located((By.CSS_SELECTOR, s)))
					elem = browser.find_element(By.CSS_SELECTOR, s)
				picked = s
				break
			except Exception:
				try:
					if _is_xpath(s):
						xp = s[6:] if s.lower().startswith("xpath=") else s
						WebDriverWait(browser, max(2, int(wait_visible_secs // 4))).until(EC.presence_of_element_located((By.XPATH, xp)))
						elem = browser.find_element(By.XPATH, xp)
					else:
						WebDriverWait(browser, max(2, int(wait_visible_secs // 4))).until(EC.presence_of_element_located((By.CSS_SELECTOR, s)))
						elem = browser.find_element(By.CSS_SELECTOR, s)
					picked = s
					break
				except Exception:
					continue
		text = ""
		if elem is not None:
			try:
				try:
					browser.execute_script("arguments[0].scrollIntoView({block: 'center'});", elem)
				except Exception:
					pass
				text = (elem.text or elem.get_attribute("textContent") or "").strip()
			except Exception:
				text = ""
		# JS fallback for reading text
		if not text:
			try:
				if _is_xpath(picked):
					xp = picked[6:] if picked.lower().startswith("xpath=") else picked
					js_text = browser.execute_script(
						"var xp=arguments[0]; var n=document.evaluate(xp, document, null, XPathResult.FIRST_ORDERED_NODE_TYPE, null).singleNodeValue; return n ? (n.innerText || n.textContent || n.value || '') : '';",
						xp,
					)
				else:
					js_text = browser.execute_script(
						"var el=document.querySelector(arguments[0]); return el ? (el.innerText || el.textContent || el.value || '') : '';",
						picked,
					)
				if isinstance(js_text, str) and js_text.strip():
					text = js_text.strip()
			except Exception:
				pass
		return {"url": start_url, "selector": picked, "text": text}
	finally:
		if browser is not None:
			try:
				browser.quit()
			except Exception:
				pass

def form_fill_fetch(
	start_url: str,
	input_selector: str,
	text: str,
	submit_selector: str | None = None,
	target_selector: str | None = None,
	wait_visible_secs: int = 8,
) -> Dict[str, Any]:
	"""
	Open start_url, type text into input_selector, optionally click submit_selector,
	then read target_selector (or the input's value if not provided). Returns {url, selector, text}.
	"""
	if not _is_host_allowed(start_url):
		return {"url": start_url, "selector": target_selector or input_selector, "text": "", "blocked": True}
	try:
		import undetected_chromedriver as uc  # type: ignore
		from selenium.webdriver.common.by import By  # type: ignore
		from selenium.webdriver.support.ui import WebDriverWait  # type: ignore


						xp,
					)
				else:
					js_text = browser.execute_script(
						"var el=document.querySelector(arguments[0]); return el ? (el.innerText || el.textContent || el.value || '') : '';",
						picked,
					)
				if isinstance(js_text, str) and js_text.strip():
					text = js_text.strip()
			except Exception:
				pass
		return {"url": start_url, "selector": picked, "text": text}
	finally:
		if browser is not None:
			try:
				browser.quit()
			except Exception:
				pass

def form_fill_fetch(
	start_url: str,
	input_selector: str,
	text: str,
	submit_selector: str | None = None,
	target_selector: str | None = None,
	wait_visible_secs: int = 8,
) -> Dict[str, Any]:
	"""
	Open start_url, type text into input_selector, optionally click submit_selector,
	then read target_selector (or the input's value if not provided). Returns {url, selector, text}.
	"""
	if not _is_host_allowed(start_url):
		return {"url": start_url, "selector": target_selector or input_selector, "text": "", "blocked": True}
	try:
		import undetected_chromedriver as uc  # type: ignore
		from selenium.webdriver.common.by import By  # type: ignore
		from selenium.webdriver.support.ui import WebDriverWait  # type: ignore
		from selenium.webdriver.support import expected_conditions as EC  # type: ignore
		from selenium.common.exceptions import TimeoutException  # type: ignore
	except Exception:
		raise RuntimeError("undetected-chromedriver/selenium not installed")

	opts = uc.ChromeOptions()
	opts.add_argument("--headless=new")
	opts.add_argument("--no-sandbox")
	opts.add_argument("--disable-dev-shm-usage")
	try:
		opts.page_load_strategy = 'eager'  # type: ignore[attr-defined]
	except Exception:
		pass

	browser = None
	try:
		browser = uc.Chrome(options=opts)
		try:
			browser.set_page_load_timeout(_get_timeout_secs())
		except Exception:
			pass
		_install_net_guards(browser)
		try:
			browser.get(start_url)
		except TimeoutException:
			try:
				browser.execute_script("window.stop();")
			except Exception:
				pass
		# Fill the input (support multi-selector and XPath)
		inp = None
		picked_input = input_selector
		for s in _split_selectors(input_selector or ""):
			try:
				if _is_xpath(s):
					xp = s[6:] if s.lower().startswith("xpath=") else s
					WebDriverWait(browser, int(max(1, wait_visible_secs))).until(EC.element_to_be_clickable((By.XPATH, xp)))
					inp = browser.find_element(By.XPATH, xp)
				else:
					WebDriverWait(browser, int(max(1, wait_visible_secs))).until(EC.element_to_be_clickable((By.CSS_SELECTOR, s)))
					inp = browser.find_element(By.CSS_SELECTOR, s)
				picked_input = s
				break
			except Exception:
				continue
		if inp is None:
			return {"url": start_url, "selector": picked_input, "text": ""}
		try:
			inp.clear()
		except Exception:
			pass
		try:
			browser.execute_script("arguments[0].scrollIntoView({block: 'center'});", inp)
		except Exception:
			pass
		inp.send_keys(text)
		# Optionally submit
		if submit_selector:
			clicked = False
			for s in _split_selectors(submit_selector or ""):
				try:
					if _is_xpath(s):
						xp = s[6:] if s.lower().startswith("xpath=") else s
						WebDriverWait(browser, int(max(1, wait_visible_secs))).until(EC.element_to_be_clickable((By.XPATH, xp)))
						btn = browser.find_element(By.XPATH, xp)
					else:
						WebDriverWait(browser, int(max(1, wait_visible_secs))).until(EC.element_to_be_clickable((By.CSS_SELECTOR, s)))
						btn = browser.find_element(By.CSS_SELECTOR, s)
					btn.click()
					clicked = True
					break
				except Exception:
					continue
		# Read target value/text
		sel = target_selector or picked_input
		val = ""
		el = None
		for s in _split_selectors(sel or ""):
			try:
				if _is_xpath(s):
					xp = s[6:] if s.lower().startswith("xpath=") else s
					WebDriverWait(browser, int(max(1, wait_visible_secs))).until(EC.presence_of_element_located((By.XPATH, xp)))
					el = browser.find_element(By.XPATH, xp)
				else:
					WebDriverWait(browser, int(max(1, wait_visible_secs))).until(EC.presence_of_element_located((By.CSS_SELECTOR, s)))
					el = browser.find_element(By.CSS_SELECTOR, s)
					sel = s
				break
			except Exception:
				continue
		if el is not None:
			try:
				val = el.text or el.get_attribute("textContent") or el.get_attribute("value") or ""
				val = (val or "").strip()
			except Exception:
				val = ""
		# JS fallback to read value/text if empty
		if not val:
			try:
				if _is_xpath(sel):
					xp = sel[6:] if sel.lower().startswith("xpath=") else sel
					js_val = browser.execute_script(
						"var xp=arguments[0]; var n=document.evaluate(xp, document, null, XPathResult.FIRST_ORDERED_NODE_TYPE, null).singleNodeValue; return n ? (n.innerText || n.textContent || n.value || '') : '';",
						xp,
					)
				else:
					js_val = browser.execute_script(
						"var el=document.querySelector(arguments[0]); return el ? (el.innerText || el.textContent || el.value || '') : '';",
						sel,
					)
				if isinstance(js_val, str) and js_val.strip():
					val = js_val.strip()
			except Exception:
				pass
		return {"url": start_url, "selector": sel, "text": val}
	finally:
		if browser is not None:
			try:
				browser.quit()
			except Exception:
				pass


def fetch_text(url: str, css_selector: str = "body", wait_visible_secs: int | None = None) -> Dict[str, Any]:
	# Egress allowlist enforcement
	if not _is_host_allowed(url):
		return {"url": url, "selector": css_selector, "text": "", "blocked": True}
	with trace_span("dom_fetch", {"url": url, "selector": css_selector}):
		try:
			import undetected_chromedriver as uc  # type: ignore
			from selenium.common.exceptions import SessionNotCreatedException  # type: ignore
			exists = True
		except Exception:
			exists = False
		if not exists:
			# Graceful fallback to HTTP fetch when ChromeDriver/UC is unavailable
			return _http_fetch(url, css_selector)
		from selenium.webdriver.common.by import By  # type: ignore
		from selenium.webdriver.support.ui import WebDriverWait  # type: ignore
		from selenium.webdriver.support import expected_conditions as EC  # type: ignore
		from selenium.common.exceptions import StaleElementReferenceException, TimeoutException  # type: ignore
		import re

		def _make_opts():
			# Create a FRESH options object each time; uc forbids reusing the same instance
			o = uc.ChromeOptions()
			o.add_argument("--headless=new")
			o.add_argument("--no-sandbox")
			o.add_argument("--disable-dev-shm-usage")
			# Reduce blocking on full page load
			try:
				# 'eager' returns after DOMContentLoaded
				o.page_load_strategy = 'eager'  # type: ignore[attr-defined]
			except Exception:
				pass
			return o
		browser = None
		try:
			# Attempt multiple cold-driver attempts to mitigate flakiness
			retries = max(1, _get_retries())
			for attempt in range(retries):
				try:
					try:
						# Fresh options per attempt to avoid reuse error
						opts = _make_opts()
						_apply_proxy_to_options(opts)
						chrome_binary = _get_chrome_binary()
						if chrome_binary:
							try:
								opts.binary_location = chrome_binary  # type: ignore[attr-defined]
							except Exception:
								pass
						cd_path = _get_chromedriver_path()
						if cd_path:
							browser = uc.Chrome(driver_executable_path=cd_path, options=opts)
						else:


import logging
from typing import Any, Dict, Optional, List, Tuple
from pathlib import Path
from datetime import datetime

from agi_dw.core.llm.hf_client import HFClient
from agi_dw.core.utils.prompt_logger import PromptLogger, get_prompt_logger  # type: ignore
from agi_dw.core.ops.tracing import trace_span, meter_cost  # type: ignore
from agi_dw.tools.artifact_cache import is_enabled as cache_enabled, get_cached_text as cache_get, set_cached_text as cache_set  # type: ignore

try:
	import yaml  # type: ignore
except Exception:
	yaml = None

# Reuse actuator parsing/coercion to salvage slightly malformed outputs
try:
	from agi_dw.core.actuator.parse import coerce_flat_yaml  # type: ignore
except Exception:
	def coerce_flat_yaml(_: str) -> Dict[str, Any]:  # fallback no-op
		return {}

import re
import json
import os

def verify_sequence(
	sequence: Dict[str, Any],
	state_requirements: Optional[Dict[int, List[str]]] = None,
	validation_points: Optional[List[Dict[str, Any]]] = None
) -> Dict[str, Any]:
	"""Verify sequence execution against requirements and validation points."""
	results = {
		"state_requirements": {
			"satisfied": [],
			"missing": []
		},
		"validation_points": []
	}

	# Check state requirements
	if state_requirements:
		for step_idx, required in state_requirements.items():
			step = sequence.get("steps", [])[step_idx] if step_idx < len(sequence.get("steps", [])) else None
			if step:
				state = step.get("state", {})
				for req in required:
					if _check_state_requirement(state, req):
						results["state_requirements"]["satisfied"].append(req)
					else:
						results["state_requirements"]["missing"].append(req)

	# Check validation points
	if validation_points:
		for point in validation_points:
			step_idx = point.get("step")
			check = point.get("check")
			step = sequence.get("steps", [])[step_idx] if step_idx < len(sequence.get("steps", [])) else None

			if step and check:
				passed, details = _check_validation_point(step, check)
				results["validation_points"].append({
					"step": step_idx,
					"check": check,
					"passed": passed,
					"details": details
				})

	return results

def _check_state_requirement(state: Dict[str, Any], requirement: str) -> bool:
	"""Check if a state requirement is satisfied."""
	# Handle form state
	if requirement.startswith("form."):
		field = requirement.split(".", 1)[1]
		return bool(state.get("form_state", {}).get(field))

	# Handle navigation state
	if requirement.startswith("nav."):
		page = requirement.split(".", 1)[1]
		return page in state.get("navigation_history", [])

	# Handle element state
	if requirement.startswith("element."):
		selector = requirement.split(".", 1)[1]
		return bool(state.get("dom_state", {}).get(selector))

	return False

def _check_validation_point(step: Dict[str, Any], check: str) -> Tuple[bool, str]:
	"""Check if a validation point passes."""
	state = step.get("state", {})
	observation = step.get("observation", {})

	# Handle form validation
	if check.startswith("form."):
		field = check.split(".", 1)[1]
		if field in state.get("form_state", {}):
			return True, f"Form field {field} is set"
		return False, f"Form field {field} not found"

	# Handle element validation
	if check.startswith("element."):
		selector = check.split(".", 1)[1]
		if selector in state.get("dom_state", {}):
			return True, f"Element {selector} found"
		return False, f"Element {selector} not found"

	# Handle content validation
	if check.startswith("content."):
		text = check.split(".", 1)[1]
		if text.lower() in str(observation.get("content", "")).lower():
			return True, f"Content '{text}' found"
		return False, f"Content '{text}' not found"

	return False, "Unknown validation check"

def verify_trace_snippet(
	trace: Dict[str, Any],
	model: str = "meta-llama/Llama-3.1-8B-Instruct",
	timeout_sec: int = 8,
	use_llm: bool = True,
	require_llm: bool = False,
	backend: str = "hf",
	log_prompts: bool = False,
	adapter_dir: str | None = None,
	structured_mode: str = "none",
) -> Dict[str, Any]:
	"""
	LLM verifier with optional strict mode. HF backend only.
	Prompts require YAML output; we parse YAML (or JSON fallback) to a dict.
	"""
	if use_llm:
		with trace_span("verify_trace", {"backend": backend, "structured": structured_mode}):
			logger = get_prompt_logger("verifier", bool(log_prompts), echo=bool(log_prompts))
			def _log(kind: str, text: str) -> None:
				logger.log_text(kind, text)
			def _red(s: str) -> str:
				try:
					from agi_dw.core.utils.redact import redact_text  # type: ignore
					rs, _ = redact_text(s)
					return rs
				except Exception:
					return s
			if backend == "hf":
				try:
					# Prefer local PEFT adapter if provided
					if adapter_dir:
						try:
							from agi_dw.core.llm.adapter_cache import AdapterCache  # type: ignore
							tok, peft_model = AdapterCache.get(model, adapter_dir)
							prompt = (
								"You are a strict verifier. Given a trace snippet, output ONLY a YAML mapping with keys"
								" success_prob (float 0..1), risk (float 0..1), critique (short string). No prose.\n"
								"Example:\nsuccess_prob: 0.72\nrisk: 0.18\ncritique: brief reason\n\n"
								"Trace (YAML or JSON content below; do NOT echo it):\n" + str(trace)
							)
							if log_prompts:
								_log("prompt", prompt)
							enc = tok(prompt, return_tensors="pt").to(peft_model.device)
							with torch.inference_mode():
								out_ids = peft_model.generate(
									**enc,
									max_new_tokens=200,
									do_sample=False,
									num_beams=1,
								)
							text = tok.decode(out_ids[0], skip_special_tokens=True)
							try:
								meter_cost("verify_llm", 1.0)
							except Exception:
								pass
							if log_prompts:
								_log("response", text)
							parsed: Dict[str, Any] = _robust_struct_parse(text)
						except Exception as e:
							if log_prompts:
								print(f"[HF VERIFY WARN] Adapter load failed: {e}")
							parsed = {}
							text = ""
					else:
						parsed = {}
						text = ""
					# Fallback to shared HF client (base model)
					if not parsed:
						# Optional structured decoding via Outlines (JSON schema)
						if structured_mode == "json":
							try:
								# Lazy import outlines; optional dependency
								from outlines import models as _out_models  # type: ignore
								from outlines import generate as _out_generate  # type: ignore
								schema = {
									"type": "object",
									"properties": {
										"success_prob": {"type": "number", "minimum": 0.0, "maximum": 1.0},
										"risk": {"type": "number", "minimum": 0.0, "maximum": 1.0},
										"critique": {"type": "string"},
										"validation_results": {
											"type": "object",
											"properties": {
												"state_requirements": {
													"type": "object",
													"properties": {
														"satisfied": {"type": "array", "items": {"type": "string"}},
														"missing": {"type": "array", "items": {"type": "string"}}
													}
												},
												"validation_points": {
													"type": "array",
													"items": {
														"type": "object",
														"properties": {
															"step": {"type": "integer"},
															"check": {"type": "string"},
															"passed": {"type": "boolean"},
															"details": {"type": "string"}
														},
														"required": ["step", "check", "passed"]
													}
												}
											}
										}
									},
									"required": ["success_prob", "risk", "critique"],
									"additionalProperties": False,
								}
								base_prompt = (
									"You are a strict verifier. Given a trace snippet, return ONLY a JSON object with:\n"
									"1. success_prob (0..1): Overall success probability\n"
									"2. risk (0..1): Overall risk assessment\n"
									"3. critique: Brief explanation\n"
									"4. validation_results:\n"
									"   - state_requirements: Track satisfied and missing state elements\n"
									"   - validation_points: Array of validation checks with step, check, passed, and details\n"
									"Example:\n"
									'{"success_prob":0.8,"risk":0.2,"critique":"Login successful","validation_results":{'
									'"state_requirements":{"satisfied":["login_page"],"missing":[]},'
									'"validation_points":[{"step":0,"check":"Form visible","passed":true,"details":"Found #login-form"}]}}\n\n'
									"Trace (YAML/JSON below; do NOT echo it):\n" + str(trace)
								)
								if log_prompts:
									_log("prompt", base_prompt)
								# Optional cache for structured verify
								use_cache = cache_enabled("AGI_VERIFY_CACHE")
								ttl = int(os.environ.get("AGI_VERIFY_CACHE_TTL", "1800") or 1800) if use_cache else 0
								text = ""
								if use_cache:
									cached = cache_get("verifier", [model, "json", base_prompt], ttl)
									if isinstance(cached, str) and cached.strip():
										text = cached
								if not text:
									mdl = _out_models.transformers(model)
									generator = _out_generate.json(mdl, schema)
									text = generator(base_prompt)
									try:
										meter_cost("verify_structured", 1.0)
									except Exception:
										pass
									if use_cache and text:
										cache_set("verifier", [model, "json", base_prompt], str(text))
									if log_prompts:
										_log("response", text)
									parsed = _robust_json_parse(text)
							except Exception:
								parsed = {}
								text = ""
						# Standard YAML-first path
						if not parsed:
							client = HFClient.get_cached(model)
							base_prompt = (
								"You are a strict verifier. Given a trace snippet, output ONLY a YAML mapping with keys"
								" success_prob (float 0..1), risk (float 0..1), critique (short string). No prose.\n"
								"Example:\n"
								"success_prob: 0.72\n"
								"risk: 0.18\n"
								"critique: brief reason\n\n"
								"Trace (YAML or JSON content below; do NOT echo it):\n" + str(trace)
							)
							if log_prompts:
								_log("prompt", base_prompt)
							# Optional cache for YAML verify
							use_cache = cache_enabled("AGI_VERIFY_CACHE")
							ttl = int(os.environ.get("AGI_VERIFY_CACHE_TTL", "1800") or 1800) if use_cache else 0
							text = ""
							if use_cache:
								cached = cache_get("verifier", [model, "yaml", base_prompt], ttl)
								if isinstance(cached, str) and cached.strip():
									text = cached
							if not text:
								text = client.generate(base_prompt, max_new_tokens=200, temperature=0.0)
								try:
									meter_cost("verify_yaml", 1.0)
								except Exception:
									pass
								if use_cache and text:
									cache_set("verifier", [model, "yaml", base_prompt], str(text))
							if log_prompts:
								_log("response", text)
							parsed = _robust_struct_parse(text)
				if "success_prob" not in parsed or "risk" not in parsed:
					retry_prompt = (
						"Return ONLY YAML with exactly these keys: success_prob, risk, critique. No backticks.\n"
						"Example:\nsuccess_prob: 0.5\nrisk: 0.5\ncritique: short\n\nNow respond:"
					)
					if log_prompts:
						_log("prompt", retry_prompt)
					# Optional cache for retry prompt
					use_cache = cache_enabled("AGI_VERIFY_CACHE")
					ttl = int(os.environ.get("AGI_VERIFY_CACHE_TTL", "1800") or 1800) if use_cache else 0
					text_retry = ""
					if use_cache:
						cached = cache_get("verifier", [model, "retry", retry_prompt], ttl)
						if isinstance(cached, str) and cached.strip():
							text_retry = cached
					if not text_retry:
						text_retry = client.generate(retry_prompt, max_new_tokens=120, temperature=0.0)
						try:
							meter_cost("verify_retry", 1.0)
						except Exception:
							pass
					if log_prompts:
						_log("response", text_retry)
					parsed = _robust_struct_parse(text_retry)
					if ("success_prob" not in parsed or "risk" not in parsed) and text_retry:
						coerced = coerce_flat_yaml(text_retry)
						if "success_prob" in coerced and "risk" in coerced:
							parsed = coerced
					# Final regex fallback before failing strict
					if ("success_prob" not in parsed or "risk" not in parsed) and (text_retry or text):
						fallback_src = text_retry or text
						m1 = re.search(r"success_prob\s*[:=]\s*([0-9]*\.?[0-9]+)", fallback_src)
						m2 = re.search(r"risk\s*[:=]\s*([0-9]*\.?[0-9]+)", fallback_src)
						if m1 and m2:
							parsed = {
								"success_prob": float(m1.group(1)),
								"risk": float(m2.group(1)),
								"critique": "",
							}
				# If LLM responded but still unparsable, return conservative defaults instead of failing hard
				if ("success_prob" not in parsed or "risk" not in parsed) and (text_retry or text):
					status = (trace.get("result") or {}).get("status")
					fallback_risk = 0.2 if status == "ok" else 0.6
					fallback_sp = 0.8 if status == "ok" else 0.3
					parsed = {"success_prob": fallback_sp, "risk": fallback_risk, "critique": "unparsable-llm"}
				if "success_prob" in parsed and "risk" in parsed:
					return {
						"success_prob": float(parsed.get("success_prob", 0.5)),
						"risk": float(parsed.get("risk", 0.5)),
						"critique": str(parsed.get("critique", "")),
					}
				if require_llm:
					raise RuntimeError("HF LLM returned no valid YAML verdict")
			except Exception as e:
				if require_llm:
					raise RuntimeError(f"HF LLM call failed: {e}")
		else:
			# HF-only
			pass
	# Heuristic fallback
	return {"success_prob": 0.5, "risk": 0.5, "critique": "no-llm"}

def _robust_json_parse(text: str) -> Dict[str, Any]:
	try:
		return json.loads(text)
	except Exception:
		return {}

def _robust_struct_parse(text: str) -> Dict[str, Any]:
	# YAML first
	if yaml is not None:
		try:
			y = yaml.safe_load(text)
			if isinstance(y, dict):
				return y
		except Exception:
			pass
	# JSON fallback
	t = text.strip()
	if t.startswith("{") and t.endswith("}"):
		try:
			return json.loads(t)
		except Exception:
			return {}
	start = t.find("{")
	end = t.rfind("}")
	if start != -1 and end != -1 and end > start:
		try:
			return json.loads(t[start : end + 1])
		except Exception:
			return {}
	return {}


from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, List, Tuple, Optional, Callable

import time
import hashlib

from agi_dw.bench.common.pipeline import (
	try_generate_python_body,
	looks_like_python_code,
	attach_adapter_from_args,
	parse_k_list,
	evaluate_samples,
	read_jsonl as _read_jsonl_shared,
	build_results_by_task as _build_results_by_task_shared,
	dedupe_by_passed as _dedupe_by_passed_shared,
	shard_task_ids,
	compute_sharded_outpath,
	write_trace,
	repair_code_failures,
	write_suite_out,
	write_run_artifact,
	evaluate_candidate_signals,
	maybe_self_refine_body,
	wm_rerank_candidates,
    get_python_version,
    build_env_fingerprint,
    is_primary_process,
    seed_everything,
    get_shard_tag,
    init_trace_paths,
    init_suite_temp_paths,
    read_pass_cache,
    update_pass_cache,
    build_generation_params,
    summarize_telemetry,
    online_update_and_rerun_failures,
    emit_training_datasets,
    write_samples_and_verbose,
    store_passed_solutions_in_memory,
)
from agi_dw.core.ops.tracing import trace_span, meter_cost  # type: ignore

def default_env_fingerprint() -> Dict[str, Any]:
    return build_env_fingerprint(["torch", "datasets", "transformers", "human_eval", "evalplus"])  # conservative superset

def resolve_default_workers() -> int:
    try:
        import os
        cpu_cnt = int((os.cpu_count() or 4))
        return min(16, max(1, cpu_cnt))
    except Exception:
        return 4

DEFAULT_WORKERS = resolve_default_workers()

def get_llm_and_basics(args: Any, root: Path, suite: str, default_adapter: Optional[Path] = None):
    from agi_dw.core.llm.hf_client import HFClient  # type: ignore
    from agi_dw.core.prompts.bench import build_prompt  # type: ignore
    from agi_dw.scripts.bench.cache import LLMCache  # type: ignore
    from agi_dw.core.utils.prompt_logger import PromptLogger  # type: ignore
    from agi_dw.core.utils.critic import get_critic  # type: ignore
    from agi_dw.core.utils.bench_utils import ensure_safe_env  # type: ignore

    ensure_safe_env()
    llm = HFClient.get_cached(str(getattr(args, "model", "meta-llama/Llama-3.2-3B")))
    try:
        attach_adapter_from_args(llm, args, root, default_adapter=default_adapter)
    except Exception:
        pass
    cache = LLMCache(root / "data" / "bench" / "cache")
    from agi_dw.core.utils.prompt_logger import get_prompt_logger  # type: ignore
    logger = get_prompt_logger(suite, bool(getattr(args, "log_prompts", False)))
    critic = None
    if bool(getattr(args, "critic", False)):
        cm = getattr(args, "critic_model", None)
        if isinstance(cm, str) and cm.strip().lower() in ("", "none", "null"):
            cm = None
        critic = get_critic(cm)
    return llm, cache, logger, critic

def generate_candidates_for_prompt(
    llm: Any,
    cache: Any,
    logger: Any,
    critic: Any,
    prompt: str,
    inp: str,
    sanitize_fn: Callable[[str, str], str],
    task_id: str,
    args: Any,
    suite: str,
    env_fp: Dict[str, Any],
    trace_path: Path,
    _prompts_hasher: Any,
) -> Tuple[List[Dict[str, Any]], str]:
    params, grammar_constrained, n_samples, n_candidates, retries, retry_backoff = build_generation_params(args)
    try:
        _prompts_hasher.update(inp.encode("utf-8", errors="ignore"))
    except Exception:
        pass
    from agi_dw.core.utils.bench_utils import strip_fences, precheck_code, retry_with_backoff  # type: ignore
    try:
        from agi_dw.core.verifier.llm_verifier import verify_trace_snippet  # type: ignore
    except Exception:
        verify_trace_snippet = None  # type: ignore

    lines: List[Dict[str, Any]] = []
    candidates: List[Dict[str, Any]] = []
    chosen: str = ""

    if n_candidates == 1:
        clean: str = ""
        last_err: Optional[str] = None
        for attempt in range(max(1, retries + 1)):
            comp = None
            if attempt == 0:
                comp = cache.get(str(getattr(args, "model")), inp, params)
                if comp is None:
                    def _gen() -> str:
                        return try_generate_python_body(llm, inp, params, grammar_constrained)
                    import time as _time
                    _t0 = _time.time()
                    text = retry_with_backoff(_gen, retries, retry_backoff) or ""
                    _latency_ms = int((_time.time() - _t0) * 1000)
                    comp = text
            clean_try = sanitize_fn(strip_fences(str(comp or "")), prompt)
            if bool(getattr(args, "precheck", False)):
                ok, err_msg = precheck_code(clean_try)
                if not ok:
                    last_err = err_msg or "precheck failed"
                    if attempt < retries:
                        continue
            if not looks_like_python_code(clean_try):
                last_err = "rejected_noncode"
                if attempt < retries:
                    continue
            if critic and clean_try:
                ok, _ = critic.review(clean_try)
                if not ok:
                    last_err = "critic rejected"
                    if attempt < retries:
                        continue
            if not clean_try.strip() and attempt < retries:
                last_err = "empty after sanitize"
                continue
            clean = clean_try
            if attempt == 0 and comp and clean.strip():
                cache.put(str(getattr(args, "model")), inp, params, str(comp))
            break
        # Optional self-eval/refinement loop for single-candidate path
        if bool(getattr(args, "pav_enable", False)) and clean.strip():
            try:
                max_iters = max(1, int(getattr(args, "pav_max_iters", 2) or 2))
                for _iter in range(max_iters):
                    # Evaluate signals
                    ch_pre_ok = False
                    ch_crit_ok = False
                    ch_v_risk = 0.5
                    try:
                        ch_pre_ok, _ = precheck_code(clean)
                    except Exception:
                        ch_pre_ok = False
                    if critic and clean:
                        try:
                            ch_crit_ok, _ = critic.review(clean)
                        except Exception:
                            ch_crit_ok = False
                    if verify_trace_snippet is not None and clean:
                        try:
                            vres = verify_trace_snippet(
                                trace={"obs": {"kind": "code", "content": prompt}, "plan": {}, "action": {"completion": clean}, "result": {"status": "pending"}},
                                model=str(getattr(args, "verifier_model", "meta-llama/Llama-3.2-3B")),
                                backend=str(getattr(args, "verifier_backend", "hf")),
                                use_llm=True,
                                timeout=_safe_int(getattr(args, "timeout", 15) or 15, 15),
                            )
                            ch_v_risk = float(vres.get("risk", 0.5))
                        except Exception:
                            ch_v_risk = 0.5
                    need_refine = (not ch_pre_ok) or (critic is not None and not ch_crit_ok) or (ch_v_risk > float(getattr(args, "verify_risk_threshold", 0.6) or 0.6)) or bool(getattr(args, "self_eval", False))
                    if not need_refine:
                        break
                    review_system = "You are a Python coding assistant. Improve the function body minimally to pass the tests. Output ONLY the corrected function body."
                    review_user = (
                        "Task prompt (do not rewrite def):\n" + prompt
                        + "\n\nCurrent attempt body:\n" + clean
                        + "\n\nInstruction: Return a corrected function body only."
                    )
                    try:
                        messages = [{"role": "system", "content": review_system}, {"role": "user", "content": review_user}]
                        text = str(llm.chat(messages, max_new_tokens=params.get("max_new_tokens", 256), temperature=params.get("temperature", 0.2), top_p=params.get("top_p"), top_k=params.get("top_k")))
                    except Exception:
                        text = str(llm.generate(review_user, **params))
                    rev = sanitize_fn(strip_fences(text), prompt)
                    if looks_like_python_code(rev):
                        rev_pre_ok = False
                        try:
                            rev_pre_ok, _ = precheck_code(rev)
                        except Exception:
                            rev_pre_ok = False
                        rev_v_risk = ch_v_risk
                        if verify_trace_snippet is not None:
                            try:
                                vres2 = verify_trace_snippet(
                                    trace={"obs": {"kind": "code", "content": prompt}, "plan": {}, "action": {"completion": rev}, "result": {"status": "pending"}},
                                    model=str(getattr(args, "verifier_model", "meta-llama/Llama-3.2-3B")),
                                    backend=str(getattr(args, "verifier_backend", "hf")),
                                    use_llm=True,
                                    timeout=_safe_int(getattr(args, "timeout", 15) or 15, 15),
                                )
                                rev_v_risk = float(vres2.get("risk", ch_v_risk))
                            except Exception:
                                rev_v_risk = ch_v_risk
                        if rev_pre_ok or (rev_v_risk <= ch_v_risk):
                            clean = rev
                            write_trace(trace_path, suite, task_id, "self_eval", None, (int(getattr(args, "seed")) if getattr(args, "seed", None) is not None else None), env_fp, inp, clean, meta={"rev_pre_ok": bool(rev_pre_ok), "rev_v_risk": float(rev_v_risk)})
                        else:
                            break
            except Exception:
                pass
        write_trace(trace_path, suite, task_id, 0, int(locals().get("_latency_ms", 0)), (int(getattr(args, "seed")) if getattr(args, "seed", None) is not None else None), env_fp, inp, clean)
        chosen = clean
        lines.append({"task_id": task_id, "completion": clean, "orig_prompt": prompt, "input": inp, "role": "chosen", "nbest_rank": 0, "latency_ms": int(locals().get("_latency_ms", 0)), "attempt_idx": 0})
        return lines, chosen

    # n-best path
    for ci in range(n_candidates):
        comp = None
        cand_latency_ms: Optional[int] = None
        if ci == 0:
            comp = cache.get(str(getattr(args, "model")), inp, params)
        if comp is None:
            def _gen_multi() -> str:
                return try_generate_python_body(llm, inp, params, grammar_constrained)
            import time as _time
            _t0 = _time.time()
            text = retry_with_backoff(_gen_multi, retries, retry_backoff) or ""
            cand_latency_ms = int((_time.time() - _t0) * 1000)
            comp = text
            if ci == 0 and comp:
                cache.put(str(getattr(args, "model")), inp, params, str(comp))
        cand = sanitize_fn(strip_fences(str(comp or "")), prompt)
        pre_ok, crit_ok, v_risk = evaluate_candidate_signals(cand, prompt, precheck_code, critic, verify_trace_snippet, args)
        candidates.append({"body": cand, "precheck": pre_ok, "critic": crit_ok, "length": len(cand or ""), "v_risk": v_risk, "latency_ms": cand_latency_ms})
    from agi_dw.bench.common.pipeline import pick_best_candidate  # local import to avoid circulars
    chosen = pick_best_candidate(candidates, float(getattr(args, "verify_risk_threshold", 0.6) or 0.6))
    from agi_dw.bench.common.pipeline import rank_candidates  # type: ignore
    _ranks = rank_candidates(candidates)
    # Adaptive extra sampling if chosen looks risky/low-quality
    try:
        ch_pre_ok = False
        ch_crit_ok = False
        ch_v_risk = 0.5
        if chosen:
            try:
                ch_pre_ok, _ = precheck_code(chosen)
            except Exception:
                ch_pre_ok = False
            if critic:
                try:
                    ch_crit_ok, _ = critic.review(chosen)
                except Exception:
                    ch_crit_ok = False
            if verify_trace_snippet is not None:
                try:
                    vres = verify_trace_snippet(
                        trace={"obs": {"kind": "code", "content": prompt}, "plan": {}, "action": {"completion": chosen}, "result": {"status": "pending"}},
                        model=str(getattr(args, "verifier_model", "meta-llama/Llama-3.2-3B")),
                        backend=str(getattr(args, "verifier_backend", "hf")),
                        use_llm=True,
                        timeout=_safe_int(getattr(args, "timeout", 15) or 15, 15),
                    )
                    ch_v_risk = float(vres.get("risk", 0.5))
                except Exception:
                    ch_v_risk = 0.5
        bad = (not looks_like_python_code(chosen)) or (not ch_pre_ok) or (critic is not None and not ch_crit_ok) or (ch_v_risk > float(getattr(args, "verify_risk_threshold", 0.6) or 0.6))
        if bad:
            M = max(1, _safe_int(getattr(args, "adaptive_extra_candidates", 2) or 2, 2))
            for _ in range(M):
                text = retry_with_backoff(lambda: try_generate_python_body(llm, inp, params, grammar_constrained), retries, retry_backoff) or ""
                cand_try = sanitize_fn(strip_fences(text), prompt)
                if cand_try.strip():
                    try:
                        pre_ok, _ = precheck_code(cand_try)
                    except Exception:
                        pre_ok = False
                    crit_ok = False
                    if critic and cand_try:
                        try:
                            crit_ok, _ = critic.review(cand_try)
                        except Exception:
                            crit_ok = False
                    v_risk = 0.5
                    if verify_trace_snippet is not None:
                        try:
                            vres3 = verify_trace_snippet(
                                trace={"obs": {"kind": "code", "content": prompt}, "plan": {}, "action": {"completion": cand_try}, "result": {"status": "pending"}},
                                model=str(getattr(args, "verifier_model", "meta-llama/Llama-3.2-3B")),
                                backend=str(getattr(args, "verifier_backend", "hf")),
                                use_llm=True,
                                timeout=_safe_int(getattr(args, "timeout", 15) or 15, 15),
                            )
                            v_risk = float(vres3.get("risk", 0.5))
                        except Exception:
                            v_risk = 0.5
                    candidates.append({"body": cand_try, "precheck": bool(pre_ok and looks_like_python_code(cand_try)), "critic": crit_ok, "length": len(cand_try or ""), "v_risk": v_risk, "latency_ms": None})
            from agi_dw.bench.common.pipeline import pick_best_candidate as _pick_best
            chosen = _pick_best(candidates, float(getattr(args, "verify_risk_threshold", 0.6) or 0.6))
    except Exception:
        pass

    # Optional WM re-ranking
    wm_model = None
    wm_enabled = bool(getattr(args, "wm_prior", False))
    if wm_enabled:
        try:
            from agi_dw.core.world_model.api import WorldModelPrior  # type: ignore
        except Exception:
            WorldModelPrior = None  # type: ignore
        try:
            wm_path = Path(getattr(args, "wm_model", Path(Path(__file__).resolve().parents[2]) / "models" / "wm_mlp" / "wm_mlp.joblib"))
            if 'WorldModelPrior' in globals() and WorldModelPrior is not None and wm_path.exists():
                wm_model = WorldModelPrior.load(wm_path)
        except Exception:
            wm_model = None
        if wm_model is not None:
            try:
                chosen = wm_rerank_candidates(wm_model, prompt, candidates, chosen, float(getattr(args, "wm_threshold", 0.6) or 0.6))
            except Exception:
                pass

    # Optional self-evaluation/refinement loop before committing chosen
    if bool(getattr(args, "pav_enable", False)):
        try:
            max_iters = max(1, _safe_int(getattr(args, "pav_max_iters", 2) or 2, 2))
            for _iter in range(max_iters):
                ch_pre_ok = False
                ch_crit_ok = False
                ch_v_risk = 0.5
                try:
                    ch_pre_ok, _ = precheck_code(chosen)
                except Exception:
                    ch_pre_ok = False
                if critic and chosen:
                    try:
                        ch_crit_ok, _ = critic.review(chosen)
                    except Exception:
                        ch_crit_ok = False
                if verify_trace_snippet is not None and chosen:
                    try:
                        vres = verify_trace_snippet(
                            trace={"obs": {"kind": "code", "content": prompt}, "plan": {}, "action": {"completion": chosen}, "result": {"status": "pending"}},
                            model=str(getattr(args, "verifier_model", "meta-llama/Llama-3.2-3B")),
                            backend=str(getattr(args, "verifier_backend", "hf")),
                            use_llm=True,
                            timeout=_safe_int(getattr(args, "timeout", 15) or 15, 15),
                        )
                        ch_v_risk = float(vres.get("risk", 0.5))
                    except Exception:
                        ch_v_risk = 0.5
                need_refine = (not ch_pre_ok) or (critic is not None and not ch_crit_ok) or (ch_v_risk > float(getattr(args, "verify_risk_threshold", 0.6) or 0.6)) or bool(getattr(args, "self_eval", False))
                if not need_refine:
                    break
                review_system = "You are a Python coding assistant. Improve the function body minimally to pass the tests. Output ONLY the corrected function body."
                review_user = (
                    "Task prompt (do not rewrite def):\n" + prompt
                    + "\n\nCurrent attempt body:\n" + chosen
                    + "\n\nInstruction: Return a corrected function body only."
                )
                try:
                    messages = [{"role": "system", "content": review_system}, {"role": "user", "content": review_user}]
                    text = str(llm.chat(messages, max_new_tokens=params.get("max_new_tokens", 256), temperature=params.get("temperature", 0.2), top_p=params.get("top_p"), top_k=params.get("top_k")))
                except Exception:
                    text = str(llm.generate(review_user, **params))
                rev = sanitize_fn(strip_fences(text), prompt)
                if looks_like_python_code(rev):
                    rev_pre_ok = False
                    try:
                        rev_pre_ok, _ = precheck_code(rev)
                    except Exception:
                        rev_pre_ok = False
                    rev_v_risk = ch_v_risk
                    if verify_trace_snippet is not None:
                        try:
                            vres2 = verify_trace_snippet(
                                trace={"obs": {"kind": "code", "content": prompt}, "plan": {}, "action": {"completion": rev}, "result": {"status": "pending"}},
                                model=str(getattr(args, "verifier_model", "meta-llama/Llama-3.2-3B")),
                                backend=str(getattr(args, "verifier_backend", "hf")),
                                use_llm=True,
                                timeout=_safe_int(getattr(args, "timeout", 15) or 15, 15),
                            )
                            rev_v_risk = float(v
# ... truncated ...


	wm_rerank_candidates,
    get_python_version,
    build_env_fingerprint,
    is_primary_process,
    seed_everything,
    get_shard_tag,
    init_trace_paths,
    init_suite_temp_paths,
    read_pass_cache,
    update_pass_cache,
    build_generation_params,
    summarize_telemetry,
    online_update_and_rerun_failures,
    emit_training_datasets,
    write_samples_and_verbose,
    store_passed_solutions_in_memory,
)
from agi_dw.core.ops.tracing import trace_span, meter_cost  # type: ignore

def default_env_fingerprint() -> Dict[str, Any]:
    return build_env_fingerprint(["torch", "datasets", "transformers", "human_eval", "evalplus"])  # conservative superset

def resolve_default_workers() -> int:
    try:
        import os
        cpu_cnt = int((os.cpu_count() or 4))
        return min(16, max(1, cpu_cnt))
    except Exception:
        return 4

DEFAULT_WORKERS = resolve_default_workers()

def get_llm_and_basics(args: Any, root: Path, suite: str, default_adapter: Optional[Path] = None):
    from agi_dw.core.llm.hf_client import HFClient  # type: ignore
    from agi_dw.core.prompts.bench import build_prompt  # type: ignore
    from agi_dw.scripts.bench.cache import LLMCache  # type: ignore
    from agi_dw.core.utils.prompt_logger import PromptLogger  # type: ignore
    from agi_dw.core.utils.critic import get_critic  # type: ignore


    seed_everything,
    get_shard_tag,
    init_trace_paths,
    init_suite_temp_paths,
    read_pass_cache,
    update_pass_cache,
    build_generation_params,
    summarize_telemetry,
    online_update_and_rerun_failures,
    emit_training_datasets,
    write_samples_and_verbose,
    store_passed_solutions_in_memory,
)
from agi_dw.core.ops.tracing import trace_span, meter_cost  # type: ignore

def default_env_fingerprint() -> Dict[str, Any]:
    return build_env_fingerprint(["torch", "datasets", "transformers", "human_eval", "evalplus"])  # conservative superset

def resolve_default_workers() -> int:
    try:
        import os
        cpu_cnt = int((os.cpu_count() or 4))
        return min(16, max(1, cpu_cnt))
    except Exception:
        return 4

DEFAULT_WORKERS = resolve_default_workers()

def get_llm_and_basics(args: Any, root: Path, suite: str, default_adapter: Optional[Path] = None):
    from agi_dw.core.llm.hf_client import HFClient  # type: ignore
    from agi_dw.core.prompts.bench import build_prompt  # type: ignore
    from agi_dw.scripts.bench.cache import LLMCache  # type: ignore
    from agi_dw.core.utils.prompt_logger import PromptLogger  # type: ignore
    from agi_dw.core.utils.critic import get_critic  # type: ignore
    from agi_dw.core.utils.bench_utils import ensure_safe_env  # type: ignore

    ensure_safe_env()
    llm = HFClient.get_cached(str(getattr(args, "model", "meta-llama/Llama-3.2-3B")))
    try:
        attach_adapter_from_args(llm, args, root, default_adapter=default_adapter)
    except Exception:
        pass
    cache = LLMCache(root / "data" / "bench" / "cache")


)
from agi_dw.core.ops.tracing import trace_span, meter_cost  # type: ignore

def default_env_fingerprint() -> Dict[str, Any]:
    return build_env_fingerprint(["torch", "datasets", "transformers", "human_eval", "evalplus"])  # conservative superset

def resolve_default_workers() -> int:
    try:
        import os
        cpu_cnt = int((os.cpu_count() or 4))
        return min(16, max(1, cpu_cnt))
    except Exception:
        return 4

DEFAULT_WORKERS = resolve_default_workers()

def get_llm_and_basics(args: Any, root: Path, suite: str, default_adapter: Optional[Path] = None):
    from agi_dw.core.llm.hf_client import HFClient  # type: ignore
    from agi_dw.core.prompts.bench import build_prompt  # type: ignore
    from agi_dw.scripts.bench.cache import LLMCache  # type: ignore
    from agi_dw.core.utils.prompt_logger import PromptLogger  # type: ignore
    from agi_dw.core.utils.critic import get_critic  # type: ignore
    from agi_dw.core.utils.bench_utils import ensure_safe_env  # type: ignore

    ensure_safe_env()
    llm = HFClient.get_cached(str(getattr(args, "model", "meta-llama/Llama-3.2-3B")))
    try:
        attach_adapter_from_args(llm, args, root, default_adapter=default_adapter)
    except Exception:
        pass
    cache = LLMCache(root / "data" / "bench" / "cache")
    from agi_dw.core.utils.prompt_logger import get_prompt_logger  # type: ignore
    logger = get_prompt_logger(suite, bool(getattr(args, "log_prompts", False)))
    critic = None
    if bool(getattr(args, "critic", False)):
        cm = getattr(args, "critic_model", None)
        if isinstance(cm, str) and cm.strip().lower() in ("", "none", "null"):
            cm = None
        critic = get_critic(cm)
    return llm, cache, logger, critic

def generate_candidates_for_prompt(
    llm: Any,
    cache: Any,
    logger: Any,
    critic: Any,
    prompt: str,
    inp: str,
    sanitize_fn: Callable[[str, str], str],
    task_id: str,
    args: Any,
    suite: str,
    env_fp: Dict[str, Any],
    trace_path: Path,
    _prompts_hasher: Any,
) -> Tuple[List[Dict[str, Any]], str]:
    params, grammar_constrained, n_samples, n_candidates, retries, retry_backoff = build_generation_params(args)
    try:
        _prompts_hasher.update(inp.encode("utf-8", errors="ignore"))


    from agi_dw.core.utils.bench_utils import ensure_safe_env  # type: ignore

    ensure_safe_env()
    llm = HFClient.get_cached(str(getattr(args, "model", "meta-llama/Llama-3.2-3B")))
    try:
        attach_adapter_from_args(llm, args, root, default_adapter=default_adapter)
    except Exception:
        pass
    cache = LLMCache(root / "data" / "bench" / "cache")
    from agi_dw.core.utils.prompt_logger import get_prompt_logger  # type: ignore
    logger = get_prompt_logger(suite, bool(getattr(args, "log_prompts", False)))
    critic = None
    if bool(getattr(args, "critic", False)):
        cm = getattr(args, "critic_model", None)
        if isinstance(cm, str) and cm.strip().lower() in ("", "none", "null"):
            cm = None
        critic = get_critic(cm)
    return llm, cache, logger, critic

def generate_candidates_for_prompt(
    llm: Any,
    cache: Any,
    logger: Any,
    critic: Any,
    prompt: str,
    inp: str,
    sanitize_fn: Callable[[str, str], str],
    task_id: str,
    args: Any,
    suite: str,
    env_fp: Dict[str, Any],
    trace_path: Path,
    _prompts_hasher: Any,
) -> Tuple[List[Dict[str, Any]], str]:
    params, grammar_constrained, n_samples, n_candidates, retries, retry_backoff = build_generation_params(args)
    try:
        _prompts_hasher.update(inp.encode("utf-8", errors="ignore"))
    except Exception:
        pass
    from agi_dw.core.utils.bench_utils import strip_fences, precheck_code, retry_with_backoff  # type: ignore
    try:
        from agi_dw.core.verifier.llm_verifier import verify_trace_snippet  # type: ignore
    except Exception:
        verify_trace_snippet = None  # type: ignore

    lines: List[Dict[str, Any]] = []
    candidates: List[Dict[str, Any]] = []
    chosen: str = ""

    if n_candidates == 1:
        clean: str = ""
        last_err: Optional[str] = None
        for attempt in range(max(1, retries + 1)):
            comp = None
            if attempt == 0:
                comp = cache.get(str(getattr(args, "model")), inp, params)
                if comp is None:
                    def _gen() -> str:
                        return try_generate_python_body(llm, inp, params, grammar_constrained)
                    import time as _time
                    _t0 = _time.time()
                    text = retry_with_backoff(_gen, retries, retry_backoff) or ""
                    _latency_ms = int((_time.time() - _t0) * 1000)
                    comp = text
            clean_try = sanitize_fn(strip_fences(str(comp or "")), prompt)
            if bool(getattr(args, "precheck", False)):
                ok, err_msg = precheck_code(clean_try)
                if not ok:
                    last_err = err_msg or "precheck failed"
                    if attempt < retries:
                        continue
            if not looks_like_python_code(clean_try):
                last_err = "rejected_noncode"
                if attempt < retries:
                    continue
            if critic and clean_try:
                ok, _ = critic.review(clean_try)
                if not ok:
                    last_err = "critic rejected"
                    if attempt < retries:
                        continue
            if not clean_try.strip() and attempt < retries:
                last_err = "empty after sanitize"
                continue
            clean = clean_try
            if attempt == 0 and comp and clean.strip():
                cache.put(str(getattr(args, "model")), inp, params, str(comp))
            break
        # Optional self-eval/refinement loop for single-candidate path
        if bool(getattr(args, "pav_enable", False)) and clean.strip():
            try:
                max_iters = max(1, int(getattr(args, "pav_max_iters", 2) or 2))
                for _iter in range(max_iters):
                    # Evaluate signals
                    ch_pre_ok = False
                    ch_crit_ok = False
                    ch_v_risk = 0.5
                    try:
                        ch_pre_ok, _ = precheck_code(clean)
                    except Exception:
                        ch_pre_ok = False
                    if critic and clean:
                        try:
                            ch_crit_ok, _ = critic.review(clean)
                        except Exception:
                            ch_crit_ok = False
                    if verify_trace_snippet is not None and clean:
                        try:
                            vres = verify_trace_snippet(
                                trace={"obs": {"kind": "code", "content": prompt}, "plan": {}, "action": {"completion": clean}, "result": {"status": "pending"}},
                                model=str(getattr(args, "verifier_model", "meta-llama/Llama-3.2-3B")),
                                backend=str(getattr(args, "verifier_backend", "hf")),
                                use_llm=True,
                                timeout=_safe_int(getattr(args, "timeout", 15) or 15, 15),
                            )
                            ch_v_risk = float(vres.get("risk", 0.5))
                        except Exception:
                            ch_v_risk = 0.5
                    need_refine = (not ch_pre_ok) or (critic is not None and not ch_crit_ok) or (ch_v_risk > float(getattr(args, "verify_risk_threshold", 0.6) or 0.6)) or bool(getattr(args, "self_eval", False))
                    if not need_refine:
                        break
                    review_system = "You are a Python coding assistant. Improve the function body minimally to pass the tests. Output ONLY the corrected function body."
                    review_user = (
                        "Task prompt (do not rewrite def):\n" + prompt
                        + "\n\nCurrent attempt body:\n" + clean
                        + "\n\nInstruction: Return a corrected function body only."
                    )
                    try:
                        messages = [{"role": "system", "content": review_system}, {"role": "user", "content": review_user}]
                        text = str(llm.chat(messages, max_new_tokens=params.get("max_new_tokens", 256), temperature=params.get("temperature", 0.2), top_p=params.get("top_p"), top_k=params.get("top_k")))
                    except Exception:
                        text = str(llm.generate(review_user, **params))
                    rev = sanitize_fn(strip_fences(text), prompt)
                    if looks_like_python_code(rev):
                        rev_pre_ok = False
                        try:
                            rev_pre_ok, _ = precheck_code(rev)
                        except Exception:
                            rev_pre_ok = False
                        rev_v_risk = ch_v_risk
                        if verify_trace_snippet is not None:
                            try:
                                vres2 = verify_trace_snippet(
                                    trace={"obs": {"kind": "code", "content": prompt}, "plan": {}, "action": {"completion": rev}, "result": {"status": "pending"}},
                                    model=str(getattr(args, "verifier_model", "meta-llama/Llama-3.2-3B")),
                                    backend=str(getattr(args, "verifier_backend", "hf")),
                                    use_llm=True,
                                    timeout=_safe_int(getattr(args, "timeout", 15) or 15, 15),
                                )
                                rev_v_risk = float(vres2.get("risk", ch_v_risk))
                            except Exception:
                                rev_v_risk = ch_v_risk
                        if rev_pre_ok or (rev_v_risk <= ch_v_risk):
                            clean = rev
                            write_trace(trace_path, suite, task_id, "self_eval", None, (int(getattr(args, "seed")) if getattr(args, "seed", None) is not None else None), env_fp, inp, clean, meta={"rev_pre_ok": bool(rev_pre_ok), "rev_v_risk": float(rev_v_risk)})
                        else:
                            break
            except Exception:
                pass
        write_trace(trace_path, suite, task_id, 0, int(locals().get("_latency_ms", 0)), (int(getattr(args, "seed")) if getattr(args, "seed", None) is not None else None), env_fp, inp, clean)
        chosen = clean
        lines.append({"task_id": task_id, "completion": clean, "orig_prompt": prompt, "input": inp, "role": "chosen", "nbest_rank": 0, "latency_ms": int(locals().get("_latency_ms", 0)), "attempt_idx": 0})
        return lines, chosen

    # n-best path
    for ci in range(n_candidates):
        comp = None
        cand_latency_ms: Optional[int] = None
        if ci == 0:
            comp = cache.get(str(getattr(args, "model")), inp, params)
        if comp is None:
            def _gen_multi() -> str:
                return try_generate_python_body(llm, inp, params, grammar_constrained)
            import time as _time
            _t0 = _time.time()
            text = retry_with_backoff(_gen_multi, retries, retry_backoff) or ""
            cand_latency_ms = int((_time.time() - _t0) * 1000)
            comp = text
            if ci == 0 and comp:
                cache.put(str(getattr(args, "model")), inp, params, str(comp))
        cand = sanitize_fn(strip_fences(str(comp or "")), prompt)
        pre_ok, crit_ok, v_risk = evaluate_candidate_signals(cand, prompt, precheck_code, critic, verify_trace_snippet, args)
        candidates.append({"body": cand, "precheck": pre_ok, "critic": crit_ok, "length": len(cand or ""), "v_risk": v_risk, "latency_ms": cand_latency_ms})
    from agi_dw.bench.common.pipeline import pick_best_candidate  # local import to avoid circulars
    chosen = pick_best_candidate(candidates, float(getattr(args, "verify_risk_threshold", 0.6) or 0.6))
    from agi_dw.bench.common.pipeline import rank_candidates  # type: ignore
    _ranks = rank_candidates(candidates)
    # Adaptive extra sampling if chosen looks risky/low-quality
    try:
        ch_pre_ok = False
        ch_crit_ok = False
        ch_v_risk = 0.5
        if chosen:
            try:
                ch_pre_ok, _ = precheck_code(chosen)
            except Exception:
                ch_pre_ok = False
            if critic:
                try:
                    ch_crit_ok, _ = critic.review(chosen)
                except Exception:
                    ch_crit_ok = False
            if verify_trace_snippet is not None:
                try:
                    vres = verify_trace_snippet(
                        trace={"obs": {"kind": "code", "content": prompt}, "plan": {}, "action": {"completion": chosen}, "result": {"status": "pending"}},
                        model=str(getattr(args, "verifier_model", "meta-llama/Llama-3.2-3B")),
                        backend=str(getattr(args, "verifier_backend", "hf")),
                        use_llm=True,
                        timeout=_safe_int(getattr(args, "timeout", 15) or 15, 15),
                    )
                    ch_v_risk = float(vres.get("risk", 0.5))
                except Exception:
                    ch_v_risk = 0.5
        bad = (not looks_like_python_code(chosen)) or (not ch_pre_ok) or (critic is not None and not ch_crit_ok) or (ch_v_risk > float(getattr(args, "verify_risk_threshold", 0.6) or 0.6))
        if bad:
            M = max(1, _safe_int(getattr(args, "adaptive_extra_candidates", 2) or 2, 2))
            for _ in range(M):
                text = retry_with_backoff(lambda: try_generate_python_body(llm, inp, params, grammar_constrained), retries, retry_backoff) or ""
                cand_try = sanitize_fn(strip_fences(text), prompt)
                if cand_try.strip():
                    try:
                        pre_ok, _ = precheck_code(cand_try)
                    except Exception:
                        pre_ok = False
                    crit_ok = False
                    if critic and cand_try:
                        try:
                            crit_ok, _ = critic.review(cand_try)
                        except Exception:
                            crit_ok = False
                    v_risk = 0.5
                    if verify_trace_snippet is not None:
                        try:
                            vres3 = verify_trace_snippet(
                                trace={"obs": {"kind": "code", "content": prompt}, "plan": {}, "action": {"completion": cand_try}, "result": {"status": "pending"}},
                                model=str(getattr(args, "verifier_model", "meta-llama/Llama-3.2-3B")),
                                backend=str(getattr(args, "verifier_backend", "hf")),
                                use_llm=True,
                                timeout=_safe_int(getattr(args, "timeout", 15) or 15, 15),
                            )
                            v_risk = float(vres3.get("risk", 0.5))
                        except Exception:
                            v_risk = 0.5
                    candidates.append({"body": cand_try, "precheck": bool(pre_ok and looks_like_python_code(cand_try)), "critic": crit_ok, "length": len(cand_try or ""), "v_risk": v_risk, "latency_ms": None})
            from agi_dw.bench.common.pipeline import pick_best_candidate as _pick_best
            chosen = _pick_best(candidates, float(getattr(args, "verify_risk_threshold", 0.6) or 0.6))
    except Exception:
        pass

    # Optional WM re-ranking
    wm_model = None
    wm_enabled = bool(getattr(args, "wm_prior", False))
    if wm_enabled:
        try:
            from agi_dw.core.world_model.api import WorldModelPrior  # type: ignore
        except Exception:
            WorldModelPrior = None  # type: ignore
        try:
            wm_path = Path(getattr(args, "wm_model", Path(Path(__file__).resolve().parents[2]) / "models" / "wm_mlp" / "wm_mlp.joblib"))
            if 'WorldModelPrior' in globals() and WorldModelPrior is not None and wm_path.exists():
                wm_model = WorldModelPrior.load(wm_path)
        except Exception:
            wm_model = None
        if wm_model is not None:
            try:
                chosen = wm_rerank_candidates(wm_model, prompt, candidates, chosen, float(getattr(args, "wm_threshold", 0.6) or 0.6))
            except Exception:
                pass

    # Optional self-evaluation/refinement loop before committing chosen
    if bool(getattr(args, "pav_enable", False)):
        try:
            max_iters = max(1, _safe_int(getattr(args, "pav_max_iters", 2) or 2, 2))
            for _iter in range(max_iters):
                ch_pre_ok = False
                ch_crit_ok = False
                ch_v_risk = 0.5
                try:
                    ch_pre_ok, _ = precheck_code(chosen)
                except Exception:
                    ch_pre_ok = False
                if critic and chosen:
                    try:
                        ch_crit_ok, _ = critic.review(chosen)
                    except Exception:
                        ch_crit_ok = False
                if verify_trace_snippet is not None and chosen:
                    try:
                        vres = verify_trace_snippet(
                            trace={"obs": {"kind": "code", "content": prompt}, "plan": {}, "action": {"completion": chosen}, "result": {"status": "pending"}},
                            model=str(getattr(args, "verifier_model", "meta-llama/Llama-3.2-3B")),
                            backend=str(getattr(args, "verifier_backend", "hf")),
                            use_llm=True,
                            timeout=_safe_int(getattr(args, "timeout", 15) or 15, 15),
                        )
                        ch_v_risk = float(vres.get("risk", 0.5))
                    except Exception:
                        ch_v_risk = 0.5
                need_refine = (not ch_pre_ok) or (critic is not None and not ch_crit_ok) or (ch_v_risk > float(getattr(args, "verify_risk_threshold", 0.6) or 0.6)) or bool(getattr(args, "self_eval", False))
                if not need_refine:
                    break
                review_system = "You are a Python coding assistant. Improve the function body minimally to pass the tests. Output ONLY the corrected function body."
                review_user = (
                    "Task prompt (do not rewrite def):\n" + prompt
                    + "\n\nCurrent attempt body:\n" + chosen
                    + "\n\nInstruction: Return a corrected function body only."
                )
                try:
                    messages = [{"role": "system", "content": review_system}, {"role": "user", "content": review_user}]
                    text = str(llm.chat(messages, max_new_tokens=params.get("max_new_tokens", 256), temperature=params.get("temperature", 0.2), top_p=params.get("top_p"), top_k=params.get("top_k")))
                except Exception:
                    text = str(llm.generate(review_user, **params))
                rev = sanitize_fn(strip_fences(text), prompt)
                if looks_like_python_code(rev):
                    rev_pre_ok = False
                    try:
                        rev_pre_ok, _ = precheck_code(rev)
                    except Exception:
                        rev_pre_ok = False
                    rev_v_risk = ch_v_risk
                    if verify_trace_snippet is not None:
                        try:
                            vres2 = verify_trace_snippet(
                                trace={"obs": {"kind": "code", "content": prompt}, "plan": {}, "action": {"completion": rev}, "result": {"status": "pending"}},
                                model=str(getattr(args, "verifier_model", "meta-llama/Llama-3.2-3B")),
                                backend=str(getattr(args, "verifier_backend", "hf")),
                                use_llm=True,
                                timeout=_safe_int(getattr(args, "timeout", 15) or 15, 15),
                            )
                            rev_v_risk = float(vres2.get("risk", ch_v_risk))
                        except Exception:
                            rev_v_risk = ch_v_risk
                    if rev_pre_ok or (rev_v_risk <= ch_v_risk):
                        chosen = rev
                        write_trace(trace_path, suite, task_id, "self_eval", None, (int(getattr(args, "seed")) if getattr(args, "seed", None) is not None else None), env_fp, inp, chosen, meta={"rev_pre_ok": bool(rev_pre_ok), "rev_v_risk": float(rev_v_risk)})
                    else:
                        break
        except Exception:
            pass

    # Logger
    if logger:
        try:
            logger.log_text("prompt", inp)
            logger.log_text("completion", chosen)
        except Exception:
            pass

    # Determine chosen latency if available
    chosen_latency = None
    for c in candidates:
        if c.get("body", "") == chosen:
            chosen_latency = c.get("latency_ms")
            break

    write_trace(trace_path, suite, task_id, 0, (int(chosen_latency) if chosen_latency is not None else None), (int(getattr(args, "seed")) if getattr(args, "seed", None) is not None else None), env_fp, inp, chosen)
    lines.append({"task_id": task_id, "completion": chosen, "orig_prompt": prompt, "input": inp, "role": "chosen", "nbest_rank": 0, "latency_ms": (int(chosen_latency) if chosen_latency is not None else None), "attempt_idx": 0})
    for ri, cand in enumerate(_ranks):
        body = str(cand.get("body", ""))
        if body == chosen:
            continue
        write_trace(trace_path, suite, task_id, (ri + 1), (int(cand.get("latency_ms")) if cand.get("latency_ms") is not None else None), (int(getattr(args, "seed")) if getattr(args, "seed", None) is not None else None), env_fp, inp, body)
        lines.append({"task_id": task_id, "completion": body, "orig_prompt": prompt, "input": inp, "role": "alt", "nbest_rank": (ri if ri < len(_ranks) else 9999), "latency_ms": (int(cand.get("latency_ms")) i
# ... truncated ...


    # Determine chosen latency if available
    chosen_latency = None
    for c in candidates:
        if c.get("body", "") == chosen:
            chosen_latency = c.get("latency_ms")
            break

    write_trace(trace_path, suite, task_id, 0, (int(chosen_latency) if chosen_latency is not None else None), (int(getattr(args, "seed")) if getattr(args, "seed", None) is not None else None), env_fp, inp, chosen)
    lines.append({"task_id": task_id, "completion": chosen, "orig_prompt": prompt, "input": inp, "role": "chosen", "nbest_rank": 0, "latency_ms": (int(chosen_latency) if chosen_latency is not None else None), "attempt_idx": 0})
    for ri, cand in enumerate(_ranks):
        body = str(cand.get("body", ""))
        if body == chosen:
            continue
        write_trace(trace_path, suite, task_id, (ri + 1), (int(cand.get("latency_ms")) if cand.get("latency_ms") is not None else None), (int(getattr(args, "seed")) if getattr(args, "seed", None) is not None else None), env_fp, inp, body)
        lines.append({"task_id": task_id, "completion": body, "orig_prompt": prompt, "input": inp, "role": "alt", "nbest_rank": (ri if ri < len(_ranks) else 9999), "latency_ms": (int(cand.get("latency_ms")) if cand.get("latency_ms") is not None else None), "attempt_idx": (ri + 1)})
    return lines, chosen

def run_codebody_suite(
    args: Any,
    suite: str,
    problems: Dict[str, Dict[str, Any]],
    prompt_builder_key: str,
    sanitize_fn: Callable[[str, str], str],
    evaluate_functional_correctness: Any,
    evalplus_evaluate: Optional[Any] = None,
) -> int:
    """Generic end-to-end runner for code-body benchmarks.

    Parameters
    - args: argparse-like object with benchmark options
    - suite: suite name (e.g., "humaneval", "mbpp")
    - problems: mapping of task_id -> {"prompt": str, ...}
    - prompt_builder_key: key for build_prompt (e.g., "humaneval", "mbpp")
    - sanitize_fn: function(text, prompt) -> cleaned body
    - evaluate_functional_correctness: evaluator callable for this suite
    - evalplus_evaluate: optional evalplus alternative callable
    """
    from agi_dw.core.prompts.bench import build_prompt  # type: ignore
    from agi_dw.core.utils.bench_utils import ensure_safe_env, inject_similar_functions  # type: ignore

    root = Path(__file__).resolve().parents[2]
    run_started_ts = time.time()
    traces_dir, trace_path, ts_str = init_trace_paths(root, suite, run_started_ts)
    env_fp = default_env_fingerprint()
    default_workers = DEFAULT_WORKERS
    ensure_safe_env()

    # Seed and tasks
    try:
        seed_everything(getattr(args, "seed", None))
    except Exception:
        pass
    task_ids: List[str] = sorted(list(problems.keys()))
    try:
        if int(getattr(args, "limit", 0) or 0) > 0:
            task_ids = task_ids[: int(getattr(args, "limit"))]
    except Exception:
        pass
    task_ids = shard_task_ids(task_ids, args)

    # LLM and helpers
    llm, cache, logger, critic = get_llm_and_basics(args, root, suite, default_adapter=None)
    try:
        _crit2 = critic
        from agi_dw.bench.common.pipeline import build_critic_from_args as _build_crit  # type: ignore
        _crit2 = _build_crit(args)
        if _crit2 is not None:
            critic = _crit2
    except Exception:
        pass

    # Code memory & index
    try:
        from agi_dw.bench.common.pipeline import load_code_memory_from_args as _load_mem, resolve_code_index as _resolve_ci  # type: ignore
    except Exception:
        _load_mem = None  # type: ignore
        _resolve_ci = None  # type: ignore
    code_memory = (_load_mem(args, root) if _load_mem is not None else None)
    code_index: Optional[Dict[str, Any]] = None
    index_k: int = 0
    index_path: Optional[str] = None
    if _resolve_ci is not None:
        try:
            code_index, index_k, index_path = _resolve_ci(args, root)
        except Exception:
            code_index, index_k, index_path = None, 0, None

    # Pass cache and temp paths
    pass_cache, pass_cache_path = read_pass_cache(root, suite)
    try:
        shard_tag = get_shard_tag(args)
    except Exception:
        shard_tag = ""
    samples_path, verbose_path, errors_path = init_suite_temp_paths(root, suite, shard_tag)

    # Prompt checksum
    prompts_hasher = hashlib.sha256()
    def _build_input(base: str) -> str:
        try:
            return str(build_prompt(prompt_builder_key, base))
        except Exception:
            return base

    def _generate_for_task(tid: str) -> List[Dict[str, Any]]:
        prompt = str(problems[tid]["prompt"]) if isinstance(problems.get(tid), dict) else str(problems[tid])
        base = prompt
        if code_memory is not None:
            try:
                base = code_memory.enhance_prompt(base, problem_id=tid)
            except Exception:
                pass
        elif code_index and index_k > 0:
            try:
                funcs_map = dict(code_index.get("functions", {})) if isinstance(code_index, dict) else {}
                n_funcs = sum(len(v) for v in funcs_map.values()) if funcs_map else 0
                if n_funcs >= int(getattr(args, "planner_index_min_funcs", 50) or 50):
                    base = inject_similar_functions(base, prompt, code_index, index_k)
            except Exception:
                pass
        inp = _build_input(base)
        try:
            from agi_dw.bench.common.pipeline import maybe_short_circuit_from_pass_cache as _shortc, annotate_prompt_with_cached_solution as _anno  # type: ignore
        except Exception:
            _shortc = None  # type: ignore
            _anno = None  # type: ignore
        if _shortc is not None:
            sc = _shortc(suite, tid, inp, pass_cache, trace_path, env_fp, args)
            if sc is not None:
                return sc
        cached = pass_cache.get(tid)
        if cached and str(cached.get("completion", "")).strip() and _anno is not None:
            try:
                base = _anno(base, str(cached.get("completion", "")))
            except Exception:
                pass
        inp = _build_input(base)
        try:
            prompts_hasher.update(inp.encode("utf-8", errors="ignore"))
        except Exception:
            pass
        lines, _chosen = generate_candidates_for_prompt(
            llm=llm,
            cache=cache,
            logger=logger,
            critic=critic,
            prompt=prompt,
            inp=inp,
            sanitize_fn=sanitize_fn,
            task_id=tid,
            args=args,
            suite=suite,
            env_fp=env_fp,
            trace_path=trace_path,
            _prompts_hasher=prompts_hasher,
        )
        return lines

    # Parallel generation
    max_workers = max(1, int(getattr(args, "max_workers", default_workers) or default_workers))
    all_rows: List[Dict[str, Any]] = []
    from concurrent.futures import ThreadPoolExecutor as _TPE, as_completed as _as_completed  # type: ignore
    with trace_span("generate_candidates", {"suite": suite, "n_tasks": len(task_ids), "max_workers": max_workers}):
        with _TPE(max_workers=max_workers) as ex:
            futs = {ex.submit(_generate_for_task, tid): tid for tid in task_ids}
            for fut in _as_completed(futs):
                _tid = futs[fut]
                rows: List[Dict[str, Any]] = []
                try:
                    rows = fut.result()
                except Exception as e:
                    try:
                        import traceback, sys  # type: ignore
                        trace = traceback.format_exc()
                        msg = {"task_id": _tid, "error": str(e), "traceback": trace}
                        with errors_path.open("a", encoding="utf-8") as ef:
                            ef.write(__import__("json").dumps(msg) + "\n")
                        print(__import__("json").dumps({"task_id": _tid, "error": str(e)}), file=sys.stderr)
                    except Exception:
                        pass
                    rows = [{"task_id": _tid, "completion": "", "orig_prompt": str(problems.get(_tid, {}).get("prompt", "")), "input": "", "role": "chosen", "nbest_rank": 0}]
                all_rows.extend(rows)

    # Write samples and sidecar
    write_samples_and_verbose(all_rows, samples_path, verbose_path)

    # Evaluate
    try:
        with trace_span("evaluate", {"suite": suite}):
            k_list = parse_k_list(args)
            use_evalplus = bool(getattr(args, "use_evalplus", False)) and (evalplus_evaluate is not None)
            if use_evalplus:
                res, results_path = evaluate_samples(samples_path, k_list, default_workers, args, evalplus_evaluate)  # type: ignore
            else:
                res, results_path = evaluate_samples(samples_path, k_list, default_workers, args, evaluate_functional_correctness)
    except Exception as e:
        print(__import__("json").dumps({"ok": False, "error": f"evaluation failed: {e}"}))
        return 2

    # Optional online update and rerun
    if bool(getattr(args, "online_update", True)) and is_primary_process():
        with trace_span("online_update", {"suite": suite}):
            try:
                _res2 = online_update_and_rerun_failures(
                    args=args,
                    root=root,
                    suite=suite,
                    default_workers=default_workers,
                    samples_path=samples_path,
                    verbose_path=verbose_path,
                    results_path=results_path,
                    k_list=k_list,
                    regen_failed_fn=_generate_for_task,
                    evalplus_evaluate=evalplus_evaluate,
                    evaluate_functional_correctness=evaluate_functional_correctness,
                )
                if _res2 is not None:
                    res = _res2
            except Exception:
                pass

    # Build per-task results
    results_path = Path(str(samples_path) + "_results.jsonl")
    results_by_task: Dict[str, bool] = _build_results_by_task_shared(results_path)
    # Code memory store (skip in benchmark_mode by default)
    try:
        if not bool(getattr(args, "benchmark_mode", True)):
            store_passed_solutions_in_memory(code_memory, results_path, problems, suite)
    except Exception:
        pass

    # Optional repair loop
    if bool(getattr(args, "repair", False)):
        with trace_span("repair_loop", {"suite": suite}):
            failed_rows: List[Dict[str, str]] = []
            for obj in _read_jsonl_shared(verbose_path):
                if not results_by_task.get(str(obj.get("task_id", "")), False):
                    failed_rows.append(obj)
            try:
                meter_cost("repair_failures", float(len(failed_rows)))
            except Exception:
                pass
            try:
                from agi_dw.core.utils.bench_utils import strip_fences as _strip_f  # type: ignore
            except Exception:
                _strip_f = lambda x: x  # type: ignore
            repair_code_failures(
                suite=suite,
                failed_rows=failed_rows,
                args=args,
                llm=llm,
                sanitize_completion=sanitize_fn,
                strip_fences_fn=_strip_f,
                samples_path=samples_path,
                verbose_path=verbose_path,
                root=root,
            )
            try:
                k_list = parse_k_list(args)
                use_evalplus = bool(getattr(args, "use_evalplus", False)) and (evalplus_evaluate is not None)
                if use_evalplus:
                    res, _ = evaluate_samples(samples_path, k_list, default_workers, args, evalplus_evaluate)  # type: ignore
                else:
                    res, _ = evaluate_samples(samples_path, k_list, default_workers, args, evaluate_functional_correctness)
            except Exception:
                pass

    # Update pass cache (skip in benchmark_mode by default)
    try:
        if not bool(getattr(args, "benchmark_mode", True)):
            latest_passes: Dict[str, Dict[str, str]] = {}
            for row in _read_jsonl_shared(results_path):
                if bool(row.get("passed", False)):
                    tid = str(row.get("task_id", ""))
                    comp = str(row.get("completion", ""))
                    if tid and comp:
                        latest_passes[tid] = {"completion": comp, "prompt": str(problems.get(tid, {}).get("prompt", ""))}
            if latest_passes:
                try:
                    update_pass_cache(pass_cache_path, pass_cache, latest_passes)
                except Exception:
                    pass
    except Exception:
        pass

    # Telemetry and dedupe
    results_by_task = {}
    for row in _read_jsonl_shared(results_path):
        tid = str(row.get("task_id", ""))
        passed = bool(row.get("passed", False))
        if tid:
            results_by_task[tid] = bool(results_by_task.get(tid, False) or passed)
    # Emit indirect learning dataset by default (trace-only, no labels)
    try:
        from agi_dw.bench.common.pipeline import emit_indirect_trace_datasets  # type: ignore
    except Exception:
        emit_indirect_trace_datasets = None  # type: ignore
    try:
        if emit_indirect_trace_datasets is not None and bool(getattr(args, "emit_indirect", True)):
            emit_indirect_trace_datasets(root, suite, trace_path, results_by_task)
    except Exception:
        pass
    telemetry, pass1_rate = summarize_telemetry(task_ids, results_by_task)
    try:
        _dedupe_by_passed_shared(samples_path, verbose_path, results_path)
    except Exception:
        pass
    try:
        flake_detection(samples_path, args, default_workers, task_ids, results_by_task, suite, root)
    except Exception:
        pass

    # Write per-suite JSONL and run artifact
    with trace_span("write_artifacts", {"suite": suite}):
        outp = write_suite_out(str(getattr(args, "out")), args, suite, task_ids, results_by_task)
        try:
            args_summary = build_codebody_args_summary(args)
            paths_dict = {"samples": str(samples_path), "sidecar": str(verbose_path), "results": str(results_path), "out": str(outp)}
            metrics_dict = {"pass1_rate": pass1_rate, **(telemetry if isinstance(telemetry, dict) else {}), **({"pass_at_k": res} if isinstance(res, dict) else {})}
            write_run_artifact(
                root=root,
                suite=suite,
                run_started_ts=run_started_ts,
                env=env_fp,
                prompts_hasher=prompts_hasher,
                args_summary=args_summary,
                paths=paths_dict,
                aggregate=res,
                metrics=metrics_dict,
                model=str(getattr(args, "model", "")),
                adapter_dir=str(getattr(args, "adapter_dir", "") or ""),
            )
        except Exception:
            pass

    print(__import__("json").dumps({
        "ok": True,
        "out": str(outp),
        "n": len(task_ids),
        **({"aggregate": res} if isinstance(res, dict) else {}),
        **({"telemetry": telemetry} if telemetry else {}),
        **({"pass1_rate": pass1_rate} if pass1_rate is not None else {}),
        **({"pass_at_k": res} if isinstance(res, dict) else {}),
    }))
    try:
        if not bool(getattr(args, "benchmark_mode", True)):
            emit_training_datasets(root, suite, verbose_path, results_path)
    except Exception:
        pass
    return 0