import os
import glob
from smolagents import CodeAgent
from prompt.prompt import get_prompt

class QuestionAnsweringAgent(CodeAgent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
    
    def get_current_files(self):
        """Get a set of all files in the current working directory"""
        try:
            # Get all files in current directory (including hidden files)
            all_files = set()
            for pattern in ['*', '.*']:
                all_files.update(glob.glob(pattern))
            # Filter to only include actual files (not directories)
            files = {f for f in all_files if os.path.isfile(f)}
            return files
        except Exception as e:
            print(f"Error getting current files: {e}")
            return set()
    
    def cleanup_created_files(self, files_before):
        """Remove files that were created during execution"""
        try:
            files_after = self.get_current_files()
            newly_created_files = files_after - files_before
            
            for file_path in newly_created_files:
                try:
                    if os.path.exists(file_path):
                        os.remove(file_path)
                        print(f"Cleaned up file: {file_path}")
                except Exception as e:
                    print(f"Error cleaning up file {file_path}: {e}")
                    
        except Exception as e:
            print(f"Error during cleanup: {e}")
        
    def __call__(self, question_text, file_name) -> str:
        # Take snapshot of files before execution
        files_before = self.get_current_files()
        
        try:
            enhanced_question = get_prompt(question_text, file_name)
            response = self.run(enhanced_question, reset=True)
            return response
        finally:
            # Always clean up files, even if there's an exception
            self.cleanup_created_files(files_before)