from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Awaitable, Callable, Dict, List, Optional
import json
import os
import re
import subprocess
import sys
import time

import questionary
from openai import AsyncOpenAI
from pydantic import BaseModel

from agents import Agent, OpenAIChatCompletionsModel, Runner, SQLiteSession, function_tool, set_tracing_disabled

from armorer.agent_sdk import attach_tool_metadata, get_tool_metadata, tool, tool_name
from armorer.config import resolve_armorer_config_dir
from armorer.profiles import NetworkPolicy
from armorer.security.policy import PolicyGuard, get_policy_guard, policy_scope
from armorer.training_data import record_training_trace
from armorer.tools.interaction import ask_selection
from armorer.tools.knowledge import get_knowledge_topic, list_knowledge_topics, search_knowledge
from armorer.tools.search import google_search, is_available as is_search_available

MAX_CHARS = 20000


@tool(category="system_io", verification_level="runtime_command")
def shell(command: str, work_dir: str = None, timeout: Optional[int] = None) -> str:
    """Execute a shell command on the system."""
    cwd = Path.cwd().resolve()
    installed_dir = cwd / "installed"
    guard = get_policy_guard()

    if "/tmp/autogpt-deployment" in command:
        return (
            "Error: You are attempting to operate in a temporary directory outside the project. "
            "Armorer MUST operate within the 'installed/' directory."
        )

    if work_dir:
        wd_path = Path(work_dir).resolve()
        if not str(wd_path).startswith(str(installed_dir)) and not str(wd_path).startswith(str(cwd)):
            pass

    if "/home/user" in command or "/Users/" in command:
        return (
            "Error: You are using an absolute path that does not exist here. "
            "Use relative paths inside the project root."
        )

    try:
        run_dir = work_dir or str(cwd)
        cmd_decision = guard.check_command(command, cwd=Path(run_dir))
        if not cmd_decision.allowed:
            return f"Error: Security Policy Violation. {cmd_decision.reason}"

        path_decision = guard.check_path(Path(run_dir), op="exec")
        if not path_decision.allowed:
            return f"Error: Security Policy Violation. {path_decision.reason}"

        run_kwargs = {
            "shell": True,
            "cwd": run_dir,
            "capture_output": True,
            "text": True,
        }
        if timeout is not None and timeout > 0:
            run_kwargs["timeout"] = timeout
        result = subprocess.run(command, **run_kwargs)
        output = result.stdout
        if result.stderr:
            output += f"\nSTDERR: {result.stderr}"
        output = guard.sanitize_text(output)
        if len(output) > MAX_CHARS:
            output = output[:MAX_CHARS] + f"\n... [Output truncated. Total length {len(output)} > {MAX_CHARS} chars.]"
        return output or f"(Command exited with code {result.returncode})"
    except Exception as e:
        return f"Error executing shell: {str(e)}"


@tool(category="system_io", verification_level="file_write")
def file_write(path: str, content: str) -> str:
    """Write content to an allowed file path."""
    try:
        guard = get_policy_guard()
        p = Path(path).resolve()
        cwd = Path.cwd().resolve()
        installed_dir = cwd / "installed"
        armorer_config = resolve_armorer_config_dir()

        is_safe = False
        if str(p).startswith(str(installed_dir)):
            is_safe = True
        elif str(p).startswith(str(armorer_config)):
            is_safe = True
        elif p.parent == cwd and p.name in ["fix_compose.py", "autogpt_deployment_plan.md", "autogpt_deployment_recommendations.md"]:
            is_safe = True

        if not is_safe:
            return f"Error: Security Policy Violation. Write access restricted to '{installed_dir}' or '{armorer_config}'. Attempted: {p}"
        path_decision = guard.check_path(p, op="write")
        if not path_decision.allowed:
            return f"Error: Security Policy Violation. {path_decision.reason}"

        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(guard.sanitize_text(content))
        return f"Successfully wrote to {path}"
    except Exception as e:
        return f"Error writing file: {str(e)}"


@tool(category="system_io", verification_level="file_read", safe_for_retrieval=True)
def file_read(path: str) -> str:
    """Read file contents from allowed project locations."""
    try:
        guard = get_policy_guard()
        p = Path(path).resolve()
        cwd = Path.cwd().resolve()
        if not str(p).startswith(str(cwd)):
            armorer_config = resolve_armorer_config_dir()
            if not str(p).startswith(str(armorer_config)):
                return f"Error: Security Policy Violation. Read access restricted to project root. Attempted: {p}"
        path_decision = guard.check_path(p, op="read")
        if not path_decision.allowed:
            return f"Error: Security Policy Violation. {path_decision.reason}"

        if not p.exists():
            return f"Error: File {path} does not exist."

        content = p.read_text()
        content = guard.sanitize_text(content)
        if len(content) > MAX_CHARS:
            content = content[:MAX_CHARS] + f"\n... [Output truncated. Total length {len(content)} > {MAX_CHARS} chars.]"
        return content
    except Exception as e:
        return f"Error reading file: {str(e)}"


def _extract_json(text: str) -> Optional[dict]:
    if not text:
        return None
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        match = re.search(r"```json\s*(.*?)\s*```", text, re.DOTALL)
        if match:
            try:
                return json.loads(match.group(1))
            except json.JSONDecodeError:
                pass
        start = text.find("{")
        end = text.rfind("}")
        if start != -1 and end != -1:
            try:
                return json.loads(text[start : end + 1])
            except json.JSONDecodeError:
                pass
    return None


_BACKGROUND_JOB_TOOLS = {
    "start_install_agent_job",
    "start_run_agent_job",
    "start_uninstall_agent_job",
}


def _coerce_mapping(value: Any) -> Optional[dict]:
    if isinstance(value, dict):
        return value
    if isinstance(value, str):
        return _extract_json(value)
    return None


def _extract_response_payload(res: Any) -> dict:
    final = getattr(res, "final_output", None)
    if isinstance(final, dict):
        payload = dict(final)
        payload["message"] = str(payload.get("message", "") or "").strip()
        return payload

    text = str(final or "").strip()
    if text:
        parsed = _coerce_mapping(text)
        if isinstance(parsed, dict):
            payload = dict(parsed)
            payload["message"] = str(payload.get("message", "") or "").strip()
            return payload
        return {"message": text}

    fallback = _extract_text_from_run_result(res)
    return {"message": fallback}


def _extract_text_from_run_result(res: Any) -> str:
    final = getattr(res, "final_output", None)
    text = str(final or "").strip()
    if text:
        parsed = _extract_json(text)
        if isinstance(parsed, dict):
            nested = str(parsed.get("message", "") or "").strip()
            if nested:
                return nested
        else:
            return text

    # Fallback for runs where final_output is unexpectedly empty but the
    # underlying raw responses still contain assistant text.
    for response in reversed(getattr(res, "raw_responses", []) or []):
        for item in getattr(response, "output", []) or []:
            if getattr(item, "type", None) != "message":
                continue
            parts = []
            for content in getattr(item, "content", []) or []:
                part = getattr(content, "text", None)
                if part:
                    parts.append(str(part))
            combined = "".join(parts).strip()
            if combined:
                return combined

    return ""


def _input_to_history(user_input: str | list[dict[str, Any]]) -> list[dict[str, Any]]:
    if isinstance(user_input, str):
        return [{"role": "user", "content": user_input}]
    return list(user_input)


def _latest_user_text(user_input: str | list[dict[str, Any]]) -> str:
    if isinstance(user_input, str):
        return user_input
    for message in reversed(user_input):
        if str(message.get("role", "")).strip().lower() != "user":
            continue
        content = message.get("content", "")
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            parts = []
            for item in content:
                if isinstance(item, dict):
                    text = item.get("text")
                    if text:
                        parts.append(str(text))
            return "".join(parts)
    return ""


def _is_lightweight_conversation(user_input: str | list[dict[str, Any]]) -> bool:
    text = str(_latest_user_text(user_input) or "").strip().lower()
    if not text:
        return False
    simple = {
        "hi", "hello", "hey", "yo", "sup",
        "thanks", "thank you", "ok", "okay",
        "good morning", "good afternoon", "good evening",
        "how are you", "who are you",
    }
    if text in simple:
        return True
    if len(text.split()) <= 4 and re.fullmatch(r"[a-zA-Z ?!'.]+", text):
        return text.endswith("?") or text in simple
    return False


def _build_self_retry_input(res: Any, attempt: int) -> list[dict[str, Any]]:
    history = res.to_input_list()
    if attempt == 1:
        correction = (
            "Your previous attempt did not produce a usable final answer. "
            "Failure: empty or malformed final message. "
            "Answer the user's actual question directly now. "
            "Use tools if needed. Do not describe your internal tools, orchestration, or capabilities unless the user explicitly asked about them. "
            "Do not return JSON unless explicitly asked."
        )
    elif attempt == 2:
        correction = (
            "Your previous retry still did not produce a usable final answer. "
            "This is a correction pass. "
            "If the question depends on system state, inspect it with tools and answer from observed results. "
            "Do not answer from your role description, available tool names, or generic capabilities. "
            "Respond concisely and directly to the user's request."
        )
    else:
        correction = (
            "Final correction attempt. "
            "You must provide a direct user-facing answer now. "
            "If needed, call the relevant tools again and then summarize the observed result in plain language. "
            "Do not output an empty message. "
            "Do not discuss internal orchestration."
        )
    history.append(
        {
            "role": "user",
            "content": correction,
        }
    )
    return history


def _build_tool_error_retry_input(
    user_input: str | list[dict[str, Any]],
    error_message: str,
    attempt: int,
) -> list[dict[str, Any]]:
    history = _input_to_history(user_input)
    extra_guidance = ""
    lowered = str(error_message or "").lower()
    if "tool job_response not found" in lowered or "tool track_job not found" in lowered:
        extra_guidance = (
            " `track_job` is not a callable tool. "
            "It must be returned inside your final JSON as `ui_action`, alongside a non-empty `message`. "
            "Never call `job_response`, `track_job`, `ui_action`, or `message` as tools. "
            "If you already started a background job, stop calling tools and respond with JSON like "
            '{"message":"Installing openclaw...","ui_action":{"type":"track_job","job_id":"<job_id>","message":"Installing openclaw..."}}.'
        )
    history.append(
        {
            "role": "user",
            "content": (
                f"Your previous attempt failed with this tool error: {error_message}. "
                f"Retry attempt {attempt}. "
                "Only use tools that are actually available in this session. "
                "Do not invent tool names. "
                "If this task does not require tools, answer directly. "
                "If it does require tools, choose from the registered tools only."
                f"{extra_guidance}"
            ),
        }
    )
    return history


def _build_lightweight_chat_retry_input(user_input: str | list[dict[str, Any]]) -> list[dict[str, Any]]:
    history = _input_to_history(user_input)
    history.append(
        {
            "role": "user",
            "content": (
                "This is lightweight conversation. "
                "Answer directly and naturally. "
                "Do not inspect system state, do not use tools, and do not describe your capabilities unless asked."
            ),
        }
    )
    return history


def _model_retry_suggestion(provider: str, model_name: str) -> str:
    if (provider or "").lower() == "ollama":
        return (
            f"I retried this request multiple times with `{model_name}` but it still failed to produce a usable final answer. "
            "Try a stronger model in `armorer config`, or switch temporarily to a hosted model for this task."
        )
    return (
        f"I retried this request multiple times with `{model_name}` but it still failed to produce a usable final answer. "
        "Try a different model and rerun the request."
    )


def _item_value(item: Any, key: str) -> Any:
    if isinstance(item, dict):
        return item.get(key)
    return getattr(item, key, None)


def _normalized_tool_item(item: Any) -> Optional[dict[str, Any]]:
    raw = _item_value(item, "raw_item")
    if isinstance(raw, dict):
        return raw
    to_input_item = getattr(item, "to_input_item", None)
    if callable(to_input_item):
        try:
            normalized = to_input_item()
        except Exception:
            normalized = None
        if isinstance(normalized, dict):
            return normalized
    if isinstance(item, dict):
        return item
    return None


def _extract_tool_events(res: Any) -> list[dict[str, Any]]:
    events: list[dict[str, Any]] = []
    call_names: dict[str, str] = {}
    for item in getattr(res, "new_items", []) or []:
        source = _normalized_tool_item(item)
        item_type = str(_item_value(item, "type") or (source or {}).get("type") or "")
        if item_type in {"tool_call_item", "function_call"}:
            name = ""
            call_id = ""
            if isinstance(source, dict):
                name = str(source.get("name") or source.get("tool_name") or "")
                call_id = str(source.get("call_id") or source.get("tool_call_id") or source.get("id") or "")
                fn = source.get("function")
                if not name and isinstance(fn, dict):
                    name = str(fn.get("name") or "")
            if not name:
                name = str(_item_value(item, "name") or _item_value(item, "tool_name") or "")
            if not call_id:
                call_id = str(_item_value(item, "call_id") or _item_value(item, "tool_call_id") or _item_value(item, "id") or "")
            if call_id and name:
                call_names[call_id] = name
            events.append(
                {
                    "type": "tool_call",
                    "tool": name,
                    "call_id": call_id,
                }
            )
        elif item_type in {"tool_call_output_item", "function_call_output"}:
            call_id = str(_item_value(item, "call_id") or _item_value(item, "tool_call_id") or _item_value(item, "id") or "")
            if not call_id and isinstance(source, dict):
                call_id = str(source.get("call_id") or source.get("tool_call_id") or source.get("id") or "")
            output = _item_value(item, "output")
            if output is None and isinstance(source, dict):
                output = source.get("output")
            events.append(
                {
                    "type": "tool_output",
                    "tool": call_names.get(call_id, ""),
                    "call_id": call_id,
                    "output": output,
                }
            )
    return events


def _extract_background_job_action(tool_events: list[dict[str, Any]]) -> Optional[dict]:
    for event in reversed(tool_events):
        if event.get("type") != "tool_output":
            continue
        tool_name = str(event.get("tool", "") or "").strip()
        if tool_name not in _BACKGROUND_JOB_TOOLS:
            continue
        payload = _coerce_mapping(event.get("output"))
        if not isinstance(payload, dict):
            continue
        job_id = str(payload.get("job_id", "") or "").strip()
        if not job_id or payload.get("ok") is False:
            continue
        message = str(payload.get("message", "") or "").strip() or "Background job started."
        return {
            "type": "track_job",
            "job_id": job_id,
            "message": message,
        }
    return None


def _enrich_payload_with_background_job(payload: dict, tool_events: list[dict[str, Any]]) -> dict:
    existing = payload.get("ui_action")
    if isinstance(existing, dict) and str(existing.get("type", "")).strip().lower() == "track_job":
        return payload
    action = _extract_background_job_action(tool_events)
    if not action:
        return payload
    enriched = dict(payload)
    enriched["ui_action"] = action
    if not str(enriched.get("message", "") or "").strip():
        enriched["message"] = str(action.get("message", "") or "").strip()
    return enriched


@dataclass
class ManagedAgent:
    name: str
    system_prompt: str
    tools: List[Any]
    model: OpenAIChatCompletionsModel
    session_id: str = ""
    hooks: Any = None


class BaseAgent:
    """OpenAI Agents SDK wrapper used by Armorer."""

    def __init__(self, provider: Optional[str] = None, model_name: Optional[str] = None):
        from armorer.config import ConfigManager

        cm = ConfigManager()
        self.provider = provider or cm.get_key("MODEL_PROVIDER") or "openai"
        self.model_name = model_name or cm.get_key("Armorer_MODEL")

        if not self.model_name:
            if self.provider == "openrouter":
                self.model_name = "qwen/qwen-2.5-coder-32b-instruct"
            elif self.provider == "ollama":
                self.model_name = "llama3"
            elif self.provider == "gemini":
                self.model_name = "gemini-2.0-flash"
            else:
                self.model_name = "gpt-4o-mini"

        self._persistent_agent: Optional[ManagedAgent] = None
        self.policy_guard = PolicyGuard.default()
        self._config_manager = cm
        self._strong_local_model = cm.get_key("ARMORER_STRONG_LOCAL_MODEL") or ""
        self._hosted_fallback_model = cm.get_key("ARMORER_HOSTED_FALLBACK_MODEL") or "gpt-4o-mini"
        self._allow_hosted_fallback = (cm.get_key("ALLOW_HOSTED_FALLBACK") or "").lower() == "true"

        from armorer.knowledge import KnowledgeBase

        self.kb = KnowledgeBase()

    def set_policy_guard(self, guard: PolicyGuard) -> None:
        self.policy_guard = guard

    def _build_openai_client(self, provider: Optional[str] = None) -> AsyncOpenAI:
        cm = self._config_manager
        provider = (provider or self.provider).lower()
        timeout_raw = cm.get_key("ARMORER_MODEL_TIMEOUT")
        timeout_s: Optional[float] = float(timeout_raw) if timeout_raw else None

        if provider == "ollama":
            base_url = cm.get_key("OLLAMA_BASE_URL") or "http://localhost:11434/v1/"
            return AsyncOpenAI(api_key="ollama", base_url=base_url, timeout=timeout_s)

        if provider == "openrouter":
            key = cm.get_key("OPENROUTER_API_KEY")
            if not key:
                raise ValueError("Missing OPENROUTER_API_KEY. Run 'armorer config' to set it.")
            return AsyncOpenAI(api_key=key, base_url="https://openrouter.ai/api/v1", timeout=timeout_s)

        if provider == "gemini":
            key = cm.get_key("GEMINI_API_KEY")
            if not key:
                raise ValueError("Missing GEMINI_API_KEY. Run 'armorer config' to set it.")
            return AsyncOpenAI(api_key=key, base_url="https://generativelanguage.googleapis.com/v1beta/openai/", timeout=timeout_s)

        if provider in {"anthropic", "bedrock"}:
            raise ValueError(f"Provider '{provider}' is not OpenAI-compatible in this backend. Use openai/openrouter/ollama/gemini.")

        key = cm.get_key("OPENAI_API_KEY")
        if not key:
            raise ValueError("Missing OPENAI_API_KEY. Run 'armorer config' to set it.")
        return AsyncOpenAI(api_key=key, timeout=timeout_s)

    def _create_model_for(self, provider: str, model_name: str):
        client = self._build_openai_client(provider=provider)
        return OpenAIChatCompletionsModel(model=model_name, openai_client=client)

    def _create_model(self):
        return self._create_model_for(self.provider, self.model_name)

    def _wrap_tools(self, tools: List[Any]) -> List[Any]:
        wrapped = []
        for t in tools:
            if hasattr(t, "run") and hasattr(t, "name"):
                wrapped.append(t)
                continue
            if callable(t):
                wrapped.append(function_tool(t, strict_mode=False))
        return wrapped

    def get_agent(
        self,
        system_prompt: str,
        tools: Optional[List] = None,
        status_obj: Any = None,
        trace_tags: Optional[List[str]] = None,
        session_id: Optional[str] = None,
        agent_name: Optional[str] = None,
        session_manager: Any = None,
        model: Any = None,
        structured_output_model: type[BaseModel] | None = None,
        hooks: Any = None,
    ) -> ManagedAgent:
        del status_obj, trace_tags, session_manager

        chosen_model = model if model is not None else self._create_model()
        if isinstance(chosen_model, str):
            old = self.model_name
            self.model_name = chosen_model
            chosen_model = self._create_model()
            self.model_name = old

        if tools is None:
            tool_fns = [
                file_read,
                file_write,
                ask_selection,
                shell,
                search_knowledge,
                get_knowledge_topic,
                list_knowledge_topics,
            ]
            if is_search_available():
                tool_fns.append(google_search)
        else:
            tool_fns = tools

        return ManagedAgent(
            name=agent_name or "armorer-agent",
            system_prompt=system_prompt,
            tools=self._wrap_tools(tool_fns),
            model=chosen_model,
            session_id=session_id or "",
            hooks=hooks,
        )

    def init_agent(
        self,
        system_prompt: str,
        tools: Optional[List] = None,
        status_obj: Any = None,
        session_id: Optional[str] = None,
        agent_name: Optional[str] = None,
        hooks: Any = None,
    ) -> None:
        self._persistent_agent = self.get_agent(
            system_prompt,
            tools=tools,
            status_obj=status_obj,
            session_id=session_id,
            agent_name=agent_name,
            hooks=hooks,
        )

    def _session_for_id(self, session_id: str):
        safe_session_id = str(session_id or "").strip()
        if not safe_session_id:
            return None
        db_path = resolve_armorer_config_dir() / "agent_sessions.sqlite"
        return SQLiteSession(safe_session_id, db_path=db_path)

    def _run_with_agent(self, managed: ManagedAgent, user_input: str, structured_output_model: Optional[Any] = None):
        set_tracing_disabled(True)
        agent = self._sdk_agent_from_managed(managed, structured_output_model=structured_output_model)
        result = Runner.run_sync(
            agent,
            user_input,
            max_turns=20,
            session=self._session_for_id(managed.session_id),
            hooks=managed.hooks,
        )
        return result

    def _sdk_agent_from_managed(
        self,
        managed: ManagedAgent,
        *,
        structured_output_model: Optional[Any] = None,
    ) -> Agent:
        return Agent(
            name=managed.name,
            instructions=managed.system_prompt,
            tools=managed.tools,
            model=managed.model,
            output_type=structured_output_model,
        )

    def create_agent_tool(
        self,
        *,
        system_prompt: str,
        tool_name: str,
        tool_description: str,
        tools: Optional[List[Any]] = None,
        agent_name: Optional[str] = None,
        session_id: Optional[str] = None,
        custom_output_extractor: Optional[Callable[[Any], Awaitable[str]]] = None,
    ):
        try:
            managed = self.get_agent(
                system_prompt=system_prompt,
                tools=tools,
                session_id=session_id,
                agent_name=agent_name or tool_name,
            )
            sdk_agent = self._sdk_agent_from_managed(managed)
            return attach_tool_metadata(
                sdk_agent.as_tool(
                    tool_name=tool_name,
                    tool_description=tool_description,
                    custom_output_extractor=custom_output_extractor,
                ),
                tool_name=tool_name,
                description=tool_description,
                category="specialist",
                verification_level="specialist_agent",
                safe_for_retrieval=False,
            )
        except Exception as exc:
            async def _unavailable_specialist(*args, **kwargs) -> str:
                del args, kwargs
                return (
                    f"Specialist tool '{tool_name}' is currently unavailable because its model backend "
                    f"could not be initialized: {exc}"
                )

            _unavailable_specialist.__name__ = tool_name
            _unavailable_specialist.__doc__ = tool_description
            return attach_tool_metadata(
                function_tool(_unavailable_specialist, strict_mode=False),
                name=tool_name,
                description=tool_description,
                category="specialist",
                verification_level="specialist_unavailable",
                safe_for_retrieval=False,
            )

    def _consent_for_hosted_fallback(self) -> bool:
        if self._allow_hosted_fallback:
            return True
        if not sys.stdin.isatty():
            return False
        answer = questionary.confirm(
            "Primary local model failed. Allow hosted fallback model for this run?",
            default=False,
        ).ask()
        return bool(answer)

    def execute(self, managed_agent: ManagedAgent, user_input: str | list[dict[str, Any]], structured_output_model: Optional[Any] = None) -> Any:
        _alerts = []
        if isinstance(user_input, str):
            filtered_input, _alerts = self.policy_guard.filter_prompt_injection(user_input)
        else:
            filtered_input = []
            for msg in user_input:
                if isinstance(msg, dict) and msg.get("role") == "user" and isinstance(msg.get("content"), str):
                    f_content, msg_alerts = self.policy_guard.filter_prompt_injection(msg["content"])
                    filtered_input.append({**msg, "content": f_content})
                    _alerts.extend(msg_alerts)
                else:
                    filtered_input.append(msg)
        last_err: Optional[Exception] = None
        res = None
        lightweight_mode = bool(not structured_output_model and _is_lightweight_conversation(filtered_input))
        trace_attempts: list[dict[str, Any]] = []
        failure_labels: list[str] = []
        all_tool_events: list[dict[str, Any]] = []
        final_response_text = ""
        outcome = "failed"
        last_payload: dict[str, Any] = {"message": ""}

        # Primary candidate
        candidates: List[ManagedAgent] = [managed_agent]
        # Stronger local fallback
        if self.provider.lower() == "ollama" and self._strong_local_model and self._strong_local_model != self.model_name:
            candidates.append(
                ManagedAgent(
                    name=managed_agent.name,
                    system_prompt=managed_agent.system_prompt,
                    tools=managed_agent.tools,
                    model=self._create_model_for("ollama", self._strong_local_model),
                )
            )

        for candidate in candidates:
            try:
                current_input: str | list[dict[str, Any]] = (
                    _build_lightweight_chat_retry_input(filtered_input)
                    if lightweight_mode
                    else filtered_input
                )
                max_attempts = 1 if structured_output_model else (2 if lightweight_mode else 3)
                for attempt in range(max_attempts):
                    try:
                        with policy_scope(self.policy_guard):
                            res = self._run_with_agent(candidate, current_input, structured_output_model=structured_output_model)
                    except Exception as e:
                        last_err = e
                        err_msg = str(e)
                        label = "unknown_tool" if "Tool " in err_msg and " not found" in err_msg else "runtime_error"
                        failure_labels.append(label)
                        trace_attempts.append(
                            {
                                "attempt": attempt + 1,
                                "input": current_input,
                                "error": err_msg,
                                "failure_label": label,
                                "new_items": [item.to_input_item() for item in getattr(res, "new_items", [])] if res else [],
                            }
                        )
                        if "Tool " in err_msg and " not found" in err_msg and attempt < max_attempts - 1:
                            current_input = _build_tool_error_retry_input(current_input, err_msg, attempt + 1)
                            continue
                        raise
                    if structured_output_model:
                        break
                    tool_events = _extract_tool_events(res)
                    if tool_events:
                        all_tool_events.extend(tool_events)
                    payload = _enrich_payload_with_background_job(_extract_response_payload(res), all_tool_events)
                    last_payload = payload
                    extracted = str(payload.get("message", "") or "").strip()
                    attempt_record = {
                        "attempt": attempt + 1,
                        "input": current_input,
                        "tool_events": tool_events,
                        "final_output": getattr(res, "final_output", None),
                        "extracted_text": extracted,
                        "response_payload": payload,
                        "new_items": [item.to_input_item() for item in getattr(res, "new_items", [])],
                    }
                    if extracted:
                        trace_attempts.append(attempt_record)
                        break
                    if attempt < max_attempts - 1:
                        attempt_record["failure_label"] = "empty_final"
                        trace_attempts.append(attempt_record)
                        failure_labels.append("empty_final")
                        if lightweight_mode:
                            current_input = _build_lightweight_chat_retry_input(current_input)
                        else:
                            current_input = _build_self_retry_input(res, attempt + 1)
                    else:
                        attempt_record["failure_label"] = "empty_final"
                        trace_attempts.append(attempt_record)
                        failure_labels.append("empty_final")
                if structured_output_model or str(last_payload.get("message", "") or "").strip():
                    break
            except Exception as e:
                last_err = e
                continue

        # Hosted fallback only after local failures and explicit consent.
        if res is None and self._consent_for_hosted_fallback():
            hosted_candidate = ManagedAgent(
                name=managed_agent.name,
                system_prompt=managed_agent.system_prompt,
                tools=managed_agent.tools,
                model=self._create_model_for("openai", self._hosted_fallback_model),
            )
            with policy_scope(self.policy_guard):
                res = self._run_with_agent(hosted_candidate, filtered_input, structured_output_model=structured_output_model)

        if res is None:
            if last_err:
                raise last_err
            raise RuntimeError("No model candidate could execute the request.")
        if structured_output_model:
            final = res.final_output
            if isinstance(final, structured_output_model):
                self._write_training_trace(
                    managed_agent,
                    filtered_input,
                    lightweight_mode,
                    trace_attempts,
                    str(final),
                    "success",
                    failure_labels,
                )
                return final
            text = str(final)
            parsed = _extract_json(text)
            if parsed:
                try:
                    structured = structured_output_model(**parsed)
                    self._write_training_trace(
                        managed_agent,
                        filtered_input,
                        lightweight_mode,
                        trace_attempts,
                        text,
                        "success",
                        failure_labels,
                    )
                    return structured
                except Exception:
                    self._write_training_trace(
                        managed_agent,
                        filtered_input,
                        lightweight_mode,
                        trace_attempts,
                        text,
                        "success",
                        failure_labels,
                    )
                    return parsed
            self._write_training_trace(
                managed_agent,
                filtered_input,
                lightweight_mode,
                trace_attempts,
                text,
                "success",
                failure_labels,
            )
            return text
        payload = _enrich_payload_with_background_job(
            _extract_response_payload(res),
            all_tool_events,
        )
        extracted = str(payload.get("message", "") or "").strip()
        if isinstance(payload.get("ui_action"), dict):
            if not extracted:
                action_message = str(payload["ui_action"].get("message", "") or "").strip()
                if action_message:
                    payload = dict(payload)
                    payload["message"] = action_message
                    extracted = action_message
            final_response_text = extracted or "Background job started."
            outcome = "success"
            self._write_training_trace(
                managed_agent,
                filtered_input,
                lightweight_mode,
                trace_attempts,
                final_response_text,
                outcome,
                failure_labels,
            )
            return payload
        if extracted:
            final_response_text = extracted
            outcome = "success"
            self._write_training_trace(
                managed_agent,
                filtered_input,
                lightweight_mode,
                trace_attempts,
                final_response_text,
                outcome,
                failure_labels,
            )
            if isinstance(payload.get("ui_action"), dict):
                return payload
            return extracted
        final_response_text = _model_retry_suggestion(self.provider, self.model_name)
        outcome = "model_switch_suggested"
        failure_labels.append("model_switch_suggested")
        self._write_training_trace(
            managed_agent,
            filtered_input,
            lightweight_mode,
            trace_attempts,
            final_response_text,
            outcome,
            failure_labels,
        )
        return final_response_text

    def _write_training_trace(
        self,
        managed_agent: ManagedAgent,
        user_input: list[dict[str, Any]],
        lightweight_mode: bool,
        trace_attempts: list[dict[str, Any]],
        final_response_text: str,
        outcome: str,
        failure_labels: list[str],
    ) -> None:
        try:
            tool_names = [tool_name(tool) or getattr(tool, "name", getattr(tool, "__name__", "")) for tool in managed_agent.tools]
            trace_id = managed_agent.session_id or f"{managed_agent.name}-{int(time.time() * 1000)}"
            tool_categories = []
            for tool in managed_agent.tools:
                meta = get_tool_metadata(tool)
                if meta:
                    tool_categories.append(
                        {
                            "name": meta.name,
                            "category": meta.category,
                            "verification_level": meta.verification_level,
                        }
                    )
            
            messages = list(user_input)
            if trace_attempts:
                last_attempt = trace_attempts[-1]
                messages = list(last_attempt.get("input", user_input))
                new_items = last_attempt.get("new_items", [])
                
                # Check if new_items already contains the final text
                has_final = False
                for item in reversed(new_items):
                    if item.get("role") == "assistant" and isinstance(item.get("content"), str) and item.get("content").strip():
                        has_final = True
                        break
                
                messages.extend(new_items)
                if not has_final and final_response_text:
                    messages.append({"role": "assistant", "content": final_response_text})
            elif final_response_text:
                messages.append({"role": "assistant", "content": final_response_text})

            trace = record_training_trace(
                trace_id=trace_id,
                session_id=managed_agent.session_id or managed_agent.name,
                messages=messages,
                provider=self.provider,
                model=self.model_name,
                system_prompt=managed_agent.system_prompt,
                task_type="lightweight_chat" if lightweight_mode else "operational_or_factual",
                tool_names_exposed=[name for name in tool_names if name],
                attempts=trace_attempts,
                final_response=final_response_text,
                outcome=outcome,
                failure_labels=failure_labels,
                guard=self.policy_guard,
            )
            if tool_categories:
                try:
                    trace_path = Path(trace["raw_path"])
                    raw = json.loads(trace_path.read_text(encoding="utf-8"))
                    raw["tool_metadata"] = tool_categories
                    trace_path.write_text(json.dumps(raw, indent=2, ensure_ascii=True), encoding="utf-8")
                except Exception:
                    pass
        except Exception:
            pass

    def __call__(self, user_input: str | list[dict[str, Any]]) -> str:
        return self.run(user_input)

    def run(
        self,
        user_input: str | list[dict[str, Any]],
        status_obj: Any = None,
        status_file: Optional[Path] = None,
        silent: bool = False,
        structured_output_model: Optional[Any] = None,
    ) -> Any:
        del status_obj, status_file, silent

        if not self._persistent_agent:
            raise RuntimeError("Agent not initialized. Call init_agent() before run().")

        if not isinstance(self._persistent_agent, ManagedAgent):
            result = self._persistent_agent(user_input, structured_output_model=structured_output_model)
            if structured_output_model and hasattr(result, "structured_output"):
                return result.structured_output
            return result

        try:
            return self.execute(self._persistent_agent, user_input, structured_output_model=structured_output_model)
        except Exception as e:
            err_msg = str(e)
            if any(x in err_msg.lower() for x in ["429", "rate limit", "throttling", "quota exceeded"]):
                if sys.stdin.isatty():
                    choice = questionary.select(
                        "How would you like to proceed?",
                        choices=["Wait 30s and retry", "Abort and Fail"],
                    ).ask()
                    if choice == "Wait 30s and retry":
                        time.sleep(30)
                        return self.execute(self._persistent_agent, user_input, structured_output_model=structured_output_model)
                else:
                    time.sleep(10)
                    return self.execute(self._persistent_agent, user_input, structured_output_model=structured_output_model)
            raise

    def clear_history(self):
        # Runner.run_sync is stateless in this wrapper.
        return
