import argparse
import asyncio

from dotenv import load_dotenv
from langsmith import Client, aevaluate
from loguru import logger

from universal_mcp.agentr.registry import AgentrRegistry
from universal_mcp.agents.codeact0 import CodeActPlaybookAgent

load_dotenv()

DEFAULT_MODEL = "anthropic:claude-sonnet-4-5"  # "gemini:gemini-3-pro-preview"
RECURSION_LIMIT = 75


def setup_logger():
    logger.remove()
    logger.add("evals.log", rotation="10 MB", compression="zip", level="INFO")


async def run_agent_eval(inputs: dict, instructions: str, eval_mode: bool = False):
    """Run agent evaluation with specified configuration."""
    api_key = inputs.get("api_key")
    registry = AgentrRegistry(api_key=api_key) if api_key else AgentrRegistry()

    base_agent = CodeActPlaybookAgent(
        "CodeAct Agent",
        instructions=instructions,
        model=DEFAULT_MODEL,
        registry=registry,
        eval_mode=eval_mode,
    )
    agent = await base_agent._build_graph()
    result = await agent.ainvoke(inputs, config={"recursion_limit": RECURSION_LIMIT})
    return result["messages"]


async def target_auto(inputs: dict):
    """Evaluation target for basic code agent tests."""
    return await run_agent_eval(
        inputs,
        instructions="Be very concise in your answers. DO NOT STOP or ASK the user any questions. Assume details if required. You are being run in develop mode, so do not use upload_file. Save files locally if required.",
    )


async def target_auto_playbook(inputs: dict):
    """Evaluation target for agent builder tests."""
    return await run_agent_eval(
        inputs,
        instructions="Be very concise in your answers. DO NOT STOP or ASK the user any questions. Assume details if required. Only test the agent once, ideally with non-default parameters to check robustness",
        eval_mode=True,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run agent evaluations")
    parser.add_argument(
        "target",
        choices=["auto", "builder"],
        help="Evaluation target to run: 'auto' for basic tests, 'builder' for agent builder tests",
    )
    args = parser.parse_args()

    client = Client()

    if args.target == "auto":
        asyncio.run(
            aevaluate(
                target_auto,
                data=client.list_examples(dataset_name="codeagent-tests-dec", splits=["3"]),
                evaluators=[],
                experiment_prefix="ci-auto-eval",
                num_repetitions=1,
                max_concurrency=10,  # TODO: Increase this, currently multiple concurrent runs have shared stdout and hence the issue
            )
        )
    elif args.target == "builder":
        asyncio.run(
            aevaluate(
                target_auto_playbook,
                data=client.list_examples(dataset_name="agent-builder", splits=["gemini"]),
                evaluators=[],
                experiment_prefix="ci-builder-eval",
                num_repetitions=1,
                max_concurrency=10,
            )
        )
    else:
        raise ValueError(f"Invalid target: {args.target}. Must be 'auto' or 'builder'.")
