import math
import os
import string
import re
from typing import Any, Dict
import ray
import asyncio
import uuid

import sympy
from autogen_agentchat.agents import AssistantAgent
from autogen_core.models import ModelFamily
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams


# MCP server configuration
calculator_mcp_server = StdioServerParams(command="uvx", args=["mcp-server-calculator"])

# Reuse evaluation functions from calc_agent.py
def normalize_option(option: str) -> str:
    return re.sub(r"(\s+|\(|\))", "", option)

def is_option_result(result: str) -> bool:
    return normalize_option(result) in list(string.ascii_letters)

def float_eval(input_str: str) -> float:
    if " = around " in input_str:
        input_str = input_str.split(" = around ")[0]
    expr = sympy.parse_expr(input_str, evaluate=True)
    return float(expr.evalf())

def scalar_are_results_same(pred_result: str, true_result: str, rel_tol: float) -> bool:
    pred_result = str(pred_result) if pred_result is not None else ""
    true_result = str(true_result) if true_result is not None else ""

    if pred_result.strip() == true_result.strip():
        return True

    if is_option_result(true_result):
        true_result = normalize_option(true_result)
        pred_result = normalize_option(pred_result)
        return pred_result == true_result

    try:
        pred_float = float_eval(pred_result)
        true_float = float_eval(true_result)
        return math.isclose(pred_float, true_float, rel_tol=rel_tol)
    except Exception:
        pass

    return False

def get_agent(model, openai_base_url, temperature, workbench, task_id=None):
    try: 
        model_client = OpenAIChatCompletionClient(
            model=model,
            base_url=openai_base_url,
            api_key=os.environ.get("OPENAI_API_KEY", task_id),
            model_info={
                "vision": False,
                "function_calling": True,
                "json_output": False,
                "family": ModelFamily.UNKNOWN,
                "structured_output": False,
            },
            temperature=temperature,
        )
    except KeyError as e:
        print(f"Error: Missing required environment variable: {e}")
        raise
    except Exception as e:
        print(f"Error initializing OpenAIChatCompletionClient: {e}")
        raise

    calc_agent = AssistantAgent(
        name="calc",
        model_client=model_client,
        workbench=workbench,
        reflect_on_tool_use=True,
    )
    return calc_agent

class CalcAgentResult:
    """Result container for calc agent rollouts.
    
    DEPRECATED: This class is no longer used. Agent rollout functions now return
    Rollout objects with triplets for PPO training compatibility.
    """
    def __init__(self, answer: str, reward: float, question: str, ground_truth: str):
        self.answer = answer
        self.reward = reward
        self.question = question
        self.ground_truth = ground_truth
        self.response_length = len(answer.split()) if answer else 0
        self.turn_count = 1


############################### 
# Integration with Agentarium #
############################### 

@ray.remote(num_cpus=0.02)
def calcagent_rollout(task: Dict[str, Any], endpoint: str = "https://api.openai.com/v1", model: str = "gpt-4o-mini", task_id: str = None) -> Dict[str, Any]:
    """
    Ray remote function for calc agent rollout.
    
    Args:
        task: Dictionary containing the task data (must have 'question' key)
        endpoint: LLM endpoint URL
        model: Model name to use for the agent
        
    Returns:
        Dict: Dictionary containing reward and other necessary data for orchestrator to create Rollout
    """
    
    async def async_rollout():
        try:
            # Extract task information
            question = task.get("question", "")
            ground_truth = task.get("result", "")
            print(f"🍉 question: {question}")
            
            # Set up the agent
            async with McpWorkbench(calculator_mcp_server) as workbench:
                calc_agent = get_agent(
                    model=model,
                    openai_base_url=endpoint,
                    temperature=0.7,
                    workbench=workbench,
                    task_id=task_id
                )
                # Perform the calculation
                output_format = "Output the answer when you are ready. The answer should be surrounded by three sharps (`###`), in the form of ### ANSWER: <answer> ###."
                prompt = question + " " + output_format
                result = await calc_agent.run(task=prompt)
                
                # Extract answer
                response_text = result.messages[-1].content
                
                # Calculate reward
                reward_value = float(scalar_are_results_same(response_text, str(ground_truth), 1e-2))
                
                # Extract clean answer
                answer_match = re.search(r"###\s*ANSWER:\s*(.+?)(\s*###|$)", response_text)
                if answer_match:
                    answer = answer_match.group(1)
                else:
                    answer = response_text
                    
                # Calculate final reward
                final_reward_value = float(scalar_are_results_same(answer, str(ground_truth), 1e-2))
                
                print(f"🍊 reward: {final_reward_value}. answer: {answer}. ground_truth: {ground_truth}.")
                
                # Return data for orchestrator to create Rollout
                return {
                    "reward": final_reward_value,
                    "prompt": prompt,
                    "response_text": response_text,
                    "question": question,
                    "ground_truth": ground_truth,
                    "answer": answer,
                    "model": model,
                    "endpoint": endpoint,
                    "metadata": {}
                }
                
        except Exception as e:
            print(f"Calc agent rollout failed: {str(e)}")
            # Return failure data for orchestrator
            return {
                "reward": 0.0,
                "prompt": task.get("question", ""),
                "response_text": "Error: Failed to compute",
                "question": task.get("question", ""),
                "ground_truth": task.get("result", ""),
                "answer": "Error: Failed to compute",
                "model": model,
                "endpoint": endpoint,
                "metadata": {"error": str(e)}
            }
    
    # Run the async function
    return asyncio.run(async_rollout())

 