import torch
from typing import Dict, List, Any, Optional, Union

class ReasoningEngine:
    """Reasoning engine for the autonomous AI agent
    
    This module provides advanced reasoning capabilities including:
    1. Chain-of-thought reasoning
    2. Task decomposition
    3. Self-reflection
    4. Decision making
    """
    
    def __init__(self, model, tokenizer, device="cpu"):
        """Initialize the reasoning engine
        
        Args:
            model: The language model to use for reasoning
            tokenizer: The tokenizer for the language model
            device: The device to run the model on (cpu or cuda)
        """
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
    
    def generate_text(self, prompt: str, max_length: int = 512, temperature: float = 0.7) -> str:
        """Generate text using the language model
        
        Args:
            prompt: The input prompt for the model
            max_length: Maximum length of the generated text
            temperature: Temperature for sampling (higher = more random)
            
        Returns:
            Generated text response
        """
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
        
        # Generate response
        with torch.no_grad():
            outputs = self.model.generate(
                inputs["input_ids"],
                max_length=max_length,
                num_return_sequences=1,
                temperature=temperature,
                top_p=0.9,
                do_sample=True
            )
        
        # Decode and return the response
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response
    
    def chain_of_thought(self, query: str) -> Dict[str, str]:
        """Implement chain-of-thought reasoning
        
        Args:
            query: User query to reason about
            
        Returns:
            Dictionary containing reasoning steps and final answer
        """
        # Construct a prompt that encourages step-by-step reasoning
        reasoning_prompt = f"""I am an autonomous AI agent called ResuRank. I need to think through this step by step to provide the best response.
        
Query: {query}

Let me reason through this carefully:
1. First, I'll identify the key parts of this query.
2. Then, I'll consider what information I need to answer it.
3. Next, I'll analyze the implications and context.
4. Finally, I'll formulate a comprehensive response.

Step-by-step reasoning:
1. """
        
        # Generate the reasoning steps
        reasoning = self.generate_text(reasoning_prompt, max_length=1024)
        
        # Analyze the reasoning for completeness
        if len(reasoning.split('\n')) < 3 or len(reasoning) < 100:
            # If reasoning is too short, try to expand it
            expansion_prompt = f"""My current reasoning is:
{reasoning}

Let me expand on this with more detailed analysis:
"""
            additional_reasoning = self.generate_text(expansion_prompt, max_length=512)
            reasoning = f"{reasoning}\n\n{additional_reasoning}"
        
        # Extract the final answer after reasoning
        answer_prompt = f"""Based on my detailed reasoning:
{reasoning}

I will now provide a clear, helpful, and comprehensive response to the query: {query}

My final answer is:"""
        
        final_answer = self.generate_text(answer_prompt, max_length=768)
        
        # Check if the answer addresses the query adequately
        if len(final_answer) < 50:
            # If answer is too short, try to improve it
            improvement_prompt = f"""My current answer is:
{final_answer}

This answer is too brief. Let me provide a more comprehensive response to the query: {query}

Improved answer:"""
            final_answer = self.generate_text(improvement_prompt, max_length=768)
        
        return {
            "reasoning": reasoning,
            "answer": final_answer
        }
    
    def decompose_task(self, task_description: str) -> List[str]:
        """Decompose a complex task into smaller subtasks
        
        Args:
            task_description: Description of the task to decompose
            
        Returns:
            List of subtask descriptions
        """
        decomposition_prompt = f"""I need to break down this complex task into smaller, manageable subtasks:

Task: {task_description}

Subtasks:
1. """
        
        decomposition = self.generate_text(decomposition_prompt, max_length=1024)
        
        # Parse the decomposition into a list of subtasks
        subtasks = []
        for line in decomposition.split('\n'):
            line = line.strip()
            if line and (line[0].isdigit() or line.startswith('- ')):
                # Remove numbering or bullet points
                cleaned_line = re.sub(r'^\d+\.\s*|^-\s*', '', line).strip()
                if cleaned_line:
                    subtasks.append(cleaned_line)
        
        return subtasks
    
    def self_reflect(self, action: str, outcome: str) -> Dict[str, str]:
        """Perform self-reflection on an action and its outcome
        
        Args:
            action: The action that was taken
            outcome: The outcome of the action
            
        Returns:
            Dictionary containing reflection and improvement suggestions
        """
        reflection_prompt = f"""I need to reflect on this action and its outcome:

Action: {action}
Outcome: {outcome}

Reflection:
"""
        
        reflection = self.generate_text(reflection_prompt, max_length=512)
        
        improvement_prompt = f"""Based on my reflection:
{reflection}

How I can improve next time:
"""
        
        improvement = self.generate_text(improvement_prompt, max_length=512)
        
        return {
            "reflection": reflection,
            "improvement": improvement
        }
    
    def make_decision(self, options: List[str], context: str) -> Dict[str, Any]:
        """Make a decision among multiple options
        
        Args:
            options: List of options to choose from
            context: Context information for the decision
            
        Returns:
            Dictionary containing the chosen option and reasoning
        """
        options_text = "\n".join([f"{i+1}. {option}" for i, option in enumerate(options)])
        
        decision_prompt = f"""I need to make a decision based on the following context and options:

Context: {context}

Options:
{options_text}

Let me analyze each option:
"""
        
        analysis = self.generate_text(decision_prompt, max_length=1024)
        
        conclusion_prompt = f"""Based on my analysis:
{analysis}

The best option is number:"""
        
        conclusion = self.generate_text(conclusion_prompt, max_length=128)
        
        # Try to extract the chosen option number
        try:
            option_num = int(re.search(r'\d+', conclusion).group()) - 1
            chosen_option = options[option_num] if 0 <= option_num < len(options) else options[0]
        except (AttributeError, ValueError, IndexError):
            # Default to first option if parsing fails
            chosen_option = options[0]
        
        return {
            "chosen_option": chosen_option,
            "analysis": analysis,
            "conclusion": conclusion
        }