import re
from typing import Any
from smolagents.tools import Tool

class FinalAnswerTool(Tool):
    name = "final_answer"
    description = "Processes and returns the final, concise answer provided by the agent."
    inputs = {'answer': {'type': 'any', 'description': 'The final answer value, potentially structured with sections.'}}
    output_type = "any"

    def forward(self, answer: Any) -> Any:
        """
        Passes the agent's final answer through without modification.
        """
        # Check if answer is a string and contains "FINAL ANSWER:" (case insensitive)
        if isinstance(answer, str):
            answer_text = answer.strip()
            final_answer_prefix = "FINAL ANSWER:"
            
            # Remove any leading/trailing newlines and normalize internal whitespace
            answer_text = '\n'.join(line.strip() for line in answer_text.splitlines()).strip()
            
            # Case insensitive check for prefix anywhere in the text
            if final_answer_prefix.upper() in answer_text.upper():
                # Split on the prefix (case insensitive) and take the last part
                # This handles cases where the prefix might appear multiple times
                parts = re.split(re.escape(final_answer_prefix), answer_text, flags=re.IGNORECASE)
                parsed_answer = parts[-1].strip()
                print(f"[FinalAnswerTool] Extracted answer after prefix: {parsed_answer[:100]}...")
                return parsed_answer

        # For non-string inputs or answers without the prefix, return as-is
        print(f"[FinalAnswerTool] Passing through raw answer: {str(answer)[:100]}...")
        return answer

        # --- PREVIOUS IMPLEMENTATION (COMMENTED OUT) ---
        # """
        # Receives the agent's final answer string, extracts the concise result
        # specifically from the '### 1. Task outcome (short version):' section if present.
        # Falls back to previous behavior or raw input otherwise.
        # """
        # import re # Ensure re is imported if uncommenting
        # if not isinstance(answer, str):
        #     print(f"[FinalAnswerTool] Warning: Input is not a string ('{type(answer)}'). Returning raw value: {str(answer)[:100]}...")
        #     return answer # Return non-strings directly

        # text = answer.strip()
        # original_text_preview = text[:100].replace('\n', '\\n') # For logging

        # # Pattern to capture content after "### 1. ...:" until the next "###" or end of string
        # # Handles variations in spacing and capitalization of the section header.
        # # Makes the "### " prefix optional.
        # # Allows one or more newlines before the next section header.
        # pattern = re.compile(r"^(?:###\s*)?1\.\s*Task outcome \(short version\):([\s\S]*?)(?=\n+(?:###\s*)?2\.|\Z)", re.IGNORECASE | re.MULTILINE)
        # match = pattern.search(text)

        # if match:
        #     # Extract the content, strip leading/trailing whitespace and newlines
        #     parsed_answer = match.group(1).strip()
        #     print(f"[FinalAnswerTool] Extracted from section 1: Raw input: '{original_text_preview}...' -> Parsed output: '{parsed_answer[:100]}...'")
        #     # Return the original full text if extraction results in an empty string (unlikely with [\s\S]*?)
        #     return parsed_answer if parsed_answer else text
        # else:
        #     # Fallback 1: Check for "FINAL ANSWER:" prefix (original behavior)
        #     print(f"[FinalAnswerTool] Info: Section '1. Task outcome (short version):' not found in '{original_text_preview}...'. Trying fallback.")
        #     final_answer_prefix = "FINAL ANSWER:"
        #     # Check from the beginning of the string for the prefix
        #     if text.upper().strip().startswith(final_answer_prefix):
        #         parsed_answer = text.strip()[len(final_answer_prefix):].strip()
        #         parsed_answer = parsed_answer if parsed_answer else text # Avoid returning empty string
        #         print(f"[FinalAnswerTool] Fallback (FINAL ANSWER:): Raw input: '{original_text_preview}...' -> Parsed output: '{parsed_answer[:100]}...'")
        #         return parsed_answer
        #     else:
        #         # Fallback 2: Return the original text if no known format is matched
        #         print(f"[FinalAnswerTool] Warning: Input missing '### 1.' section and 'FINAL ANSWER:' prefix: '{original_text_preview}...'. Returning raw value.")
        #         return text
        # --- END PREVIOUS IMPLEMENTATION ---

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.is_initialized = True