from __future__ import annotations

import asyncio
import json
from typing import Any

from autogen_agentchat.agents import AssistantAgent
from autogen_core.tools import FunctionTool
from prompt_auto_lab.core.protocol import GitHubActionEvalClient
from prompt_auto_lab.orchestrator.state_store import SharedState


async def run_rai_test(test_category: str, test_model: str, prompt_file_path: str) -> dict[str, Any]:
    # Stub: in reality this would trigger RAI test harness.
    job_id = f"rai-{test_category}-{test_model}"
    result_id = f"{job_id}-results"
    return {"result_id": result_id, "job_id": job_id, "result_folder": f"{job_id}-folder"}

class test_executor:
    def __init__(
        self,
        model_client: Any,
        client: GitHubActionEvalClient,
        shared_state: SharedState,
        allowed_tools: set[str] | None = None,
        name: str = "TestExecutor",
    ):
        tools = []
        allowed = set(allowed_tools or {"submit_candidates_eval", "run_rai_test"})
        allowed.discard("submit_aipa_eval")

        async def submit_candidates_eval(candidates_json: str) -> dict[str, Any]:
            payload = json.loads(candidates_json)
            if not isinstance(payload, dict):
                return {"ok": False, "error": "candidates_json must decode to an object"}
            candidates = payload.get("candidates")
            if not isinstance(candidates, list):
                return {"ok": False, "error": "candidates must be a list"}
            run_id = payload.get("run_id")
            dataset_name = payload.get("dataset_name")
            tasks = []
            results: list[dict[str, Any]] = []
            for cand in candidates:
                if not isinstance(cand, dict):
                    continue
                candidate_id = cand.get("candidate_id") or cand.get("id")
                branch = cand.get("branch") or cand.get("prompt_branch")
                commit_id = cand.get("commit_id")
                if run_id is None:
                    run_id = cand.get("run_id")
                if not candidate_id or not branch or not commit_id or not run_id:
                    results.append(
                        {
                            "candidate_id": candidate_id,
                            "error": "candidate_id, branch, commit_id, run_id are required",
                        }
                    )
                    continue
                payload_obj = {
                    "prompt_branch": branch,
                    "commit_id": commit_id,
                    "run_id": run_id,
                }
                if dataset_name:
                    payload_obj["dataset_name"] = dataset_name
                tasks.append((candidate_id, payload_obj))
            if not tasks:
                return {"ok": False, "error": "no valid candidates to submit", "results": results}

            async def _submit(index: int, candidate_id: str, payload_obj: dict[str, Any]) -> dict[str, Any]:
                # Wait before submitting: 0s for first, 5s for second, 10s for third, etc.
                await asyncio.sleep(index * 30)
                resp = await client.submit_aipa_eval(json.dumps(payload_obj))
                out = {"candidate_id": candidate_id}
                if isinstance(resp, dict):
                    out.update(resp)
                return out

            submissions = await asyncio.gather(
                *[_submit(i, candidate_id, payload_obj) for i, (candidate_id, payload_obj) in enumerate(tasks)]
            )
            return {"ok": True, "results": results + submissions}
        if "submit_candidates_eval" in allowed:
            tools.append(
                FunctionTool(
                    submit_candidates_eval,
                    "Submit multiple Aipa-Eval jobs for multiple candidates. "
                    "Args: candidates_json (JSON string) with run_id (optional string), dataset_name (optional string), "
                    "candidates (list). Each candidate requires candidate_id (string), branch or prompt_branch (string), "
                    "commit_id (string); run_id may be provided per candidate. "
                    "Example: {\"run_id\":\"1769245249\",\"candidates\":[{\"candidate_id\":\"c1\",\"branch\":\"autoprompt2/x\","
                    "\"commit_id\":\"480bc3bfb34ca4e44d385b9bed8875f2cff9b5f1\"}]}.",
                    strict=True,
                )
            )
        if "run_rai_test" in allowed:
            tools.append(
                FunctionTool(
                    run_rai_test,
                    "Run RAI test. Args: test_category (string), test_model (string), prompt_file_path (string). "
                    "Example: test_category=\"safety\", test_model=\"gpt-4\", prompt_file_path=\"prompts/prompt.txt\".",
                    strict=True,
                )
            )
        if "submit_rai_eval" in allowed:
            tools.append(FunctionTool(client.submit_rai_eval, "Submit RAI-Eval job", strict=True))

        system_message = (
        "You are TestExecutor.\n"
        "Goal: run evaluations and return result identifiers.\n"
        "If task_type=picasso: ALWAYS use submit_candidates_eval (even for a single candidate).\n"
        "Ensure every result includes candidate_id in the final output; default to c1 only if a single candidate is implied.\n"
        "If task_type=rai: run RAI test harness.\n"
        "Use tools when available; do not fabricate result_id.\n"
        )
        self.test_executor = AssistantAgent(
            name,
            model_client=model_client,
            tools=tools,
            system_message=system_message,
            description="Runs evaluation via Aipa-Eval (stub).",
            
        )
