from smolagents import DuckDuckGoSearchTool, HfApiModel, load_tool, CodeAgent, PythonInterpreterTool, VisitWebpageTool, \
    Tool
import hashlib
import json
from transformers import AutoTokenizer, AutoModelForCausalLM
import os


class ModelMathTool(Tool):
    name = "math_model"
    description = "Answers advanced math questions using a pretrained math model."

    inputs = {
        "problem": {
            "type": "string",
            "description": "Math problem to solve.",
        }
    }

    output_type = "string"

    def __init__(self, model_id="Qwen/Qwen2.5-Math-7B"):
        print(f"Loading math model: {model_id}")
        self.tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
        self.model = HfApiModel(model_id=model_id, max_tokens=512)

    def forward(self, problem: str) -> str:
        print(f"[MathModelTool] Question: {problem}")
        response = self.model.__call__(problem)
        return response


# (Keep Constants as is)
# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

web_search = DuckDuckGoSearchTool()
python_interpreter = PythonInterpreterTool()
visit_webpage_tool = VisitWebpageTool()
model_math_tool = ModelMathTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud'

model = HfApiModel(model_id="HuggingFaceH4/zephyr-7b-beta", max_tokens=512, token=tok)


def get_cache_key(question: str) -> str:
    return hashlib.sha256(question.encode()).hexdigest()


def load_cached_answer(question: str) -> str | None:
    key = get_cache_key(question)
    path = f"cache/{key}.json"
    if os.path.exists(path):
        with open(path, "r") as f:
            data = json.load(f)
            return data.get("answer")
    return None


def cache_answer(question: str, answer: str):
    key = get_cache_key(question)
    path = f"cache/{key}.json"
    with open(path, "w") as f:
        json.dump({"question": question, "answer": answer}, f)


class BasicAgent:
    def __init__(self):
        print("BasicAgent initialized.")
        self.agent = CodeAgent(
            model=model,
            tools=[model_math_tool],
            max_steps=1,
            verbosity_level=0,
            grammar=None,
            planning_interval=3,

        )

    def __call__(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        answer = self.agent.run(question)
        return answer



agent = BasicAgent()

response = agent.__call__(question="How much is 2*2?")
print(response)