from smolagents import CodeAgent, InferenceClientModel, WebSearchTool
import os
from dotenv import load_dotenv

load_dotenv()

# Configure the underlying LLM model for the agent
smol_model = InferenceClientModel(token=os.getenv("HF_TOKEN")) # You can choose to not pass any model_id to InferenceClientModel to use a default model
# you can also specify a particular provider e.g. provider="together" or provider="sambanova"

# Instantiate the CodeAgent with authorized imports for file operations
# WebSearchTool is kept as an example, you might want to add/remove tools as needed.
smol_code_agent_instance = CodeAgent(
    tools=[WebSearchTool()],
    model=smol_model,
    add_base_tools=True,
    additional_authorized_imports=["os", "io"] # Authorize os and io for file access
)

# This is the agent class that app.py will instantiate and use.
# It wraps the smolagents.CodeAgent.
class BasicAgent:
    def __init__(self):
        self.smol_agent = smol_code_agent_instance
        print("BasicAgent (using smolagents.CodeAgent) initialized.")

    def __call__(self, question: str, file_path: str | None = None) -> str:
        print(f"BasicAgent received question (first 100 chars): {question[:100]}...")
        
        task_description = question
        run_additional_args = {}

        if file_path:
            print(f"BasicAgent received file_path: {file_path}")
            # Augment the task description to inform the agent about the file
            # and how to access it via additional_args.
            task_description += (
                f"\n\nA file relevant to this task has been provided. "
                f"You can access this file's content using standard Python file operations (e.g., open(), read()). "
                f"The path to this file is available in the 'additional_args' dictionary, passed to the run method, under the key 'task_file'. "
                f"The value is: '{file_path}'."
            )
            run_additional_args["additional_args"] = {"task_file": file_path}
            print(f"Passing to smolagent: task='{task_description[:150]}...', additional_args={run_additional_args['additional_args']}")
        else:
            print("BasicAgent received no file_path.")
            print(f"Passing to smolagent: task='{task_description[:150]}...'")

        try:
            # The CodeAgent's .run() method returns the final answer.
            # Pass additional_args if they exist.
            if run_additional_args:
                answer = self.smol_agent.run(task_description, **run_additional_args)
            else:
                answer = self.smol_agent.run(task_description)
                
            print(f"smolagents.CodeAgent returned: {str(answer)[:200]}...") # Log a snippet of the answer
            return str(answer)  # Ensure the output is a string
        except Exception as e:
            print(f"Error during smolagents.CodeAgent run: {e}")
            # It's often good to return a more informative error message
            return f"AGENT_ERROR: An error occurred while processing with CodeAgent: {str(e)}"

# Example test run (optional, can be commented out for deployment)
# if __name__ == '__main__':
#     print("Testing CodeAgent with a simple task...")
#     test_agent = BasicAgent()
#     # Test without file
#     # response = test_agent("What is the capital of France?")
#     # print(f"Test response (no file): {response}")
#     # To test with a file, you'd create a dummy file and pass its path:
#     # with open("dummy_test_file.txt", "w") as f:
#     #     f.write("This is a test file for the agent.")
#     # response_with_file = test_agent("Summarize the content of the provided file.", file_path="dummy_test_file.txt")
#     # print(f"Test response (with file): {response_with_file}")
#     # os.remove("dummy_test_file.txt")