import os
from smolagents import CodeAgent, ToolCallingAgent
from smolagents import OpenAIServerModel
from tools.fetch import fetch_webpage
from tools.yttranscript import get_youtube_transcript, get_youtube_title_description
import myprompts
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch
# --- Basic Agent Definition ---
class BasicAgent:
    def __init__(self):
        print("BasicAgent initialized.")
    def __call__(self, question: str) -> str:

        print(f"Agent received question (first 50 chars): {question[:50]}...")

        try:
            # Use the reviewer agent to determine if the question can be answered by a model or requires code
            print("Calling reviewer agent...")
            reviewer_answer = reviewer_agent.run(myprompts.review_prompt + "\nThe question is:\n" + question)
            print(f"Reviewer agent answer: {reviewer_answer}")

            question = question + '\n' + myprompts.output_format
            fixed_answer = ""

            if reviewer_answer == "code":
                fixed_answer = gaia_agent.run(question)
                print(f"Code agent answer: {fixed_answer}")
                
            elif reviewer_answer == "model":    
                # If the reviewer agent suggests using the model, we can proceed with the model agent
                print("Using model agent to answer the question.")
                fixed_answer = model_agent.run(myprompts.model_prompt + "\nThe question is:\n" + question)
                print(f"Model agent answer: {fixed_answer}")

            return fixed_answer
        except Exception as e:
            error = f"An error occurred while processing the question: {e}"
            print(error)
            return error

        
      
    
# model = OpenAIServerModel(
#     model_id="gpt-4.1-nano",
#     api_base="https://api.openai.com/v1",
#     api_key=os.environ["OPENAI_API_KEY"],
# )


MODEL_NAME = "meta-llama/Llama-3.2-3B"  # 3B isn't released by Meta officially, but use 8B or a 3B variant like TinyLlama if needed

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model_init = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",
    torch_dtype=torch.float16  # or bfloat16
)

def model(prompt: str, max_new_tokens=512):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
    output_ids = model.generate(input_ids, max_new_tokens=max_new_tokens)
    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return output

reviewer_agent= ToolCallingAgent(model=model, tools=[])
model_agent = ToolCallingAgent(model=model,tools=[fetch_webpage])
gaia_agent = CodeAgent(tools=[fetch_webpage,get_youtube_title_description,get_youtube_transcript ], model=model)

if __name__ == "__main__":
    # Example usage
    question = "What was the actual enrollment of the Malko competition in 2023?"
    agent = BasicAgent()
    answer = agent(question)
    print(f"Answer: {answer}")