import openai
import json
import logging
from typing import Dict, Any
from dataclasses import dataclass, asdict
import os
from datetime import datetime


SYSTEM_MESSAGE = """You are a mathematical problem-solving agent. You can only use these four atomic tools to solve problems:
- add(a, b): Add two numbers
- sub(a, b): Subtract b from a  
- mul(a, b): Multiply two numbers
- div(a, b): Divide a by b

Your task is to break down complex mathematical expressions into a sequence of these atomic operations, following proper order of operations (parentheses, multiplication/division, addition/subtraction).

For each step, call the appropriate tool with the correct arguments. Work step by step, showing your reasoning.

When you have the final answer, respond with just the number."""


@dataclass
class TraceEvent:
    """Single event in the application trace"""
    event_type: str  # "llm_call", "tool_execution", "error", "init", "result_extraction"
    component: str  # "openai_api", "math_tools", "agent", "parser"
    data: Dict[str, Any]
    


@dataclass
class ToolResult:
    tool_name: str
    args: Dict[str, float]
    result: float
    step_number: int
    


class MathToolsAgent:
    def __init__(self, client, model_name: str = "gpt-4o", system_message: str = SYSTEM_MESSAGE, logdir: str = "logs"):
        """
        Initialize the LLM agent with OpenAI API
        
        Args:
            client: OpenAI client instance
            model_name: Name of the model to use
            system_message: System message for the agent
            logdir: Directory to save trace logs
        """
        
        self.client = client
        self.system_message = system_message
        self.model_name = model_name
        self.step_counter = 0
        self.traces = []
        self.logdir = logdir
        
        # Create log directory if it doesn't exist
        os.makedirs(self.logdir, exist_ok=True)
        
        # Define available tools
        self.tools = [
            {
                "type": "function",
                "function": {
                    "name": "add",
                    "description": "Add two numbers together",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "a": {"type": "number", "description": "First number"},
                            "b": {"type": "number", "description": "Second number"}
                        },
                        "required": ["a", "b"]
                    }
                }
            },
            {
                "type": "function", 
                "function": {
                    "name": "sub",
                    "description": "Subtract second number from first number",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "a": {"type": "number", "description": "Number to subtract from"},
                            "b": {"type": "number", "description": "Number to subtract"}
                        },
                        "required": ["a", "b"]
                    }
                }
            },
            {
                "type": "function",
                "function": {
                    "name": "mul", 
                    "description": "Multiply two numbers together",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "a": {"type": "number", "description": "First number"},
                            "b": {"type": "number", "description": "Second number"}
                        },
                        "required": ["a", "b"]
                    }
                }
            },
            {
                "type": "function",
                "function": {
                    "name": "div",
                    "description": "Divide first number by second number", 
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "a": {"type": "number", "description": "Number to divide (numerator)"},
                            "b": {"type": "number", "description": "Number to divide by (denominator)"}
                        },
                        "required": ["a", "b"]
                    }
                }
            }
        ]
    
    def add(self, a: float, b: float) -> float:
        """Add two numbers"""
        result = a + b

        return result
    
    def sub(self, a: float, b: float) -> float:
        """Subtract b from a"""
        result = a - b
        return result
    
    def mul(self, a: float, b: float) -> float:
        """Multiply two numbers"""
        result = a * b
        return result
    
    def div(self, a: float, b: float) -> float:
        """Divide a by b"""
        if b == 0:
            raise ValueError("Division by zero")
        result = a / b
        return result

    
    def _execute_tool_call(self, tool_call) -> str:
        """Execute a tool call and return the result"""
        
        self.traces.append(TraceEvent(
            event_type="tool_execution",
            component="math_tools",
            data={"tool_name": tool_call.function.name, 
                  "args": json.loads(tool_call.function.arguments)}
        ))
        
        function_name = tool_call.function.name
        arguments = json.loads(tool_call.function.arguments)
        
        # Execute the appropriate function
        if function_name == "add":
            result = self.add(arguments["a"], arguments["b"])
        elif function_name == "sub":
            result = self.sub(arguments["a"], arguments["b"])
        elif function_name == "mul":
            result = self.mul(arguments["a"], arguments["b"])
        elif function_name == "div":
            result = self.div(arguments["a"], arguments["b"])
        else:
            raise ValueError(f"Unknown function: {function_name}")
        
        self.traces.append(TraceEvent(
            event_type="tool_result",
            component="math_tools",
            data={"result": result,}
        ))
        
        return str(result)
    
    def export_traces_to_log(self, run_id: str, problem: str, final_result: float = None):
        """
        Export traces to a log file with run_id
        
        Args:
            run_id: Unique identifier for this run
            problem: The problem that was solved
            final_result: The final result of the computation
        """
        timestamp = datetime.now().isoformat()
        log_filename = f"run_{run_id}_{timestamp.replace(':', '-').replace('.', '-')}.json"
        log_filepath = os.path.join(self.logdir, log_filename)
        
        log_data = {
            "run_id": run_id,
            "timestamp": timestamp,
            "problem": problem,
            "final_result": final_result,
            "model_name": self.model_name,
            "traces": [asdict(trace) for trace in self.traces]
        }
        
        with open(log_filepath, 'w') as f:
            json.dump(log_data, f, indent=2)
        
        logging.info(f"Traces exported to: {log_filepath}")
        return log_filepath
    
    def solve(self, problem: str, max_iterations: int = 10, run_id: str = None) -> Dict[str, Any]:
        """
        Solve a math problem using iterative planning with LLM and atomic tools
        
        Args:
            problem: Mathematical expression or problem to solve
            max_iterations: Maximum number of LLM iterations to prevent infinite loops
            run_id: Optional run identifier. If None, generates one automatically
            
        Returns:
            Final numerical result
        """
        # Generate run_id if not provided
        if run_id is None:
            run_id = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{hash(problem) % 10000:04d}"
        
        # Reset traces for each new problem
        self.traces = []
        
        logging.info(f"Solving: {problem} (Run ID: {run_id})")
        logging.info("=" * 60)
        
        # Reset state
        self.execution_history = []
        self.step_counter = 0
        

        
        messages = [
            {"role": "system", "content": self.system_message},
            {"role": "user", "content": f"Solve this mathematical expression step by step: {problem}"}
        ]
        
        iteration = 0
        while iteration < max_iterations:
            iteration += 1
            logging.info(f"\n--- LLM Iteration {iteration} ---")
            
            try:
                self.traces.append(TraceEvent(
                    event_type="llm_call",
                    component="openai_api",
                    data={
                        "model": self.model_name,
                        "messages": messages,
                        # "tools": [tool["function"] for tool in self.tools]
                    }
                ))
                
                # Call OpenAI API with function calling
                response = self.client.chat.completions.create(
                    model=self.model_name,
                    messages=messages,
                    tools=self.tools,
                    tool_choice="auto",
                    # temperature=0
                )
                
                message = response.choices[0].message
                messages.append(message.dict())
                
                self.traces.append(TraceEvent(
                    event_type="llm_response",
                    component="openai_api",
                    data={
                        "content": message.content,
                        "tool_calls": [tool.model_dump() for tool in message.tool_calls] if message.tool_calls else [],
                    }
                ))
                
                # Check if the model wants to call functions
                if message.tool_calls:
                    logging.info(f"LLM planning: {message.content or 'Executing tools...'}")
                    
                    # Execute each tool call
                    for tool_call in message.tool_calls:
                        result = self._execute_tool_call(tool_call)
                        
                        # Add tool result to conversation
                        messages.append({
                            "role": "tool",
                            "tool_call_id": tool_call.id,
                            "content": result
                        })
                else:
                    # No more tool calls - this should be the final answer
                    logging.info(f"LLM final response: {message.content}")
                    
                    # Try to extract the numerical result
                    try:
                        # Look for a number in the response
                        import re
                        numbers = re.findall(r'-?\d+\.?\d*', message.content)
                        if numbers:
                            final_result = float(numbers[-1])  # Take the last number found
                            logging.info("=" * 60)
                            logging.info(f"Final result: {final_result}")                        
                            self.traces.append(TraceEvent(
                                event_type="result_extraction",
                                component="math_tools",
                                data={"final_result": final_result}
                            ))
                            
                            # Export traces to log file
                            log_filename = self.export_traces_to_log(run_id, problem, final_result)
                            return {"result": final_result, "log_file": log_filename}

                        else:
                            logging.info("Could not extract numerical result from LLM response")
                            break
                    except ValueError:
                        logging.info("Could not parse final result as number")
                        break
                        
            except Exception as e:
                logging.info(f"Error in iteration {iteration}: {e}")
                break
        
        logging.info("Max iterations reached or error occurred")
        # Export traces even if solve failed
        return {"result": 0, "log_file": self.export_traces_to_log(run_id, problem, None)}
    

def get_default_agent(model_name: str = "gpt-4o", logdir:str = "logs") -> MathToolsAgent:
    """Get a default instance of the MathToolsAgent with OpenAI client"""
    openai_client = openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
    return MathToolsAgent(client=openai_client, model_name=model_name, logdir=logdir)


if __name__ == "__main__":
    # Example usage
    client = openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
    agent = MathToolsAgent(client, logdir="agent_logs")
    
    problem = "((2 + 3) * 4) - (6 / 2)"
    print(f"Problem: {problem}")
    
    result = agent.solve(problem)
    print(f"Result: {result}")