from smolagents import CodeAgent, DuckDuckGoSearchTool, load_tool, tool, LiteLLMModel
import requests
import wikipedia
import openpyxl
import subprocess
import os

# ---------------------------
# 1. Base model (local Qwen)
# ---------------------------
# Example: If you've downloaded Qwen locally at ./models/qwen
# Use HfLocalModel for local inference
# model = HfLocalModel(
#     model_id="./models/Qwen2.5-Coder-14B-Instruct",  # path to local Qwen
#     max_tokens=2048,
#     temperature=0.3,
# )

model = LiteLLMModel(
    model_id="ollama_chat/qwen2:7b",  # Or try other Ollama-supported models
    api_base="http://127.0.0.1:11434",  # Default Ollama local server
    num_ctx=8192,
)

# ---------------------------
# Local Tools
# ---------------------------

@tool
def local_web_search(query: str, num_results: int = 5) -> list:
    """
    Perform a simple web search using DuckDuckGo.

    Args:
        query (str): The search query string.
        num_results (int): Number of results to return (default = 5).

    Returns:
        list: A list of dictionaries containing 'title' and 'url' for each result.
    """
    from duckduckgo_search import DDGS
    results = []
    with DDGS() as ddgs:
        for r in ddgs.text(query, max_results=num_results):
            results.append({"title": r.get("title"), "url": r.get("href")})
    return results


@tool
def local_wikipedia_search(query: str, sentences: int = 2) -> str:
    """
    Search and summarize a Wikipedia article.

    Args:
        query (str): The topic to search on Wikipedia.
        sentences (int): Number of sentences in the summary (default = 2).

    Returns:
        str: A short summary of the topic from Wikipedia.
    """
    try:
        return wikipedia.summary(query, sentences=sentences)
    except Exception as e:
        return f"Error fetching summary: {str(e)}"


# @tool
# def local_image_caption(image_path: str) -> str:
#     """
#     Generate a dummy caption for an image (placeholder).

#     Args:
#         image_path (str): Path to the image file.

#     Returns:
#         str: Caption describing the image.
#     """
#     # ⚠️ Replace with real model if available (BLIP, CLIP, etc.)
#     return f"Caption for image at {image_path}: [Image captioning not implemented]."


@tool
def local_audio_transcribe(audio_path: str) -> str:
    """
    Transcribe speech from an audio file using Whisper (requires whisper installed).

    Args:
        audio_path (str): Path to the audio file (e.g., .mp3, .wav).

    Returns:
        str: Transcribed text from the audio.
    """
    try:
        import whisper
        model = whisper.load_model("base")
        result = model.transcribe(audio_path)
        return result["text"]
    except Exception as e:
        return f"Error transcribing audio: {str(e)}"


@tool
def local_python_runner(code: str) -> str:
    """
    Execute a Python script safely.

    Args:
        code (str): Python code to execute.

    Returns:
        str: The output or error message from execution.
    """
    try:
        result = subprocess.run(
            ["python3", "-c", code],
            capture_output=True,
            text=True,
            timeout=10
        )
        return result.stdout if result.stdout else result.stderr
    except Exception as e:
        return f"Execution error: {str(e)}"


@tool
def local_excel_reader(file_path: str) -> float:
    """
    Read an Excel file and compute the sum of all numeric values.

    Args:
        file_path (str): Path to the Excel file (.xlsx).

    Returns:
        float: The sum of all numeric values in the file.
    """
    try:
        workbook = openpyxl.load_workbook(file_path)
        total_sum = 0
        for sheet in workbook.worksheets:
            for row in sheet.iter_rows():
                for cell in row:
                    if isinstance(cell.value, (int, float)):
                        total_sum += cell.value
        return total_sum
    except Exception as e:
        return f"Error reading Excel file: {str(e)}"


@tool
def check_commutativity(elements: list, table: dict) -> str:
    """
    Check for non-commutativity in a given operation table.

    Args:
        elements (list): List of elements in the operation.
        table (dict): Operation table as a nested dictionary
                      (e.g., table[a][b] = result of a*b).

    Returns:
        str: Comma-separated elements that violate commutativity.
    """
    counterexample_set = set()
    for a in elements:
        for b in elements:
            if table[a][b] != table[b][a]:
                counterexample_set.update([a, b])
    return ",".join(sorted(counterexample_set))


# ---------------------------
# 3. Build Agent
# ---------------------------
agent = CodeAgent(
    model=model,
    tools=[
        DuckDuckGoSearchTool(),
        local_wikipedia_search,
        # local_image_caption,
        local_audio_transcribe,
        local_python_runner,
        local_excel_reader,
        check_commutativity,
    ],
    add_base_tools=True,
    max_steps=8,
    planning_interval=3,
    verbosity_level=2,
)

# ---------------------------
# 4. Questions dataset
# ---------------------------

import requests

url = "https://agents-course-unit4-scoring.hf.space/questions"

headers = {
    "accept": "application/json"
}

response = requests.get(url, headers=headers)

if response.status_code == 200:
    tasks = response.json()
    print("✅ Response JSON:", tasks)
else:
    print(f"❌ Failed with status code {response.status_code}")
    print(response.text)


# ---------------------------
# 5. Run Agent and collect results
# ---------------------------
results = {
    "username": "ginnigarg",
    "agent_code": "ginniAgent_v1",
    "answers": []
}

for task in tasks:
    try:
        answer = agent.run(task["question"])
    except Exception as e:
        answer = f"Error: {str(e)}"
    results["answers"].append({
        "task_id": task["task_id"],
        "submitted_answer": str(answer)
    })

# ---------------------------
# 6. Print final JSON
# ---------------------------
import json
print(json.dumps(results, indent=2))