from __future__ import annotations

import json
from typing import Any

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

from prompt_auto_lab.core.data_models import LoopControllerOutputModel
from prompt_auto_lab.orchestrator.state_store import SharedState

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

        async def set_state(patch_json: str) -> dict[str, Any]:
            patch = json.loads(patch_json)
            if not isinstance(patch, dict):
                raise ValueError("patch_json must decode to an object")
            await shared_state.patch(patch)  # type: ignore[arg-type]
            return patch_json

        tools = []
        if "get_state" in allowed_tools:
            tools.append(FunctionTool(get_state, "Read shared state", strict=True))

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

        system_message = (
            "You are LoopController.\n"
            "Goal: decide lifecycle action for the current node (start_new_node, rollback, or stop).\n"
            "Inputs: you may read shared state via tools. Active node is iteration_idx; nodes are append-only\n"
            "  if ga.enable = 'true' (ga is Genetic Algorithm), there may be more candidates in node (except root node), you should select best score from each candidates and compare previous best score.\n"
            "Outputs: ONLY a JSON object. See examples below for the exact format for each action.\n"
            "\n"
            "## TERMINATION RULES (CRITICAL - MUST FOLLOW)\n"
            "Check task_spec.termination_rule in the state. The termination_rule defines when to STOP:\n"
            "- If termination_rule.type = 'no_progress' and termination_rule.rounds = N:\n"
            "  This means: STOP if the best overall_score has NOT improved for N consecutive rounds.\n"
            "  To check this:\n"
            "  1. Look at the 'history' array (or iterations) to find the best overall_score achieved so far.\n"
            "  2. Count how many recent rounds have passed without exceeding that best score.\n"
            "  3. If that count >= N, you MUST return action='stop' with reason explaining no progress for N rounds.\n"
            "  Example: If rounds=3 and the last 3 iterations all have scores <= the previous best, STOP.\n"
            "\n"
            "## OTHER STOP CONDITIONS\n"
            "- max_rounds reached: if current iteration_idx >= max_rounds, STOP.\n"
            "- overall_score reached maximum (perfect score, like 1.0 or 100): STOP.\n"
            "\n"
            "## ROLLBACK RULE (CRITICAL - MUST FOLLOW)\n"
            "- If the current node's best candidate overall_score is LOWER than the historical best overall_score, you MUST use 'rollback' instead of 'start_new_node'.\n"
            "  This ensures the next iteration is based on the best-performing baseline, not the worse current node.\n"
            "\n"
            "## ACTIONS AND OUTPUT EXAMPLES\n"
            "- start_new_node: continue optimizing from the current node. Only use when current score >= historical best.\n"
            '  Example: {"action": "start_new_node", "selected_candidates": ["c2"], "reason": "c2 score 0.55 improved over historical best 0.53, continue from current node."}\n'
            "- rollback: revert to previous best baseline automatically. Use when current score < historical best. No target_node or selected_candidates needed.\n"
            '  Example: {"action": "rollback", "reason": "Best candidate score 0.43 < historical best 0.53, rolling back to previous baseline."}\n'
            "- stop: terminate the optimization loop.\n"
            '  Example: {"action": "stop", "reason": "No progress for 3 consecutive rounds. Best score 0.53 not exceeded."}\n'
            "\n"
            "## CANDIDATE SELECTION\n"
            "- For start_new_node and stop: select the best candidate (max overall_score) and include its id in selected_candidates.\n"
            "- For rollback: no selected_candidates needed (the system automatically selects from the rollback target).\n"
            f"\nTools available: {tools_description}.\n"
        )

        self.loopcontroller = AssistantAgent(
            name,
            model_client=model_client,
            tools=tools,
            system_message=system_message,
            # model_client_stream=False,
            description="Controls optimization loop termination.",
            max_tool_iterations=5,
        )


if __name__ == "__main__":
    import asyncio
    import sys
    from pathlib import Path

    from autogen_agentchat.ui import Console
    from prompt_auto_lab.core.auth_token import get_azure_token_provider

    repo_root = Path(__file__).resolve().parents[2]
    sys.path.append(str(repo_root))

    from prompt_auto_lab.core.model_clients import build_azure_model_client

    async def _debug_main() -> None:
        token_provider = get_azure_token_provider("https://cognitiveservices.azure.com/.default")
        model_client = build_azure_model_client(token_provider=token_provider)

        debug_state: dict[str, Any] = {
            "task_spec": {"termination_rule": {"type": "no_progress", "rounds": 2}},
            "history": [{"version": 1, "score": 0.5}, {"version": 2, "score": 0.5}],
            "continue": True,
            "stop_reason": "",
            "record": {"version": 2, "score": 0.5},
            "max_rounds": 5,
        }

        from prompt_auto_lab.orchestrator.state_store import SharedState

        agent = LoopController(model_client=model_client, shared_state=SharedState(debug_state)).loopcontroller
        task = "Decide whether to continue given the termination_rule and history; write state.continue and state.stop_reason."
        await Console(agent.run_stream(task=task), output_stats=True)
        await model_client.close()

    asyncio.run(_debug_main())
