"""
Reasoning tools for the GAIA agent.

This module provides tools for complex reasoning tasks, including:
- Step-by-step reasoning
- Mathematical calculations
- Fact verification

All tools handle errors gracefully and provide detailed error messages.
"""

import logging
import traceback
import re
import math
import json
from typing import Dict, Any, List, Optional, Union, Tuple, Callable
from datetime import datetime

import operator
from sympy import symbols, sympify, solve, simplify
import numpy as np

from src.gaia.agent.config import get_model_config, get_tool_config
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser

logger = logging.getLogger("gaia_agent.tools.reasoning")

class StepByStepReasoner:
    """Tool for performing step-by-step reasoning."""
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """
        Initialize the step-by-step reasoner.
        
        Args:
            config: Optional configuration dictionary
        """
        self.model_config = config or get_model_config()
        self.model = ChatOpenAI(
            model=self.model_config.get("reasoning_model", "gpt-4o"),
            temperature=self.model_config.get("temperature", 0.1),
            max_tokens=self.model_config.get("max_tokens", 4096)
        )
    
    def reason(self, question: str, context: Optional[str] = None) -> Dict[str, Any]:
        """
        Perform step-by-step reasoning on a question.
        
        Args:
            question: The question to reason about
            context: Optional context information
            
        Returns:
            Dictionary containing the reasoning steps and conclusion
            
        Raises:
            Exception: If an error occurs during reasoning
        """
        
        try:
            prompt_template = """You are an expert at step-by-step reasoning.
            
            Your task is to solve the following problem by breaking it down into clear, logical steps.
            
            {context_text}
            
            Question: {question}
            
            Think through this step-by-step and provide your reasoning in the following JSON format:
            {{
                "steps": [
                    {{
                        "step_number": 1,
                        "reasoning": "First step of reasoning",
                        "intermediate_conclusion": "Conclusion from this step"
                    }},
                    ...
                ],
                "final_conclusion": "The final answer based on all steps",
                "confidence": 0.95  // A number between 0 and 1 indicating your confidence
            }}
            
            JSON Response:"""
            
            context_text = f"Context: {context}" if context else "No additional context provided."
            
            prompt = PromptTemplate.from_template(prompt_template)
            
            chain = prompt | self.model | StrOutputParser()
            
            result = chain.invoke({
                "question": question,
                "context_text": context_text
            })
            
            parsed_result = json.loads(result)
            
            
            return parsed_result
        
        except Exception as e:
            logger.error(f"Error during step-by-step reasoning: {str(e)}")
            logger.error(traceback.format_exc())
            raise Exception(f"Step-by-step reasoning failed: {str(e)}")
    
    def chain_of_thought(self, question: str, context: Optional[str] = None) -> str:
        """
        Perform chain-of-thought reasoning and return a textual response.
        
        Args:
            question: The question to reason about
            context: Optional context information
            
        Returns:
            String containing the reasoning and conclusion
            
        Raises:
            Exception: If an error occurs during reasoning
        """
        
        try:
            prompt_template = """You are an expert at chain-of-thought reasoning.
            
            Your task is to solve the following problem by thinking step-by-step.
            
            {context_text}
            
            Question: {question}
            
            Let's think through this step-by-step:"""
            
            context_text = f"Context: {context}" if context else "No additional context provided."
            
            prompt = PromptTemplate.from_template(prompt_template)
            
            chain = prompt | self.model | StrOutputParser()
            
            result = chain.invoke({
                "question": question,
                "context_text": context_text
            })
            
            
            return result
        
        except Exception as e:
            logger.error(f"Error during chain-of-thought reasoning: {str(e)}")
            logger.error(traceback.format_exc())
            raise Exception(f"Chain-of-thought reasoning failed: {str(e)}")


class MathCalculator:
    """Tool for performing mathematical calculations."""
    
    def __init__(self):
        """Initialize the math calculator."""
        self.operations = {
            '+': operator.add,
            '-': operator.sub,
            '*': operator.mul,
            '/': operator.truediv,
            '^': operator.pow,
            '**': operator.pow,
            '%': operator.mod,
            '//': operator.floordiv
        }
    
    def calculate(self, expression: str) -> Dict[str, Any]:
        """
        Evaluate a mathematical expression.
        
        Args:
            expression: The mathematical expression to evaluate
            
        Returns:
            Dictionary containing the result and steps
            
        Raises:
            Exception: If an error occurs during calculation
        """
        
        try:
            cleaned_expr = self._clean_expression(expression)
            
            result = float(sympify(cleaned_expr).evalf())
            
            if result.is_integer():
                result = int(result)
            
            return {
                "expression": expression,
                "result": result,
                "steps": [f"Evaluated expression: {cleaned_expr}", f"Result: {result}"]
            }
        
        except Exception as e:
            logger.error(f"Error calculating expression: {str(e)}")
            logger.error(traceback.format_exc())
            raise Exception(f"Mathematical calculation failed: {str(e)}")
    
    def solve_equation(self, equation: str, variable: str = 'x') -> Dict[str, Any]:
        """
        Solve a mathematical equation for a variable.
        
        Args:
            equation: The equation to solve (e.g., "x + 5 = 10")
            variable: The variable to solve for (default: 'x')
            
        Returns:
            Dictionary containing the solution and steps
            
        Raises:
            Exception: If an error occurs during solving
        """
        
        try:
            if '=' in equation:
                left_side, right_side = equation.split('=', 1)
                equation_sympy = f"({left_side.strip()}) - ({right_side.strip()})"
            else:
                equation_sympy = equation
            
            var = symbols(variable)
            
            expr = sympify(equation_sympy)
            solutions = solve(expr, var)
            
            formatted_solutions = []
            for sol in solutions:
                if sol.is_real:
                    try:
                        float_sol = float(sol)
                        if float_sol.is_integer():
                            formatted_solutions.append(int(float_sol))
                        else:
                            formatted_solutions.append(float_sol)
                    except:
                        formatted_solutions.append(str(sol))
                else:
                    formatted_solutions.append(str(sol))
            
            return {
                "equation": equation,
                "variable": variable,
                "solutions": formatted_solutions,
                "steps": [
                    f"Parsed equation: {equation}",
                    f"Rearranged to: {equation_sympy} = 0",
                    f"Solved for {variable}",
                    f"Solutions: {', '.join(map(str, formatted_solutions))}"
                ]
            }
        
        except Exception as e:
            logger.error(f"Error solving equation: {str(e)}")
            logger.error(traceback.format_exc())
            raise Exception(f"Equation solving failed: {str(e)}")
    
    def _clean_expression(self, expression: str) -> str:
        """
        Clean a mathematical expression for evaluation.
        
        Args:
            expression: The expression to clean
            
        Returns:
            Cleaned expression
        """
        expression = expression.strip()
        
        replacements = {
            'sin': 'math.sin',
            'cos': 'math.cos',
            'tan': 'math.tan',
            'log': 'math.log',
            'exp': 'math.exp',
            'sqrt': 'math.sqrt',
            'pi': 'math.pi',
            'e': 'math.e'
        }
        
        for old, new in replacements.items():
            expression = re.sub(r'\b' + old + r'\b', new, expression)
        
        return expression


class FactVerifier:
    """Tool for verifying facts and claims."""
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """
        Initialize the fact verifier.
        
        Args:
            config: Optional configuration dictionary
        """
        self.model_config = config or get_model_config()
        self.model = ChatOpenAI(
            model=self.model_config.get("reasoning_model", "gpt-4o"),
            temperature=self.model_config.get("temperature", 0.1),
            max_tokens=self.model_config.get("max_tokens", 4096)
        )
    
    def verify(self, claim: str, evidence: Optional[str] = None) -> Dict[str, Any]:
        """
        Verify a claim against evidence.
        
        Args:
            claim: The claim to verify
            evidence: Optional evidence to use for verification
            
        Returns:
            Dictionary containing verification results
            
        Raises:
            Exception: If an error occurs during verification
        """
        
        try:
            prompt_template = """You are an expert fact checker.
            
            Your task is to verify the following claim based on the provided evidence.
            
            Claim: {claim}
            
            Evidence: {evidence}
            
            Analyze the claim and evidence carefully. Provide your verification in the following JSON format:
            {{
                "is_verifiable": true/false,  // Whether the claim can be verified with the given evidence
                "verification_result": "supported"/"contradicted"/"insufficient_evidence",
                "confidence": 0.95,  // A number between 0 and 1 indicating your confidence
                "reasoning": "Your step-by-step reasoning process",
                "missing_information": "What additional information would help verify this claim"
            }}
            
            JSON Response:"""
            
            evidence_text = evidence if evidence else "No evidence provided."
            
            prompt = PromptTemplate.from_template(prompt_template)
            
            chain = prompt | self.model | StrOutputParser()
            
            result = chain.invoke({
                "claim": claim,
                "evidence": evidence_text
            })
            
            parsed_result = json.loads(result)
            
            
            return parsed_result
        
        except Exception as e:
            logger.error(f"Error during fact verification: {str(e)}")
            logger.error(traceback.format_exc())
            raise Exception(f"Fact verification failed: {str(e)}")
    
    def compare_claims(self, claim1: str, claim2: str) -> Dict[str, Any]:
        """
        Compare two claims for consistency.
        
        Args:
            claim1: The first claim
            claim2: The second claim
            
        Returns:
            Dictionary containing comparison results
            
        Raises:
            Exception: If an error occurs during comparison
        """
        
        try:
            prompt_template = """You are an expert at analyzing logical consistency.
            
            Your task is to compare the following two claims and determine if they are consistent, contradictory, or unrelated.
            
            Claim 1: {claim1}
            
            Claim 2: {claim2}
            
            Analyze the claims carefully. Provide your analysis in the following JSON format:
            {{
                "relationship": "consistent"/"contradictory"/"unrelated",
                "confidence": 0.95,  // A number between 0 and 1 indicating your confidence
                "reasoning": "Your step-by-step reasoning process",
                "implications": "What the relationship between these claims implies"
            }}
            
            JSON Response:"""
            
            prompt = PromptTemplate.from_template(prompt_template)
            
            chain = prompt | self.model | StrOutputParser()
            
            result = chain.invoke({
                "claim1": claim1,
                "claim2": claim2
            })
            
            parsed_result = json.loads(result)
            
            
            return parsed_result
        
        except Exception as e:
            logger.error(f"Error during claim comparison: {str(e)}")
            logger.error(traceback.format_exc())
            raise Exception(f"Claim comparison failed: {str(e)}")


def create_step_by_step_reasoner() -> StepByStepReasoner:
    """Create a step-by-step reasoner instance."""
    return StepByStepReasoner()

def create_math_calculator() -> MathCalculator:
    """Create a math calculator instance."""
    return MathCalculator()

def create_fact_verifier() -> FactVerifier:
    """Create a fact verifier instance."""
    return FactVerifier()