from __future__ import annotations

import json
import logging
from copy import deepcopy
from pathlib import Path
from typing import Any

from autogen_agentchat.agents import AssistantAgent
from autogen_core.tools import FunctionTool

from prompt_auto_lab.orchestrator.state_store import SharedState

logger = logging.getLogger(__name__)



class VersionManager:
    def __init__(
        self,
        model_client: Any,
        shared_state: SharedState,
        allowed_tools: set[str] | None = None,
        field_glossary: str = "",
        name: str = "VersionManager",
        store_dir: Path | str | None = None,
    ) -> AssistantAgent:
        self.store_dir = Path(store_dir) if store_dir else None
        async def get_state() -> dict[str, Any]:
            return await shared_state.read(None)

        def _current_node(snapshot: dict[str, Any]) -> tuple[str, dict[str, Any], dict[str, Any]]:
            iterations = snapshot.get("iteration") or {}
            idx_raw = snapshot.get("iteration_idx") or "root"
            node = iterations.get(str(idx_raw)) or iterations.get(idx_raw) or {}
            return str(idx_raw), iterations, node

        def _next_node_id(iterations: dict[str, Any]) -> str:
            max_id = 0
            for k in iterations.keys():
                try:
                    max_id = max(max_id, int(k))
                except Exception:
                    continue
            return str(max_id + 1)

        async def apply_instruction(instruction_json: str) -> dict[str, Any]:
            """Deterministic executor for structured instructions."""
            inst = json.loads(instruction_json) if instruction_json else {}
            if not isinstance(inst, dict):
                raise ValueError("instruction_json must decode to an object")

            snap = await shared_state.read(None)
            idx, iterations, node_infos = _current_node(snap)
            action = inst.get("action") or "patch_current_node"
            updated: dict[str, Any] = {}

            def _inherit_fields(source_id: str, inherit_keys: list[str] | None = None) -> dict[str, Any]:
                inherit = inherit_keys or [
                    "experiment_id",
                    "experiment_link",
                    "metrics",
                    "overall_score",
                    "commit_id",
                    "branch",
                ]
                source = iterations.get(str(source_id)) or {}
                return {k: source.get(k) for k in inherit}

            def _pick_best_candidate(
                node: dict[str, Any],
                selected_ids: list[str] | None = None,
                selection_strategy: str = "greedy",
                pareto_max_selected: int = 2,
            ) -> tuple[dict[str, Any] | None, str | None]:
                """Returns (best_candidate, best_id)."""
                candidates = node.get("candidates")
                if not isinstance(candidates, dict):
                    return None, None

                # If specific IDs provided, use first valid one
                if selected_ids:
                    for cid in selected_ids:
                        entry = candidates.get(str(cid))
                        if isinstance(entry, dict):
                            return entry, cid

                # Apply selection strategy
                if selection_strategy == "gepa":
                    from prompt_auto_lab.core.pareto_utils import select_pareto_candidates
                    logger.info(f"Using Pareto selection strategy (max_selected={pareto_max_selected})")
                    # Get analysis info from node for crossover
                    analysis_commit = node.get("analysis_commit_id")
                    analysis_branch = node.get("analysis_branch")
                    # Get source_iteration from node (which iteration's candidates are being selected)
                    source_iteration = node.get("current_node") or node.get("iteration_idx")
                    result = select_pareto_candidates(
                        candidates,
                        max_selected=pareto_max_selected,
                        analysis_commit=analysis_commit,
                        analysis_branch=analysis_branch,
                    )
                    logger.info(f"Pareto selection result: primary={result.get('primary')}, pareto_front={result.get('pareto_front')}, selected={result.get('selected')}")

                    # Save pareto result to pareto.jsonl for crossover (with source_iteration)
                    result["source_iteration"] = source_iteration
                    if self.store_dir is not None:
                        pareto_file = self.store_dir / "pareto.jsonl"
                        try:
                            with open(pareto_file, "a", encoding="utf-8") as f:
                                f.write(json.dumps(result, ensure_ascii=False) + "\n")
                            logger.info(f"Saved pareto result to {pareto_file} (source_iteration={source_iteration})")
                        except Exception as e:
                            logger.warning(f"Failed to save pareto result: {e}")

                    primary_id = result.get("primary")
                    if primary_id:
                        return candidates.get(primary_id), primary_id
                    return None, None
                else:
                    logger.info("Using greedy selection strategy (highest average score)")

                # Default: greedy (max overall_score)
                best = None
                best_score = None
                best_id = None
                for c_id, entry in candidates.items():
                    if not isinstance(entry, dict):
                        continue
                    score = entry.get("overall_score")
                    if not isinstance(score, (int, float)):
                        continue
                    if best_score is None or score > best_score:
                        best_score = score
                        best = entry
                        best_id = c_id
                return best, best_id

            def _apply_best_candidate_fields(
                node: dict[str, Any], best_candidate: dict[str, Any], inherit_keys: list[str]
            ) -> None:
                # Copy key fields into the selected new candidate.
                candidates = node.get("candidates")
                if isinstance(candidates, dict):
                    for candidate in candidates.values():
                        # target_id = candidate.get("id") or "c1"
                        # target = candidates.get(str(target_id))
                        # if not isinstance(target, dict):
                        #     target = {}
                        #     candidates[str(target_id)] = target
                        for key in inherit_keys:
                            if key in best_candidate and best_candidate.get(key) is not None:
                                candidate[key] = best_candidate.get(key)
                # Mirror analysis_* fields from best candidate by naming convention.
                for key in list(node.keys()):
                    if not key.startswith("analysis_"):
                        continue
                    base_key = key[len("analysis_"):]
                    if base_key and base_key in best_candidate and best_candidate.get(base_key) is not None:
                        node[key] = best_candidate.get(base_key)

            def _apply_best_from_analysis(
                *,
                analysis_source: str,
                iterations: dict[str, Any],
                fallback_node: dict[str, Any],
                new_node: dict[str, Any],
                selected_ids: list[str] | None,
                inherit_keys: list[str],
                selection_strategy: str = "greedy",
                pareto_max_selected: int = 2,
            ) -> str | None:
                """Returns best_id."""
                analysis_node = iterations.get(str(analysis_source)) or fallback_node
                best_candidate, best_id = _pick_best_candidate(analysis_node, selected_ids, selection_strategy, pareto_max_selected)
                if isinstance(best_candidate, dict):
                    # Avoid copying candidate values to node top-level fields; keep in candidates + analysis_* only.
                    for key in inherit_keys:
                        if key in new_node:
                            new_node[key] = None
                    _apply_best_candidate_fields(new_node, best_candidate, inherit_keys)
                return best_id


            def _base_node_from_analysis(iterations: dict[str, Any], node_index:str) -> dict[str, Any]:
                analysis_node = iterations.get(node_index)
                return analysis_node or {}

            if action == "patch_current_node":
                fields = inst.get("fields") or {}
                path = inst.get("path")
                value = inst.get("value")
                idx_override = None
                if path and value is not None:
                    # support /iteration/<node>/<field>
                    if not str(path).startswith("/"):
                        raise ValueError("path must start with '/'")
                    parts = str(path).strip("/").split("/")
                    if len(parts) == 3 and parts[0] == "iteration":
                        idx_override, field = parts[1], parts[2]
                        fields = {field: value}
                    elif len(parts) == 4 and parts[0] == "iteration":
                        idx_override, field = parts[2], parts[3]
                        fields = {field: value}
                    else:
                        raise ValueError("Only /iteration/<node>/<field> paths are supported")
                if not isinstance(fields, dict):
                    raise ValueError("fields must be an object")
                target_idx = str(idx_override or idx)
                patch_iter = {target_idx: fields}
                updated["iteration"] = patch_iter
                if "continue" in fields:
                    updated["continue"] = fields.get("continue")
                    updated["stop_reason"] = fields.get("reason", "")
                await shared_state.patch({"iteration": patch_iter, **({k: updated[k] for k in ("continue", "stop_reason") if k in updated})})
                return {"ok": True, "updated": updated}

            if action == "start_new_node":
                logger.info("VersionManager: executing start_new_node action")
                parent_node = idx
                analysis_source = parent_node
                # inherit_keys = inst.get("inherit_fields")
                inherit_keys = ['experiment_id', 'experiment_link', 'metrics', 'overall_score', 'commit_id', 'branch']
                new_id = _next_node_id(iterations)
                if not new_id or new_id in iterations:
                    new_id = _next_node_id(iterations)
                base_node = node_infos
                new_node = deepcopy(base_node) if isinstance(base_node, dict) else {}
                #update keys
                new_node.update({
                    "current_node": new_id,
                    "parent_node": parent_node,
                    "analysis_source_node": analysis_source,
                    "continue": True,
                    "reason": "",

                })
                # new_node.update(_inherit_fields(analysis_source, inherit_keys))
                if isinstance(new_node.get("selected_candidates"), list):
                    new_node["selected_candidates"] = []
                # Build candidates dict based on GA population size and candidate template
                ga_cfg = snap.get("ga") if isinstance(snap.get("ga"), dict) else {}
                selection_strategy = ga_cfg.get("selection_strategy", "greedy")
                pareto_max_selected = int(ga_cfg.get("pareto_max_selected") or 2)
                logger.info(f"start_new_node: selection_strategy={selection_strategy}, pareto_max_selected={pareto_max_selected}")
                population_size = int(ga_cfg.get("default_population_size") or 1)
                if population_size < 1:
                    population_size = 1
                candidates_template: dict[str, Any] = {}
                base_candidates = base_node.get("candidates") if isinstance(base_node, dict) else None
                if isinstance(base_candidates, dict) and base_candidates:
                    best_key = next(iter(base_candidates.keys()))
                    template_candidate = base_candidates.get(best_key)
                    if isinstance(template_candidate, dict):
                        candidates_template = template_candidate
                candidates: dict[str, Any] = {}
                for i in range(population_size):
                    candidate_id = f"c{i + 1}"
                    candidates[candidate_id] = deepcopy(candidates_template)
                new_node["candidates"] = candidates
                selected_ids = inst.get("selected_candidates")
                if not isinstance(selected_ids, list):
                    selected_ids = None
                best_id = _apply_best_from_analysis(
                    analysis_source=analysis_source,
                    iterations=iterations,
                    fallback_node=node_infos,
                    new_node=new_node,
                    selected_ids=selected_ids,
                    inherit_keys=inherit_keys,
                    selection_strategy=selection_strategy,
                    pareto_max_selected=pareto_max_selected,
                )
                #update best_id to analysis_candidate
                if best_id is not None:
                    new_node["analysis_candidate"] = best_id
                patch = {"iteration": {str(new_id): new_node}, "iteration_idx": str(new_id)}
                await shared_state.patch(patch)
                return {"ok": True, "new_node": str(new_id), "status": new_node}

            if action == "rollback":
                logger.info("VersionManager: executing rollback action")
                # target_node = inst.get("target_node")
                # if target_node is None:
                #     raise ValueError("rollback requires target_node")
                parent_node = idx
                #get parent_node's parent_node as analysis_source
                target_node = node_infos.get("analysis_source_node", "root")
                analysis_source =  target_node
                inherit_keys = ['experiment_id', 'experiment_link', 'metrics', 'overall_score', 'commit_id', 'branch']

                new_id = _next_node_id(iterations)
                base_node = _base_node_from_analysis(iterations, analysis_source)
                new_node = {key: None for key in base_node.keys()}
                #update keys
                new_node.update({
                    "current_node": new_id,
                    "parent_node": parent_node,
                    "analysis_source_node": analysis_source,
                    "continue": True,
                    "reason": "",

                })
                if isinstance(new_node.get("selected_candidates"), list):
                    new_node["selected_candidates"] = []
                # Build candidates dict based on GA population size and candidate template
                ga_cfg = snap.get("ga") if isinstance(snap.get("ga"), dict) else {}
                selection_strategy = ga_cfg.get("selection_strategy", "greedy")
                pareto_max_selected = int(ga_cfg.get("pareto_max_selected") or 2)
                logger.info(f"rollback: selection_strategy={selection_strategy}, pareto_max_selected={pareto_max_selected}")
                population_size = int(ga_cfg.get("default_population_size") or 1)
                if population_size < 1:
                    population_size = 1
                candidates_template: dict[str, Any] = {}
                base_candidates = base_node.get("candidates") if isinstance(base_node, dict) else None
                if isinstance(base_candidates, dict) and base_candidates:
                    best_key = next(iter(base_candidates.keys()))
                    template_candidate = base_candidates.get(best_key)
                    if isinstance(template_candidate, dict):
                        candidates_template = template_candidate
                candidates: dict[str, Any] = {}
                for i in range(population_size):
                    candidate_id = f"c{i + 1}"
                    candidates[candidate_id] = deepcopy(candidates_template)
                new_node["candidates"] = candidates
                if new_node.get("analysis_candidate") is None and candidates:
                    new_node["analysis_candidate"] = "c1"
                selected_ids = inst.get("selected_candidates")
                if not isinstance(selected_ids, list):
                    selected_ids = None
                _apply_best_from_analysis(
                    analysis_source=analysis_source,
                    iterations=iterations,
                    fallback_node=node_infos,
                    new_node=new_node,
                    selected_ids=selected_ids,
                    inherit_keys=inherit_keys,
                    selection_strategy=selection_strategy,
                    pareto_max_selected=pareto_max_selected,
                )
                patch = {"iteration": {str(new_id): new_node}, "iteration_idx": str(new_id)}
                await shared_state.patch(patch)
                return {"ok": True, "new_node": str(new_id), "status": new_node}
            
            raise ValueError(f"Unsupported action: {action}")

        tools = []
        allowed = set(allowed_tools) if allowed_tools is not None else {"get_state", "apply_instruction"}
        if "get_state" in allowed:
            tools.append(FunctionTool(get_state, "Read shared state", strict=True))
        if "apply_instruction" in allowed:
            tools.append(
                FunctionTool(
                    apply_instruction,
                    (
                        "Apply structured instruction JSON. Schemas:\n"
                        '- patch_current_node: {"action":"patch_current_node","fields":{...}}\n'
                        '-- Example for common key states update: {"action":"patch_current_node","fields":{"analysis_results_path":"<path>"}}\n'
                        '-- Example for candidate-scoped key update: {"action":"patch_current_node","fields":{"candidates":{"c1":{"branch":"<branch_name>","commit_id":"<commit_id>"},"c2":{"branch":"<branch_name>","commit_id":"<commit_id>"}}}}\n'
                        '-- Example for metrics update: {"action":"patch_current_node","fields":{"candidates":{"c1":{"dimension_metrics":{"completeness":0.942,"engagement":0.988},"overall_score":0.9658,"experiment_id":"<uuid>","experiment_link":"<url>"}}}}\n'
                        '- start_new_node: {"action":"start_new_node"}\n'
                        '- rollback: {"action":"rollback"}\n'
                        # "- Default inherit fields: experiment_id, experiment_link, metrics, overall_score, commit_id."
                        "- **IMPORTANT**: return current node valid states, new_node valid states, roll back valid states when starting new node or rolling back."
                        "- Please follow schema strictly!, for start_new_node and rollback, do NOT invent values. only action part is required. Do NOT add extra fields."
                    ),
                    strict=True,
                )
            )

        tools_description = ", ".join([f"{tool.name}:{tool._description}" for tool in tools])

        glossary_block = f"{field_glossary}\n" if field_glossary else ""
        system_message = (
            "You are VersionManager.\n"
            "The ONLY agent allowed to mutate shared_state. Execute structured instructions deterministically.\n"
            "Your goal is to maintain shared_state, extract infos from provided context.\n"
            "Instructions come as JSON with one of actions: patch_current_node, start_new_node, rollback.\n"
            "Guidelines:\n"
            "- read current state with get_state，to know the state format.\n"
            "- Use apply_instruction with the provided JSON (do not rewrite keys), if in ga mode, please be careful about keys is in candidate level or node level.\n"
            "- Follow field_glossary strictly; treat iteration[].candidates.* as the ONLY location for candidate-scoped keys.\n"
            "- Candidate-scoped keys: branch, commit_id, experiment_id, experiment_link, dimension_metrics, overall_score, if_optimized.\n"
            "- When GA is enabled or candidates exist, never write candidate-scoped keys at iteration root.\n"
            "- If candidate_id is missing or not found, return an error and request candidate_id; do not guess.\n"
            # "- start_new_node can create once the LoopController decides to start_new_node.\n"
            "- Nodes are append-only; new nodes set parent_node and analysis_source_node and inherit baseline fields.\n"
            "- If agent=GitExecutor and output indicates reset success, update analysis_commit_id and analysis_branch (not commit_id/branch).\n"
            "- If agent=MetricsExecutor, only update iteration.candidates.*.dimension_metrics and iteration.candidates.*.overall_score; do NOT update analysis_* fields.\n"
            "- If agent=TestExecutor, only update iteration.candidates.*.experiment_id and iteration.candidates.*.experiment_link.\n"
            "- PLEASE FOLLOW THE INSTRUCTION FORMATS EXACTLY TO AVOID SHARED_STATE CORRUPTION.\n"
            f"{glossary_block}"
            f"Tools available: {tools_description}.\n"
            "Return what you updated. Do NOT invent values.\n"
        )

        self.version_manager = AssistantAgent(
            name,
            model_client=model_client,
            tools=tools,
            system_message=system_message,
            description="Applies structured state mutations.",
            max_tool_iterations=3,
        )

