from __future__ import annotations

import json
import logging
import re
from dataclasses import dataclass
from datetime import datetime, timezone
import os
from pathlib import Path
from typing import Any

from autogen_agentchat.messages import BaseChatMessage, ToolCallRequestEvent, ToolCallExecutionEvent
from autogen_agentchat.ui import Console
from autogen_agentchat.agents import AssistantAgent, UserProxyAgent

from prompt_auto_lab.agents.analyzer import BadCaseAnalyzer
from prompt_auto_lab.agents.crossover import CrossoverAgent
from prompt_auto_lab.agents.git_executor import GitExecutor
from prompt_auto_lab.agents.liquid_structure_agent import LiquidStructureAgent
from prompt_auto_lab.agents.loop_controller import LoopController
from prompt_auto_lab.agents.memory_agent import LongTermMemoryAgent
from prompt_auto_lab.agents.metrics_executor import metrics_executor
from prompt_auto_lab.agents.optimizer import OptimizedPrompt
from prompt_auto_lab.agents.planner import Planner
from prompt_auto_lab.agents.test_executor import test_executor
# from prompt_auto_lab.agents.state_updater import StateUpdater

from prompt_auto_lab.agents.user_proxy import build_user_proxy_agent
from prompt_auto_lab.agents.version_manager import VersionManager
from prompt_auto_lab.core.optimizer_papyrus_client import OptimizerPapyrusClient
from prompt_auto_lab.core.protocol import GitHubActionEvalClient, BraintrustClient, CodeExecutor, GitClient, PicassoClient, PythonCodeExecutor
from prompt_auto_lab.orchestrator.state_store import SharedState
from prompt_auto_lab.orchestrator.history_utils import HistoryUtils
from prompt_auto_lab.orchestrator.workflow import WorkflowConfig, format_workflow_for_prompt, load_workflow

std_logger = logging.getLogger(__name__)


class TeamAgentStub:
    def __init__(self, name: str) -> None:
        self.name = name
        self._tools: list[Any] = []

    def run_stream(self, task: str):
        raise RuntimeError("Team agents are executed by orchestrator.")


@dataclass
class Orchestrator:
    model_client: Any
    model_think_client: Any
    aipa: GitHubActionEvalClient
    braintrust: BraintrustClient
    picasso: PicassoClient
    git_agent: GitClient
    liquid_executor: CodeExecutor
    python_executor: PythonCodeExecutor
    store_dir: Path
    use_structured_output: bool = True  # Set to False for Local/Copilot models
    analysis_agent_client: Any = None  # AnalysisAgentClient for skill-based guidance tools
    _broadcast_event: Any = None  # Callback for broadcasting events to WebSocket
    _session_id: str | None = None  # Session ID for WebSocket broadcast
    _user_input_queue: Any = None  # Queue for receiving user input from WebSocket
    _user_proxy_mode: str | None = None  # web/console/callback; defaults to web for backward compatibility
    _user_input_callback: Any = None  # Optional callback for SDK/CLI user input
    _user_proxy_timeout: float = 3600.0  # User input timeout in seconds
    _prefer_native_user_proxy: bool = False  # Prefer AutoGen native UserProxyAgent when supported

    async def _generate_liquid_structure_doc(self, req: dict[str, Any]) -> None:
        """
        Generate liquid structure documentation at initialization.

        This is a one-time operation that analyzes liquid templates and generates
        the structure documentation before any agents run.

        Args:
            req: Request dictionary containing liquid_prompt_entries or liquid_prompt_file_path
        """
        from prompt_auto_lab.clients.picasso.liquid_function_client import LiquidFunctionClient

        std_logger.info("[INIT] Generating liquid structure documentation...")

        # Broadcast to frontend
        if self._broadcast_event is not None:
            await self._broadcast_event(
                "agent_message",
                "System",
                "Generating liquid structure documentation... This may take a few minutes."
            )

        liquid_client = LiquidFunctionClient(
            repo_path=self.picasso.repo_path,
            store_dir=self.store_dir,
        )

        # Remove existing structure doc to force regeneration with new configuration
        structure_doc_path = Path(liquid_client.structure_doc_path)
        if structure_doc_path.exists():
            std_logger.info("[INIT] Removing existing structure doc to regenerate with current configuration")
            structure_doc_path.unlink()

        # Get liquid entries (required field)
        liquid_entries = req.get("liquid_prompt_entries")

        try:
            if not liquid_entries:
                std_logger.error("[INIT] ✗ liquid_prompt_entries is required but not provided")
                std_logger.error("[INIT] Please provide at least one liquid template entry in your configuration.")
                raise RuntimeError("liquid_prompt_entries is required")

            # Validate and deduplicate entries
            valid_entries = []
            seen_paths = set()
            invalid_paths = []
            duplicate_paths = []
            empty_count = 0

            # Broadcast validation start
            if self._broadcast_event is not None:
                await self._broadcast_event(
                    "agent_message",
                    "System",
                    f"Validating {len(liquid_entries)} liquid template entries..."
                )

            for i, entry in enumerate(liquid_entries):
                path = entry.get("path", "").strip()

                # Skip empty paths
                if not path:
                    std_logger.warning("[INIT] Skipping entry %d: empty path", i + 1)
                    empty_count += 1
                    continue

                # Check for duplicates
                if path in seen_paths:
                    std_logger.warning("[INIT] Skipping entry %d: duplicate path '%s'", i + 1, path)
                    duplicate_paths.append(path)
                    continue

                # Check if file exists - all paths are now repo-relative
                file_path = Path(liquid_client.repo_path) / path
                if not file_path.exists():
                    # Fallback: try templates_base for legacy short paths
                    fallback = Path(liquid_client.templates_base) / path
                    if not fallback.exists():
                        std_logger.error("[INIT] Entry %d: file not found '%s'", i + 1, path)
                        invalid_paths.append(path)

                        # Broadcast error to frontend immediately
                        if self._broadcast_event is not None:
                            await self._broadcast_event(
                                "agent_message",
                                "System",
                                f"✗ Entry {i + 1}: File not found - {path}"
                            )
                        continue

                seen_paths.add(path)
                valid_entries.append(entry)
                std_logger.info("[INIT] ✓ Entry %d: %s (analyze_structure=%s)",
                               i + 1, path, entry.get("analyze_structure", True))

                # Broadcast validation success to frontend
                if self._broadcast_event is not None:
                    await self._broadcast_event(
                        "agent_message",
                        "System",
                        f"✓ Entry {i + 1}: {path} (analyze_structure={entry.get('analyze_structure', True)})"
                    )

            # Report validation summary
            validation_summary = []
            if invalid_paths:
                error_msg = f"{len(invalid_paths)} file(s) not found: " + ", ".join(invalid_paths)
                std_logger.error("[INIT] %s", error_msg)
                validation_summary.append(f"✗ {error_msg}")

            if duplicate_paths:
                warning_msg = f"{len(duplicate_paths)} duplicate(s) skipped: " + ", ".join(duplicate_paths)
                std_logger.warning("[INIT] %s", warning_msg)
                validation_summary.append(f"⚠ {warning_msg}")

            if empty_count > 0:
                warning_msg = f"{empty_count} empty path(s) skipped"
                std_logger.warning("[INIT] %s", warning_msg)
                validation_summary.append(f"⚠ {warning_msg}")

            # Broadcast validation summary to frontend
            if validation_summary and self._broadcast_event is not None:
                await self._broadcast_event(
                    "agent_message",
                    "System",
                    "\n".join(validation_summary)
                )

            if not valid_entries:
                error_msg = "No valid liquid entries found. Please verify your liquid_prompt_entries configuration."
                std_logger.error("[INIT] ✗ %s", error_msg)

                if self._broadcast_event is not None:
                    await self._broadcast_event(
                        "agent_message",
                        "System",
                        f"✗ {error_msg}"
                    )

                raise RuntimeError("No valid liquid entries found")

            # Process valid entries
            summary_msg = f"Processing {len(valid_entries)} valid entries (filtered from {len(liquid_entries)} total)"
            std_logger.info("[INIT] %s", summary_msg)

            # Broadcast processing start
            if self._broadcast_event is not None:
                await self._broadcast_event(
                    "agent_message",
                    "System",
                    f"Analyzing {len(valid_entries)} liquid template(s)... This may take several minutes."
                )

            result = await liquid_client.build_and_write_structure_doc_from_entries(
                valid_entries, self.model_client, broadcast_callback=self._broadcast_event
            )

            if result.get("ok"):
                doc_path = result.get("file_path")
                total_templates = result.get("total_templates", 0)
                std_logger.info("[INIT] ✓ Liquid structure doc generated: %s (%d templates)", doc_path, total_templates)

                # Broadcast success to frontend
                if self._broadcast_event is not None:
                    await self._broadcast_event(
                        "agent_message",
                        "System",
                        f"✓ Successfully generated structure documentation for {total_templates} templates"
                    )

                # Update shared state
                await self.shared_state.patch({"liquid_structure_doc_path": doc_path})
            else:
                error = result.get("error", "Unknown error")
                std_logger.error("[INIT] ✗ Failed to generate liquid structure doc: %s", error)

                # Broadcast error to frontend
                if self._broadcast_event is not None:
                    await self._broadcast_event(
                        "agent_message",
                        "System",
                        f"✗ Failed to generate structure doc: {error}"
                    )

        except Exception as e:
            std_logger.error("[INIT] ✗ Exception during liquid structure doc generation: %s", e, exc_info=True)

            # Broadcast exception to frontend
            if self._broadcast_event is not None:
                await self._broadcast_event(
                    "agent_message",
                    "System",
                    f"✗ Error: {str(e)}"
                )

            raise RuntimeError(f"Failed to generate liquid structure documentation: {e}") from e

    def _init_state(self, req: dict[str, Any], workflow: WorkflowConfig | None, max_rounds: int) -> dict[str, Any]:
        state: dict[str, Any] = dict(workflow.state_keys) if workflow and workflow.state_keys else {}
        state["user_requirements"] = req

        # Set task-level fields
        state["max_rounds"] = int(max_rounds)
        state["task_type"] = req.get("task_type")
        state["run_id"] = req.get("run_id")
        # state["workflow"] = workflow
        #initialize with user request fields
        for k, v in req.items():
            if k in state["iteration"]["root"]:
                state["iteration"]["root"][k] = v
            else:
                state[k] = v

        ga_cfg = state.get("ga") or {}
        # Apply user-provided selection strategy settings to GA config
        if isinstance(ga_cfg, dict):
            if req.get("selection_strategy"):
                ga_cfg["selection_strategy"] = req["selection_strategy"]
            if req.get("pareto_max_selected") is not None:
                ga_cfg["pareto_max_selected"] = req["pareto_max_selected"]
            state["ga"] = ga_cfg
            std_logger.info(f"GA config initialized: selection_strategy={ga_cfg.get('selection_strategy', 'greedy')}, pareto_max_selected={ga_cfg.get('pareto_max_selected', 2)}, enabled={ga_cfg.get('enabled', False)}")
        if isinstance(ga_cfg, dict) and ga_cfg.get("enabled"):
            root_node = state["iteration"]["root"]
            candidates = root_node.get("candidates")
            if not isinstance(candidates, dict):
                candidates = {}
            candidate_id = "c1"
            candidate_entry = candidates.get(candidate_id)
            if not isinstance(candidate_entry, dict):
                candidate_entry = {}
            candidate_fields = [
                "experiment_id",
                "experiment_link",
                "commit_id",
                "branch",
                "metrics",
                "overall_score",
            ]
            for key in candidate_fields:
                if key in req:
                    candidate_entry[key] = req.get(key)
                if key in root_node:
                    root_node[key] = None
                if key in state:
                    state.pop(key, None)
            candidates[candidate_id] = candidate_entry
            root_node["candidates"] = candidates
            if not root_node.get("analysis_candidate"):
                root_node["analysis_candidate"] = candidate_id

            # Preserve baseline experiment fields at root level so the Planner
            # can see them and skip TestExecutor when a baseline is provided.
            for key in ("experiment_id", "experiment_link"):
                if key in req and req[key]:
                    root_node[key] = req[key]

        root_node = state["iteration"]["root"]
        for key in list(root_node.keys()):
            if not key.startswith("analysis_"):
                continue
            base_key = key[len("analysis_"):]
            if base_key in req and not root_node.get(key):
                root_node[key] = req.get(base_key)

        keys_list = list(state.keys())
        state["state_keys"] = keys_list
        return state

    async def run(
        self,
        *,
        user_input: str | dict[str, Any],
        max_rounds: int = 20,
        use_console: bool = True,
        checkpoint_dir: str | Path | None = None,
        checkpoint_each_turn: bool = False,
        resume_from: str | Path | None = None,
    ) -> None:
        if not use_console:
            raise ValueError("AutoGen Console UI streaming is required (use_console must be True).")
        req = user_input if isinstance(user_input, dict) else {"raw": user_input}

        if not isinstance(req, dict):
            raise ValueError("user_input must be dict or str")

        task_type = req.get("task_type")
        workflow: WorkflowConfig | None = None
        if isinstance(task_type, str):
            workflow = load_workflow(task_type)
        allowed_agents = {s.agent for s in workflow.stages} if workflow else set()
        if workflow and workflow.team_agents:
            allowed_agents.update(workflow.team_agents.keys())
        allowed_agents.update({"UserProxy", "STOP"})

        self.checkpoint_dir = Path(checkpoint_dir) if checkpoint_dir is not None else None
        state: dict[str, Any] = self._init_state(req, workflow, max_rounds)

        self.shared_state = SharedState(initial=state)

        # Generate liquid structure documentation at initialization (one-time operation)
        # This is done before building agents to prepare the documentation upfront
        if task_type == "picasso" and resume_from is None:
            await self._generate_liquid_structure_doc(req)

        # Broadcast orchestrator starting after liquid structure is ready
        if self._broadcast_event is not None:
            await self._broadcast_event(
                "agent_start",
                "Orchestrator",
                f"Starting optimization loop with config: {json.dumps(req, default=str)}"
            )

        # Build fresh agents
        participants, agents_by_name, tools_by_name = self.build_agents(self.shared_state, workflow, state)
        versions_path = getattr(self, "version_store", None).path if getattr(self, "version_store", None) else self.store_dir / "version.jsonl"
        self.history_utils = HistoryUtils(self.shared_state, workflow, versions_path)
        self.participants = participants
        planner = agents_by_name["Planner"]
        participant_names = {p.name for p in participants}
        planner_name = planner.name

        # Resume from checkpoint if provided
        planner_task_resume: str | None = None
        turns_resume: int | None = None
        agent_turn_resume: int | None = None
        if resume_from is not None:
            (
                planner_task_resume,
                turns_resume,
                agent_turn_resume,
                participants,
                agents_by_name,
                tools_by_name,
                workflow,
            ) = await self.load_checkpoint(resume_from, self.shared_state, workflow)
            self.participants = participants
            participant_names = {p.name for p in participants}
            planner = agents_by_name["Planner"]
            planner_name = planner.name
            allowed_agents = {s.agent for s in workflow.stages} if workflow else allowed_agents
            allowed_agents.update({"UserProxy", "STOP"})

        class _Capture:
            last: BaseChatMessage | None = None

        def _resolve_task_value(snapshot: dict[str, Any], key: str) -> str:
            value: Any = snapshot
            for part in key.split("."):
                if isinstance(value, dict):
                    value = value.get(part)
                else:
                    return ""
            if value is None:
                return ""
            if isinstance(value, (dict, list)):
                return json.dumps(value, ensure_ascii=False)
            return str(value)

        def _format_task_template(template: str, snapshot: dict[str, Any]) -> str:
            def _replace(match: re.Match[str]) -> str:
                key = match.group(1).strip()
                return _resolve_task_value(snapshot, key)
            return re.sub(r"\{([^{}]+)\}", _replace, template)

        async def _run_agent_stream(agent: Any, *, task: str, initchat: bool = False) -> BaseChatMessage | None:
            std_logger.info("[SELECTOR] Selected agent: %s", agent.name)
            std_logger.info("[Agent:%s] START", agent.name)

            # Broadcast agent start event to WebSocket
            if self._broadcast_event is not None:
                payload: dict[str, Any] | None = None
                if agent.name == "UserProxy":
                    timeout_s = getattr(agent, "_timeout", None)
                    if isinstance(timeout_s, (int, float)):
                        payload = {"timeout_s": float(timeout_s)}
                try:
                    await self._broadcast_event("agent_start", agent.name, f"Starting task: {task}", payload)
                except TypeError:
                    await self._broadcast_event("agent_start", agent.name, f"Starting task: {task}")

            capture = _Capture()
            tool_call_index: dict[str, str] = {}

            async def _tap(stream):  # type: ignore[no-untyped-def]
                async for item in stream:
                    if isinstance(item, BaseChatMessage):
                        capture.last = item
                        try:
                            text = item.to_text()
                        except Exception:
                            text = str(item)
                        std_logger.info(
                            "---------- %s (%s) ----------\n%s",
                            type(item).__name__,
                            agent.name,
                            text,
                        )
                        # Broadcast message content to WebSocket
                        if self._broadcast_event is not None:
                            msg_text = item.to_text() if hasattr(item, 'to_text') else str(item)
                            await self._broadcast_event("agent_message", agent.name, msg_text)

                    # Broadcast tool call requests
                    elif isinstance(item, ToolCallRequestEvent):
                        for call in item.content:
                            tool_call_index[call.id] = call.name
                            std_logger.info(
                                "---------- ToolCallRequestEvent (%s) ----------\n%s",
                                agent.name,
                                call,
                            )
                            std_logger.info(
                                "[ToolCallMatch] request agent=%s call_id=%s name=%s",
                                agent.name,
                                call.id,
                                call.name,
                            )
                        if self._broadcast_event is not None:
                            for call in item.content:
                                try:
                                    args = json.loads(call.arguments) if call.arguments else {}
                                except json.JSONDecodeError:
                                    args = {"raw": call.arguments or ""}
                                await self._broadcast_event(
                                    "tool_call",
                                    agent.name,
                                    f"Calling {call.name}",
                                    {
                                        "tool_call": {
                                            "name": call.name,
                                            "args": args,
                                            "call_id": call.id
                                        }
                                    }
                                )

                    # Broadcast tool execution results
                    elif isinstance(item, ToolCallExecutionEvent):
                        for result in item.content:
                            std_logger.info(
                                "---------- ToolCallExecutionEvent (%s) ----------\n%s",
                                agent.name,
                                result,
                            )
                            request_name = tool_call_index.get(result.call_id, "")
                            std_logger.info(
                                "[ToolCallMatch] execution agent=%s call_id=%s request_name=%s result_name=%s matched=%s",
                                agent.name,
                                result.call_id,
                                request_name,
                                result.name,
                                "yes" if request_name == result.name else "no",
                            )
                        if self._broadcast_event is not None:
                            for result in item.content:
                                await self._broadcast_event(
                                    "tool_result",
                                    agent.name,
                                    f"Result from {result.name}",
                                    {
                                        "tool_result": {
                                            "name": result.name,
                                            "call_id": result.call_id,
                                            "content": result.content,
                                            "is_error": result.is_error or False
                                        }
                                    }
                                )

                    yield item
            if initchat:
                std_logger.info("Initial chat history for agent %s:", agent.name)
                await agent.model_context.clear()
            stream = agent.run_stream(task=task)

            # Try to use Console output, but fallback gracefully on Windows errors
            try:
                await Console(_tap(stream))
            except OSError as e:
                # Windows console may fail with certain characters - consume stream without console output
                std_logger.warning("Console output failed (OS error), consuming stream silently: %s", e)
                async for _ in _tap(stream):
                    pass

            # Broadcast agent complete event to WebSocket
            if self._broadcast_event is not None:
                result_text = capture.last.to_text() if capture.last else "No output"
                await self._broadcast_event("agent_complete", agent.name, result_text)

            return capture.last

        # Point-to-point routing loop:
        # - Planner has global context (its own memory + state tools + prior observations we feed it)
        # - Each other agent only sees tasks Planner sends (isolated per-agent memory)
        no_memory_agents = [
            "promptoptimizer",
            "badcaseanalyzer",
            "testexecutor",
            "metricsexecutor",
            "longtermmemoryagent",
        ]

        async def _run_single_agent(agent_name: str, task: str) -> str:
            initchat_ = agent_name.lower() in no_memory_agents
            agent = agents_by_name[agent_name]
            agent_msg = await _run_agent_stream(agent, task=task, initchat=initchat_)
            return agent_msg.to_text().strip() if agent_msg is not None else ""

        async def _run_memory_agent_per_candidate(task: str) -> str:
            snapshot = await self.shared_state.read(None)
            iteration_idx = snapshot.get("iteration_idx", "root")
            node_info = (snapshot.get("iteration") or {}).get(str(iteration_idx), {})
            candidates = node_info.get("candidates") if isinstance(node_info, dict) else {}
            candidate_ids = list(candidates.keys()) if isinstance(candidates, dict) else []
            if not candidate_ids:
                return await _run_single_agent("LongTermMemoryAgent", task)
            outputs: list[str] = []
            for cid in candidate_ids:
                per_candidate_task = (
                    f"{task}\n\nCandidate to analyze: {cid}\n"
                    "Use this candidate_id for score_change_by_candidate, git_diff_by_candidate, "
                    "and write_memory_summary_by_candidate."
                )
                per_text = await _run_single_agent("LongTermMemoryAgent", per_candidate_task)
                outputs.append(f"[candidate {cid}] {per_text}")
            return "\n\n".join(outputs)

        def _build_agent_task(task_template: str, snapshot: dict[str, Any], agent_task: str) -> str:
            stage_task = _format_task_template(task_template, snapshot) if task_template else ""
            base_task = _format_task_template(agent_task, snapshot) if agent_task else ""
            if base_task and stage_task:
                return f"{base_task}\n\n{stage_task}"
            return base_task or stage_task

        async def _run_team_agents(team_name: str, agent_task: str, iteration_idx: int) -> tuple[str, str | None]:
            nonlocal turns
            team_stages = workflow.team_agents[team_name]
            team_outputs: list[str] = []
            missing_agent: str | None = None

            for stage in team_stages:
                sub_agent = agents_by_name.get(stage.agent)
                if sub_agent is None:
                    missing_agent = stage.agent
                    break
                snapshot = await self.shared_state.read(None)
                stage_task = stage.task_template or stage.description or f"Run {stage.name}"
                sub_task = stage_task + f"\n\n iteration number:{iteration_idx}" +"\n\nBackGround:" + agent_task
                if team_outputs:
                    sub_task += f"\n\nNote: Previous stage outputs:\n" + "\n".join(team_outputs)
                initchat_ = stage.agent.lower() in no_memory_agents
                sub_msg = await _run_agent_stream(sub_agent, task=sub_task, initchat=initchat_)
                # turns += 1
                sub_text = sub_msg.to_text().strip() if sub_msg is not None else ""
                team_outputs.append(f"[{stage.agent}] {sub_text}")
            return "\n".join(team_outputs), missing_agent

        async def _run_optimize_team(agent_task: str) -> tuple[str, str | None]:
            snapshot = await self.shared_state.read(None)
            ga_cfg = snapshot.get("ga") or {}
            population_size = int(ga_cfg.get("default_population_size") or 1)
            if population_size < 1:
                population_size = 1
            outputs: list[str] = []
            missing_agent: str | None = None
            for i in range(1, population_size + 1):
                run_text, missing_agent = await _run_team_agents("OptimizeTeam", agent_task, iteration_idx=i)
                if missing_agent is not None:
                    break
                outputs.append(f"candidate id: c{i}:\n{run_text}")
            if missing_agent is not None:
                return "", missing_agent
            agent_text = "\n\n".join(outputs)

            # After OptimizeTeam, check if crossover should be run
            crossover_parents = await _get_crossover_parents()
            if crossover_parents is not None:
                crossover_text, crossover_missing = await _run_crossover_team(
                    parent_commits=crossover_parents["parent_commits"],
                    analysis_commit=crossover_parents["analysis_commit"],
                    analysis_branch=crossover_parents.get("analysis_branch"),
                )
                if crossover_missing is None and crossover_text:
                    agent_text = agent_text + "\n\n" + crossover_text
                elif crossover_missing:
                    std_logger.warning(f"Crossover skipped due to missing agent: {crossover_missing}")

            return agent_text, None

        async def _get_crossover_parents() -> dict[str, Any] | None:
            """Get the crossover info from pareto.jsonl (written by VersionManager).

            Returns:
                crossover_info dict with 'parent_commits' list if crossover should happen, None otherwise.
            """
            snapshot = await self.shared_state.read(None)
            ga_cfg = snapshot.get("ga") or {}

            # Only run crossover if using gepa selection strategy
            if ga_cfg.get("selection_strategy") != "gepa":
                return None

            # No crossover on first iteration
            iteration_idx = snapshot.get("iteration_idx", "root")
            if iteration_idx == "root":
                return None

            # Get analysis_source_node to verify pareto result matches
            node = snapshot.get("iteration", {}).get(str(iteration_idx), {})
            analysis_source = node.get("analysis_source_node")
            if not analysis_source:
                return None

            # Read crossover_info from pareto.jsonl (last line)
            pareto_file = self.store_dir / "pareto.jsonl"
            if not pareto_file.exists():
                std_logger.warning(f"pareto.jsonl not found at {pareto_file}")
                return None

            try:
                with open(pareto_file, "r", encoding="utf-8") as f:
                    last_line = None
                    for line in f:
                        line = line.strip()
                        if line:
                            last_line = line
                    if not last_line:
                        return None
                    result = json.loads(last_line)
            except Exception as e:
                std_logger.warning(f"Failed to read pareto.jsonl: {e}")
                return None

            # Verify the pareto result matches current iteration's analysis_source
            source_iteration = result.get("source_iteration")
            if str(source_iteration) != str(analysis_source):
                std_logger.warning(
                    f"Pareto result source_iteration ({source_iteration}) does not match "
                    f"current analysis_source_node ({analysis_source}), skipping crossover"
                )
                return None

            crossover_info = result.get("crossover_info")
            if crossover_info and crossover_info.get("parent_commits"):
                parent_commits = crossover_info["parent_commits"]
                std_logger.info(
                    f"Crossover: {len(parent_commits)} parents {parent_commits} "
                    f"(base: {crossover_info.get('analysis_commit')}, branch: {crossover_info.get('analysis_branch')})"
                )
            else:
                crossover_info = None
            return crossover_info

        async def _run_crossover_team(
            parent_commits: list[str],
            analysis_commit: str,
            analysis_branch: str | None = None,
        ) -> tuple[str, str | None]:
            """Run CrossoverTeam to create a hybrid candidate from N parents.

            Args:
                parent_commits: List of parent commit IDs (N >= 2)
                analysis_commit: The common ancestor commit (analysis baseline)
                analysis_branch: The branch name for the analysis commit (for git_new_branch)

            Returns:
                Tuple of (output_text, missing_agent_name)
            """
            # Determine the crossover candidate ID (next after population_size)
            snapshot = await self.shared_state.read(None)
            ga_cfg = snapshot.get("ga") or {}
            population_size = int(ga_cfg.get("default_population_size") or 3)
            crossover_candidate_id = f"c{population_size + 1}"

            std_logger.info(f"Running CrossoverTeam for {crossover_candidate_id} from {len(parent_commits)} parents: {parent_commits}")

            # Build task with commit info and branch for git_new_branch
            crossover_task = (
                f"Create crossover candidate {crossover_candidate_id} by combining optimizations from {len(parent_commits)} parent commits.\n"
                f"For git_new_branch, use analysis_branch: {analysis_branch}\n"
                f"For CrossoverAgent, use get_parent_diffs with:\n"
                f"  - parent_commits: {json.dumps(parent_commits)}\n"
                f"  - analysis_commit: {analysis_commit}\n"
                f"Then merge the changes intelligently."
            )

            # Run CrossoverTeam (git_new_branch -> CrossoverAgent -> git_commit)
            run_text, missing_agent = await _run_team_agents(
                "CrossoverTeam",
                crossover_task,
                iteration_idx=population_size + 1,  # c4 when population_size=3
            )

            if missing_agent is not None:
                return "", missing_agent

            return f"crossover candidate id: {crossover_candidate_id}:\n{run_text}", None

        async def _update_state_from_agent(agent_name: str, agent_text: str) -> None:
            updater = agents_by_name.get("VersionManager")
            if updater is None or updater is planner:
                return
            await _run_agent_stream(
                updater,
                task=f"Please update states by tool apply_instruction from this agent output (agent={agent_name}):\n{agent_text}.",
            )

        async def _apply_state_patch(fields: dict[str, Any]) -> None:
            updater = agents_by_name.get("VersionManager")
            if updater is None or updater is planner:
                return
            instruction = {"action": "patch_current_node", "fields": fields}
            await _run_agent_stream(
                updater,
                task=(
                    "Use apply_instruction with this JSON only:\n"
                    + json.dumps(instruction, ensure_ascii=False)
                ),
            )

        # max_turns = max_rounds * (len(participants) + 1)
        max_turns = max_rounds
        turns = int(turns_resume or 0)
        # agent_turn tracks agent invocations separately, used for checkpoint filenames
        agent_turn = int(agent_turn_resume or 0)
        user_req = state.get("user_requirements")
        planner_task = planner_task_resume or json.dumps(user_req, ensure_ascii=False)

        while turns < max_turns:
            # comment_state = {k: v for k, v in (await self.shared_state.read(None)).items() if k not in {"iteration", "workflow"}}
            snapshot = await self.history_utils.current_snapshot()
            # planner_prompt = f"COMMENT STATE:\n{json.dumps(comment_state, ensure_ascii=False)}\n\nSTATE SNAPSHOT:\n{json.dumps(snapshot, ensure_ascii=False)}\n\n {planner_task}"
            planner_prompt = f"STATE SNAPSHOT:\n{json.dumps(snapshot, ensure_ascii=False)}\n\n {planner_task}"
            planner_msg = await _run_agent_stream(planner, task=planner_prompt)
            # turns += 1
            if planner_msg is None:
                planner_task = "Return ONLY valid JSON with a decision object."
                continue

            raw = planner_msg.to_text().strip()
            try:
                #remove ````json` code block if any
                raw = raw.replace("```json", "").replace("```", "").strip()
                data = json.loads(raw)
            except Exception:
                std_logger.info(
                    "[ROUTER] %s → NextAgent=%s | reason=%s",
                    planner_name,
                    planner_name,
                    "Planner output is not valid JSON",
                )
                planner_task = "Your previous output was not valid JSON. Return ONLY valid JSON per the schema."
                continue

            decision = data.get("decision") if isinstance(data, dict) else None
            if not isinstance(decision, dict):
                std_logger.info(
                    "[ROUTER] %s → NextAgent=%s | reason=%s",
                    planner_name,
                    planner_name,
                    "Missing decision object",
                )
                planner_task = "Missing decision object. Return ONLY valid JSON: {\"decision\": {...}}"
                continue

            next_agent = str(decision.get("next_agent") or "")
            reason = str(decision.get("reason") or "")
            agent_task = str(decision.get("task") or "").strip()
            std_logger.info(
                "[ROUTER] %s → NextAgent=%s | reason=%s",
                planner_name,
                next_agent,
                reason,
            )

            # Broadcast routing decision to WebSocket
            if self._broadcast_event is not None:
                await self._broadcast_event(
                    "routing",
                    "Planner",
                    f"Routing to {next_agent}: {reason}",
                    {"decision": {"next_agent": next_agent, "reason": reason, "task": agent_task}}
                )
        
            if next_agent == "STOP":
                break

            if next_agent not in participant_names:
                planner_task = f"Unknown next_agent={next_agent!r}. Choose one of: {sorted(participant_names)} or STOP."
                continue

            if next_agent not in allowed_agents:
                planner_task = (
                    f"next_agent={next_agent!r} is not allowed by workflow. "
                    f"Choose one of: {sorted(allowed_agents)}."
                )
                continue

            if not agent_task:
                agent_task = (
                    "No task was provided. Provide a concise task for this agent, including any required identifiers "
                    "(e.g., result_id) and required output format."
                )

            if workflow is not None and next_agent in (workflow.team_agents or {}):
                if next_agent == "OptimizeTeam":
                    agent_text, missing_agent = await _run_optimize_team(agent_task)
                    if missing_agent is not None:
                        planner_task = f"Unknown team sub-agent={missing_agent!r} for team {next_agent!r}. Choose a valid agent."
                        continue
                else:
                    agent_text, missing_agent = await _run_team_agents(next_agent, agent_task)
                    if missing_agent is not None:
                        planner_task = f"Unknown team sub-agent={missing_agent!r} for team {next_agent!r}. Choose a valid agent."
                        continue

            else:
                if next_agent.lower() == "longtermmemoryagent":
                    agent_text = await _run_memory_agent_per_candidate(agent_task)
                else:
                    agent_text = await _run_single_agent(next_agent, agent_task)

            

            # Auto-update state based on agent output before handing back to Planner
            await _update_state_from_agent(next_agent, agent_text)
            # state current state after LoopController actions
            if next_agent == "LoopController":
                result = await self.shared_state.save_state_locally(os.path.join(self.store_dir, "version.jsonl"))
                std_logger.info("State saved after LoopController action: %s", result)
                turns += 1

            planner_task = f"OBS from {next_agent}:\n{agent_text}\n\nDecide next step. Return ONLY valid JSON."

            agent_turn += 1
            if checkpoint_each_turn:
                await self.save_checkpoint(planner_task, turns, agent_turn, next_agent)

        std_logger.info("Orchestrator run completed in %s turns.", turns)
        if self.checkpoint_dir is not None and not checkpoint_each_turn:
            await self.save_checkpoint(planner_task, turns, agent_turn, "final")

    def build_agents(
        self,
        shared_state: SharedState,
        workflow_cfg: WorkflowConfig | None,
        state_snapshot: dict[str, Any],
    ) -> tuple[list[Any], dict[str, Any], dict[str, list[Any]]]:
        workflow_prompt = format_workflow_for_prompt(workflow_cfg) if workflow_cfg else "No workflow loaded."
        state_for_prompt = {k: v for k, v in state_snapshot.items() if k != "workflow" or k != "iteration"}
        state_for_prompt["workflow_task_type"] = workflow_cfg.task_type if workflow_cfg else None
        required_lines = workflow_cfg.required_state_lines() if workflow_cfg else []
        glossary_lines = workflow_cfg.field_glossary_lines() if workflow_cfg else []
        required_block = (
            "Required state fields (ask UserProxy if missing):\n- " + "\n- ".join(required_lines)
            if required_lines
            else ""
        )
        field_glossary_block = (
            "Field glossary (schema hints):\n- " + "\n- ".join(glossary_lines)
            if glossary_lines
            else ""
        )
        planner_agent = Planner(
            model_client=self.model_client,
            shared_state=shared_state,
            workflow_prompt=workflow_prompt,
            state_snapshot=json.dumps(state_for_prompt, ensure_ascii=False),
            required_block=required_block,
            field_glossary=field_glossary_block,
        ).planner
        agent_tools_cfg = workflow_cfg.agent_tools if workflow_cfg else {}

        allowed_agents = {s.agent for s in workflow_cfg.stages} if workflow_cfg else set()
        if workflow_cfg and workflow_cfg.team_agents:
            allowed_agents.update(workflow_cfg.team_agents.keys())
        allowed_agents.update({"UserProxy", "STOP"})

        def tools_for(name: str, fallback: set[str] | None = None) -> set[str]:
            tools = agent_tools_cfg.get(name)
            if tools is None:
                return set(fallback or [])
            return set(tools)

        test_exec_agent = test_executor(
            model_client=self.model_client,
            client=self.aipa,
            shared_state=shared_state,
            allowed_tools=tools_for("TestExecutor"),
        ).test_executor

        metrics_exec_agent = metrics_executor(
            model_client=self.model_client,
            client=self.braintrust,
            shared_state=shared_state,
            allowed_tools=tools_for("MetricsExecutor"),
        ).metrics_executor_agent

        analyzer_agent = BadCaseAnalyzer(
            model_client=self.model_client,
            shared_state=shared_state,
            allowed_tools=tools_for("BadCaseAnalyzer"),
            braintrust_client=self.braintrust,
            use_structured_output=self.use_structured_output,
        ).analyzer

        optimizer_agent = OptimizedPrompt(
            model_client=self.model_client,
            client=self.picasso,
            shared_state=shared_state,
            braintrust_client=self.braintrust,
            allowed_tools=tools_for("PromptOptimizer"),
            use_structured_output=self.use_structured_output,
            analysis_agent_client=self.analysis_agent_client,
        ).optimizedPrompt

        git_executor_agent = GitExecutor(
            model_client=self.model_client,
            git_client=self.git_agent,
            shared_state=shared_state,
            allowed_tools=tools_for("GitExecutor", {"git_commit", "git_push"}),
        ).git_executor

        # store = VersionStore(self.store_dir)
        # self.version_store = store
        version_mgr_agent = VersionManager(
            model_client=self.model_client,
            shared_state=shared_state,
            allowed_tools=tools_for("VersionManager"),
            field_glossary=field_glossary_block,
            store_dir=self.store_dir,
        ).version_manager

        loop_ctrl_agent = LoopController(
            model_client=self.model_client,
            shared_state=shared_state,
            allowed_tools=tools_for("LoopController"),
        ).loopcontroller
        
        memory_agent = LongTermMemoryAgent(
            model_client=self.model_client,
            shared_state=shared_state,
            git_client=self.git_agent,
            braintrust_client=self.braintrust,
            allowed_tools=tools_for("LongTermMemoryAgent"),
        ).agent

        liquid_structure_agent = LiquidStructureAgent(
            model_client=self.model_client,
            shared_state=shared_state,
            allowed_tools=tools_for("LiquidStructureAgent"),
            repo_path=self.picasso.repo_path,
            store_dir=self.store_dir,
            use_structured_output=self.use_structured_output,
            braintrust_client=self.braintrust,
        ).agent

        crossover_agent = CrossoverAgent(
            model_client=self.model_client,
            picasso_client=self.picasso,
            git_client=self.git_agent,
            shared_state=shared_state,
            allowed_tools=tools_for("CrossoverAgent"),
        ).crossover_agent

        # state_updater_agent = StateUpdater(
        #     model_client=self.model_client,
        #     shared_state=shared_state,
        #     allowed_tools=tools_for("StateUpdater"),
        # ).state_updater
        # Keep default mode as web to preserve existing API behavior.
        user_proxy_agent = build_user_proxy_agent(
            name="UserProxy",
            mode=self._user_proxy_mode or "web",
            input_queue=self._user_input_queue,
            input_callback=self._user_input_callback,
            timeout=self._user_proxy_timeout,
            prefer_native=self._prefer_native_user_proxy,
        )

        team_agents_cfg = workflow_cfg.team_agents if workflow_cfg else {}
        team_stubs = [TeamAgentStub(name) for name in team_agents_cfg]

        participants_local = [
            planner_agent,
            test_exec_agent,
            metrics_exec_agent,
            analyzer_agent,
            liquid_structure_agent,
            memory_agent,
            optimizer_agent,
            crossover_agent,
            git_executor_agent,
            version_mgr_agent,
            loop_ctrl_agent,
            # state_updater_agent,
            user_proxy_agent,
            *team_stubs,
        ]
        agents_dict = {p.name: p for p in participants_local}
        tools_dict = {p.name: getattr(p, "_tools", []) for p in participants_local}
        return participants_local, agents_dict, tools_dict

    async def save_checkpoint(self, task: str, turns: int, agent_turn: int, agent_name: str) -> None:
        if self.checkpoint_dir is None:
            return
        snapshot = await self.shared_state.read(None)
        agents_dump = {}
        for agent in getattr(self, "participants", []):
            if isinstance(agent, (AssistantAgent, UserProxyAgent)):
                # Use save_state() method which returns a Mapping
                try:
                    state = await agent.save_state()
                    agents_dump[agent.name] = dict(state) if state else {}
                except Exception:
                    agents_dump[agent.name] = {"name": agent.name}
        payload = {
            "state": snapshot,
            "planner_task": task,
            "turns": turns,
            "agent_turn": agent_turn,
            "agent_name": agent_name,
            "task_type": snapshot.get("task_type"),
            "agents": agents_dump,
            "saved_at": datetime.now(timezone.utc).isoformat(),
        }
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        # Filename format: checkpoint_{agent_turn}_{agent}.json
        checkpoint_filename = f"checkpoint_{agent_turn:04d}_{agent_name}.json"
        checkpoint_file = self.checkpoint_dir / checkpoint_filename
        checkpoint_file.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")

    async def load_checkpoint(
        self,
        resume_from: str | Path,
        shared_state: SharedState,
        workflow: WorkflowConfig | None,
    ) -> tuple[str | None, int | None, int | None, list[Any], dict[str, Any], dict[str, list[Any]], WorkflowConfig | None]:
        payload = json.loads(Path(resume_from).read_text(encoding="utf-8"))
        planner_task_resume: str | None = None
        turns_resume: int | None = None
        agent_turn_resume: int | None = None
        if isinstance(payload, dict):
            saved_state = payload.get("state")
            if isinstance(saved_state, dict):
                await shared_state.patch(saved_state)
            planner_task_resume = payload.get("planner_task")
            turns_resume = payload.get("turns")
            agent_turn_resume = payload.get("agent_turn")
            saved_task_type = (saved_state or {}).get("task_type") if isinstance(saved_state, dict) else None
            if isinstance(saved_task_type, str):
                workflow = load_workflow(saved_task_type)
            agents_dump = payload.get("agents") or {}
            participants, agents_by_name, tools_by_name = self.build_agents(shared_state, workflow, (await shared_state.read(None)))
            # Restore agent states from checkpoint
            for agent in participants:
                dump = agents_dump.get(agent.name)
                if dump and isinstance(agent, (AssistantAgent, UserProxyAgent)):
                    try:
                        await agent.load_state(dump)
                    except Exception as e:
                        std_logger.warning(
                            "Failed to restore state for agent %s: %s",
                            agent.name,
                            e,
                        )
            return planner_task_resume, turns_resume, agent_turn_resume, participants, agents_by_name, tools_by_name, workflow
        # No checkpoint found, build agents from scratch
        participants, agents_by_name, tools_by_name = self.build_agents(shared_state, workflow, (await shared_state.read(None)))
        return planner_task_resume, turns_resume, agent_turn_resume, participants, agents_by_name, tools_by_name, workflow


def find_latest_checkpoint(checkpoint_dir: Path) -> Path | None:
    """
    Find the latest checkpoint file in the checkpoint directory.
    Filename format: checkpoint_{agent_turn}_{agent}.json
    Returns the file with the highest agent_turn number.
    """
    if not checkpoint_dir.exists():
        return None
    
    checkpoint_files = list(checkpoint_dir.glob("checkpoint_*.json"))
    if not checkpoint_files:
        return None
    
    # Sort by turn number in filename, get the highest
    def extract_turn(path: Path) -> int:
        # Filename format: checkpoint_0001_AgentName.json
        name = path.stem  # checkpoint_0001_AgentName
        parts = name.split("_")
        if len(parts) >= 2:
            try:
                return int(parts[1])
            except ValueError:
                return -1
        return -1
    
    checkpoint_files.sort(key=extract_turn, reverse=True)
    return checkpoint_files[0] if checkpoint_files else None


def default_store_dir(run_id: str) -> Path:
    return Path(".") / "autopromptlab2" / "logs" / run_id
