from __future__ import annotations

import json
from datetime import datetime, timezone
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

class Planner:
    def __init__(
        self,
        model_client: Any,
        shared_state: SharedState,
        workflow_prompt: str,
        state_snapshot: str,
        required_block: str = "",
        field_glossary: str = "",
        name: str = "Planner",
    ) -> AssistantAgent:
        async def get_state_all() -> dict[str, Any]:
            return await shared_state.read(None)

        async def get_state_by_keys(keys: list[str]) -> dict[str, Any]:
            return await shared_state.read(keys)

        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 {"ok": True}

        async def build_task_spec(requirements_json: str) -> dict[str, Any]:
            requirements = json.loads(requirements_json)
            if not isinstance(requirements, dict):
                raise ValueError("requirements_json must decode to an object")
            spec = TaskSpec.from_dict(requirements)  # type: ignore[arg-type]
            return spec.to_dict()

        async def init_versioning() -> dict[str, Any]:
            snap = await shared_state.read(["version", "branch"])
            version = int(snap.get("version") or 1)
            branch = str(snap.get("branch") or f"autopromptlab2/v{version}")
            await shared_state.patch({"version": version, "branch": branch})
            return {"version": version, "branch": branch}

        glossary_block = f"{field_glossary}\n\n" if field_glossary else ""
        system_message = (
            "You are Planner. You are the ONLY decision maker.\n"
            "Use the workflow below to route:\n"
            f"{workflow_prompt}\n\n"
            # "Comment State (current snapshot; trust these values):\n"
            f"{state_snapshot}\n\n"
            f"{glossary_block}"
            f"{required_block}\n"
            
            "Rules:\n"
            "- Keep context minimal. Do NOT paste long history.\n"

            "- Follow workflow stage order: find the first stage whose produces keys are missing for the active iteration, using values already present in Comment State/history. Do NOT ask UserProxy for any field that already exists in Comment State. Only route to UserProxy when required inputs are truly missing.\n"
           
            "- Iteration boundary rule:\n"
            "  - Metrics generation marks the END of an iteration.\n"
            "  - If iteration_index == root, iteration starts from TestExecutor, except the experiment_id is existing.\n"
            "  - If iteration_index > root, iteration starts from BadCaseAnalyzer.\n"
            
            "- TestExecutor routing constraint:\n"
            "  - TestExecutor MUST NOT be selected as the first stage of a non-initial iteration.\n"
            "  - In non-initial iterations, TestExecutor can only be selected AFTER PromptOptimizer.\n"
            
            "- Only when LoopController decides to STOP (not rollback), the iteration can be considered complete. otherwise, you should continue progress.\n"
            "  - rollback means that the current iteration is discarded, and the previous iteration is resumed, DON'T means stop!.\n"
            "  - If LoopController outputs action=rollback, then continue per workflow; do NOT choose STOP just because rollback returned to a prior node (including root).\n"
            "  - start_new_node means that a new iteration is created based on the current iteration.\n"
            "  - stop means that the optimization process is complete.\n" 

            "- Retry rules:\n"
            "  - if any agent fails, check provided inputs and let that agent retry.\n"

            "- Output values MUST be concrete:\n"
            "  - Never emit placeholders like iteration.experiment_id or {iteration.experiment_id} in tasks or require_inputs.\n"
            "  - Use values from the STATE SNAPSHOT block; if a required value is missing, route to UserProxy to request it.\n"

            "- Output ONLY JSON: {\"decision\": {\"next_agent\": \"...\", \"reason\": \"...\", \"task\": \" action:... , require_inputs: \"}}. task is required unless STOP.\n"

        )

        self.planner = AssistantAgent(
            name,
            model_client=model_client,
            tools=[
                # FunctionTool(get_state_all, "Read shared state", strict=True),
                # FunctionTool(get_state_by_keys, "Read shared state by keys", strict=True),
                # FunctionTool(set_state, "Update shared state from JSON string", strict=True),
                # FunctionTool(build_task_spec, "Build TaskSpec JSON from requirements_json", strict=True),
                # FunctionTool(init_versioning, "Initialize version/branch", strict=True),
                # FunctionTool(build_eval_request, "Build EvalRequest JSON from task_spec_json and prompt", strict=True),
            ],
            system_message=system_message,
            description="Plans workflow and outputs routing decision JSON.",
        )


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)

        from prompt_auto_lab.orchestrator.state_store import SharedState

        debug_state: dict[str, Any] = {"user_requirements": {"task_type": "picasso"}}
        agent = Planner(
            model_client=model_client,
            shared_state=SharedState(debug_state),
            workflow_prompt="debug",
            state_snapshot=json.dumps(debug_state),
            field_glossary="",
        ).planner
        task = "Plan the next step."
        await Console(agent.run_stream(task=task), output_stats=True)
        await model_client.close()

    asyncio.run(_debug_main())
