from typing import Any, Dict, Union
from smolagents.tools import Tool
import re

class ParseKmReadingTool(Tool):
    name = "parse_km_reading"
    description = "Parses kilometer readings from user input."
    inputs = {
        'input_text': {'type': 'string', 'description': 'Text input from the user that may contain kilometer readings'}
    }
    output_type = "any"

    def __init__(self):
        self.is_initialized = True
        # Store the last valid kilometer reading to use as start_km for future trips
        self.last_reading = None

    def forward(self, input_text: str) -> Dict[str, Any]:
        """
        Parse kilometer readings from user input.
        
        Args:
            input_text: Text input that may contain kilometer readings
            
        Returns:
            A dictionary with parsed kilometer readings
        """
        try:
            # Clean and normalize input
            clean_text = input_text.strip().replace(',', '')
            
            # Find all numbers in the input
            numbers = re.findall(r'\b\d+\b', clean_text)
            numbers = [int(num) for num in numbers]
            
            # Determine what kind of reading this is
            result = {
                "success": True,
                "numbers_found": numbers,
                "reading_type": None
            }
            
            if len(numbers) == 0:
                result["success"] = False
                result["error"] = "No kilometer readings found in input"
                return result
            
            if len(numbers) == 1:
                # Single reading - could be start or end
                km_reading = numbers[0]
                
                if self.last_reading is None:
                    # First reading ever, must be start
                    result["reading_type"] = "start"
                    result["start_km"] = km_reading
                    self.last_reading = km_reading
                else:
                    # We have a previous reading
                    if km_reading > self.last_reading:
                        # This is likely an end reading
                        result["reading_type"] = "end"
                        result["end_km"] = km_reading
                        result["start_km"] = self.last_reading
                        result["km_traveled"] = km_reading - self.last_reading
                        self.last_reading = km_reading
                    else:
                        # This is likely a new start reading
                        result["reading_type"] = "start"
                        result["start_km"] = km_reading
                        self.last_reading = km_reading
            
            elif len(numbers) == 2:
                # Two readings - likely start and end
                start_km = min(numbers)
                end_km = max(numbers)
                
                result["reading_type"] = "complete"
                result["start_km"] = start_km
                result["end_km"] = end_km
                result["km_traveled"] = end_km - start_km
                self.last_reading = end_km
            
            else:
                # More than 2 numbers - take first and last assuming chronological order
                result["reading_type"] = "multiple"
                result["start_km"] = numbers[0]
                result["end_km"] = numbers[-1]
                result["km_traveled"] = numbers[-1] - numbers[0]
                result["all_readings"] = numbers
                self.last_reading = numbers[-1]
            
            return result
        
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }