from smolagents import Tool
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
import torch
from wikipedia_utils import *
from youtube_utils import *


class MathModelQuerer(Tool):
    name = "math_model"
    description = "Solves advanced math problems using a pretrained\
    large language model specialized in mathematics. Ideal for symbolic reasoning, \
    calculus, algebra, and other technical math queries."

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

    output_type = "string"

    def __init__(self, model_name="deepseek-ai/deepseek-math-7b-base"):
        print(f"Loading math model: {model_name}")

        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        print("loaded tokenizer")
        self.model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
        print("loaded auto model")

        self.model.generation_config = GenerationConfig.from_pretrained(model_name)
        print("loaded coonfig")

        self.model.generation_config.pad_token_id = self.model.generation_config.eos_token_id
        print("loaded pad token")

    def forward(self, problem: str) -> str:
        try:
            print(f"[MathModelTool] Question: {problem}")

            inputs = self.tokenizer(problem, return_tensors="pt")
            outputs = self.model.generate(**inputs, max_new_tokens=100)

            result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)

            return result
        except:
            return f"Failed using the tool {self.name}"


class CodeModelQuerer(Tool):
    name = "code_querer"
    description = "Generates code snippets based on a natural language description of a\
    programming task using a powerful coding-focused language model. Suitable\
    for solving coding problems, generating functions, or implementing algorithms."

    inputs = {
        "problem": {
            "type": "string",
            "description": "Description of a code sample to be generated",
        }
    }

    output_type = "string"

    def __init__(self, model_name="Qwen/Qwen2.5-Coder-32B-Instruct"):
        from smolagents import HfApiModel
        print(f"Loading llm for Code tool: {model_name}")
        self.model = HfApiModel()

    def forward(self, problem: str) -> str:
        try:
            return self.model.generate(problem, max_new_tokens=512)
        except:
            return f"Failed using the tool {self.name}"


class WikipediaPageFetcher(Tool):
    name = "wiki_page_fetcher"
    description =' Searches and fetches summaries from Wikipedia for any topic,\
    across all supported languages and versions. Only a single query string is required as input.'



    inputs = {
        "query": {
            "type": "string",
            "description": "Topic of wikipedia search",
        }
    }

    output_type = "string"

    def forward(self, query: str) -> str:
        try:
            wiki_query = query(query)
            wiki_page = fetch_wikipedia_page(wiki_query)
            return wiki_page
        except:
            return f"Failed using the tool {self.name}"


class YoutubeTranscriptFetcher(Tool):
    name = "youtube_transcript_fetcher"
    description ="Fetches the English transcript of a YouTube video using either a direct video \
    ID or a URL that includes one. Accepts a query containing the link or the raw video ID directly. Returns the transcript as plain text."

    inputs = {
        "query": {
            "type": "string",
            "description": "A query that includes youtube id."
        },
        "video_id" : {
            "type" : "string",
            "description" : "Optional string with video id from youtube.",
            "nullable"  : True
        }
    }

    output_type = "string"

    def forward(self, query: str, video_id=None) -> str:
        try:
            if video_id is None:
                video_id = get_youtube_video_id(query)

            fetched_transcript = fetch_transcript_english(video_id)

            return post_process_transcript(fetched_transcript)
        except:
            return f"Failed using the tool {self.name}"