from smolagents import (
    CodeAgent,
    InferenceClientModel,
    VisitWebpageTool,
    WebSearchTool,
    WikipediaSearchTool,
    PythonInterpreterTool,
    FinalAnswerTool,
    tool
)
from groq import Groq
from typing import Dict, Any
import os

# ---- TOOLS ----
@tool
def image_process(image_file: str) -> Dict[str, str]:
    """
    Extract text from an image file using OCR and return OCR text and base64 encoding.

    Args:
        image_file: Path to the image file

    Returns: 
        Dict with keys 'ocr_text' and 'base64_image'
    """
    try:
        import pytesseract
        from PIL import Image
        from smolagents.utils import encode_image_base64
        image = Image.open(image_file)
        base64_img = encode_image_base64(image)  # <<<< CORRECT VARIABLE
        text = pytesseract.image_to_string(image)
        return {
            "ocr_text": text,
            "base64_image": base64_img
        }
    except Exception as e:
        return {
            "ocr_text": "",
            "base64_image": "",
            "error": str(e)
        }

# ---- GROQ MODEL WRAPPER ----
class GroqModel:
    def __init__(self, model_name=""):
        self.model_name = model_name
        self.client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

    def __call__(self, prompt, max_tokens=8096):
        if isinstance(prompt, str):
            messages = [{"role": "user", "content": prompt}]
        else:
            messages = prompt
        response = self.client.chat.completions.create(
            messages=messages,
            model=self.model_name,
            stream=False,
            max_tokens=max_tokens,
        )
        return response.choices[0].message.content

    def generate(self, prompt, max_tokens=8096, **kwargs):
        # For compatibility with agent frameworks
        return self.__call__(prompt, max_tokens=max_tokens)


# ---- MULTI-AGENT SYSTEM ----
class MultyAgentSystem:
    def __init__(self):
        deepseek_model = GroqModel("deepseek-r1-distill-llama-70b")
        qwen_model = GroqModel("qwen-qwq-32b")

        # --- Web agent definition ---
        self.web_agent = CodeAgent(
            model=qwen_model,
            tools=[WebSearchTool(), VisitWebpageTool(), WikipediaSearchTool()],
            name="web_agent",
            description=(
                "You are a web browsing agent. Whenever the given {task} involves browsing "
                "the web or a specific website such as Wikipedia or YouTube, you will use "
                "the provided tools. For web-based factual and retrieval tasks, be as precise and source-reliable as possible."
            ),
            additional_authorized_imports=[
                "markdownify",
                "json",
                "requests",
                "urllib.request",
                "urllib.parse",
                "wikipedia-api",
            ],
            verbosity_level=0,
            max_steps=10,
        )

        # --- Info agent definition ---
        self.info_agent = CodeAgent(
            model=qwen_model,
            tools=[PythonInterpreterTool(), image_process],
            name="info_agent",
            description=(
                "You are an agent tasked with cleaning, parsing, calculating information, and performing OCR if images are provided in the {task}. "
                "You handle all math, code, and data manipulation. Use numpy, math, and available libraries. For image or chess tasks, use pytesseract, PIL, or chess as required."
            ),
            additional_authorized_imports=[
                "numpy",
                "math",
                "pytesseract",
                "PIL",
                "chess",
            ],
        )

        # --- Manager agent definition ---
        self.manager_agent = CodeAgent(
            model=deepseek_model,
            tools=[FinalAnswerTool()],
            managed_agents=[self.web_agent, self.info_agent],
            name="manager_agent",
            description=(
                "You are the manager. Given a {task}, plan which agent to use: "
                "If web data is needed, delegate to web_agent. If math, parsing, or code is needed, use info_agent. "
                "After collecting outputs, optionally cross-validate and check correctness, then finalize and submit the best answer using FinalAnswerTool. "
                "For each task, explicitly explain your planning steps and reasons for choosing which agent, and always prefer the most accurate and complete answer possible."
            ),
            additional_authorized_imports=[
                "json",
                "pandas",
                "numpy",
            ],
            planning_interval=5,
            verbosity_level=2,
            max_steps=20,
        )

    def __call__(self, question, **kwargs):
        
        return self.manager_agent(question, **kwargs)