from __future__ import annotations

import json
import logging
import random
import sys
from typing import Any

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.ui import Console
from autogen_core.tools import FunctionTool
from prompt_auto_lab.core.auth_token import get_azure_token_provider

from prompt_auto_lab.core.protocol import BraintrustClient
from prompt_auto_lab.core.model_clients import build_azure_model_client
from prompt_auto_lab.orchestrator.state_store import SharedState
import requests

logger = logging.getLogger(__name__)

async def fetch_uhrs_results(result_id: str) -> dict[str, Any]:
    score = random.uniform(0, 1)
    case_scores = [{"id": f"case-{i}", "score": random.uniform(0, 1)} for i in range(3)]
    type_scores = {"safety": random.uniform(0, 1), "quality": random.uniform(0, 1)}
    return {"metrics": {"overall_score": score, "cases": case_scores, "types": type_scores}}


async def get_average_score(metrics_json: str) -> dict[str, float]:
    metrics = json.loads(metrics_json)
    if not isinstance(metrics, dict):
        raise ValueError("metrics_json must decode to an object")
    numeric = {k: float(v) for k, v in metrics.items() if isinstance(v, (int, float))}
    if not numeric:
        return {"overall_score": 0.0}
    avg = sum(numeric.values()) / len(numeric)
    return {"overall_score": avg}

class metrics_executor:
    def __init__(
        self,
        model_client: Any,
        client: BraintrustClient,
        shared_state: SharedState,
        allowed_tools: set[str] | None = None,
        name: str = "MetricsExecutor",
    ):
        tools = []
        async def fetch_summary_results(experiment_id: str) -> dict[str, Any]:
            return await client.fetch_summary_results(experiment_id, if_fetch_details=True)

        async def fetch_candidates_summary_results(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"}
            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")
                experiment_id = cand.get("experiment_id")
                if not candidate_id or not experiment_id:
                    results.append(
                        {
                            "candidate_id": candidate_id,
                            "error": "candidate_id and experiment_id are required",
                        }
                    )
                    continue
                metrics = await fetch_summary_results(str(experiment_id))
                overall = await get_average_score(json.dumps(metrics, ensure_ascii=False))
                results.append(
                    {
                        "candidate_id": candidate_id,
                        "experiment_id": experiment_id,
                        "metrics": metrics,
                        "overall_score": overall.get("overall_score"),
                    }
                )
            return {"ok": True, "results": results}
        
        allowed = allowed_tools or {"fetch_event_details", "fetch_uhrs_results", "fetch_summary_results", "fetch_candidates_summary_results", "get_average_score"}
        if "fetch_event_details" in allowed:
            tools.append(
                FunctionTool(
                    client.fetch_event_details,
                    "Fetch Braintrust event details for an experiment. Args: experiment_id (string). "
                    "Returns list of events with scores. Example: experiment_id=\"4c14943b-188f-400d-bce2-2949306ed63e\".",
                    strict=True,
                )
            )
        if "fetch_summary_results" in allowed:
            tools.append(
                FunctionTool(
                    fetch_summary_results,
                    "Fetch Braintrust summary metrics for an experiment. Args: experiment_id (string). "
                    "Returns dict of metric scores. Example: experiment_id=\"4c14943b-188f-400d-bce2-2949306ed63e\".",
                    strict=True,
                )
            )
        if "fetch_candidates_summary_results" in allowed:
            tools.append(
                FunctionTool(
                    fetch_candidates_summary_results,
                    "Fetch summary metrics for multiple candidates. Args: candidates_json (JSON string) with "
                    "candidates list: {\"candidate_id\":\"id\",\"experiment_id\":\"exp\"}. "
                    "Example: {\"candidates\":[{\"candidate_id\":\"c1\",\"experiment_id\":\"4c14943b-188f-400d-bce2-2949306ed63e\"}]}.",
                    strict=True,
                )
            )
        if "fetch_uhrs_results" in allowed:
            tools.append(
                FunctionTool(
                    fetch_uhrs_results,
                    "Fetch UHRS results. Args: result_id (string). Returns metrics dict. "
                    "Example: result_id=\"uhrs-12345\".",
                    strict=True,
                )
            )
        if "get_average_score" in allowed:
            tools.append(
                FunctionTool(
                    get_average_score,
                    "Calculate the overall average score from a JSON string of metrics. Args: metrics_json (JSON object string). "
                    "Example: {\"accuracy\":0.9,\"fluency\":0.8,\"relevance\":0.85}.",
                    strict=True,
                )
            )

        tools_description= ", ".join([f"{tool.name}:{tool._description}" for tool in tools])
        
        system_message = (
            "You are MetricsExecutor. You follow Planner's requests\n"
            "Goal: fetch evaluation results using tools and any parameter you need, you can ask directly.\n"
            f"1.use tool to get overall metrics if the metrics contain multiple dimension scores, you should calculate the overall score by averaging the scores.\n"
            f"2.Must return all candidate id (s), each dimension's score and overall score.\n"
            f"example response:\n"
            f"[{{\n  \"candidate_id\": \"c1\",\n  \"metrics\": {{\"completeness\": 0.9, \"clarity & readability\": 0.8, \"engagement\": 0.9, \"factual_accuracy\": 0.7, \"groundedness\": 0.6}},\n  \"overall_score\": 0.85\n}}, {{\n  \"candidate_id\": \"c2\",\n  \"metrics\": {{\"completeness\": 0.9, \"clarity & readability\": 0.8, \"engagement\": 0.9, \"factual_accuracy\": 0.7, \"groundedness\": 0.6}},\n  \"overall_score\": 0.85\n}}]\n"
            f"Available tools: {tools_description}\n"
            "Note\n"
            "You MUST use the provided tool to get metrics.\n"
            "You are NOT allowed to fabricate metrics."
            "metrics name should always keep the same as returned by tool exactly, including case and special characters.\n"
        )

        self.metrics_executor_agent = AssistantAgent(
            name,
            model_client=model_client,
            tools=tools,
            system_message=system_message,
            model_client_stream=False,
            description="Fetches and aggregates evaluation metrics (stub).",
            max_tool_iterations=3,
        )


if __name__ == "__main__":
    # test metrics_executor
    async def test_metrics_executor():
        braintrustclient = None
        token_provider = get_azure_token_provider("https://cognitiveservices.azure.com/.default")
        model_client = build_azure_model_client(token_provider=token_provider)
        metrics_executor_agent = metrics_executor(
            model_client=model_client,
            client=braintrustclient,
            shared_state=SharedState(),
        )
        logger.info("MetricsExecutor initialized.")
        task = "please use result_id est_result_id_123 to get metrics."
        agent = metrics_executor_agent.metrics_executor_agent
        await Console(
            agent.run_stream(task=task),
            output_stats=True,
        )

    import asyncio

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
    asyncio.run(test_metrics_executor())
