import tempfile

import cv2
import torch
from pytube import YouTube


@tool
def web_search(query: str) -> str:
    """Performs a DuckDuckGo search for the given query and returns the results.

    Args:
        query: The search query.

    Returns:
        The top search results as a string.
    """
    try:
        if not query or not query.strip():
            return "Error: Search query cannot be empty"
        
        search_tool = DuckDuckGoSearchRun()
        results = search_tool.invoke(query.strip())
        
        # Clean up the results a bit
        if len(results) > 2000:  # Truncate very long results
            results = results[:2000] + "... (truncated)"
        
        return results
    except Exception as e:
        return f"Error performing web search: {str(e)}"


@tool
def get_max_bird_species_count_from_video(url: str) -> Dict:
    """
    Downloads a YouTube video and returns the maximum number of unique bird species 
    visible in any frame, along with the timestamp.

    Parameters:
        url (str): YouTube video URL

    Returns:
        dict: {
            "max_species_count": int,
            "timestamp": str,
            "species_list": List[str],
        }
    """
    # 1. Download YouTube video
    yt = YouTube(url)
    stream = yt.streams.filter(file_extension='mp4').get_highest_resolution()
    temp_video_path = os.path.join(tempfile.gettempdir(), "video.mp4")
    stream.download(filename=temp_video_path)

    # 2. Load object detection model for bird species
    # Load a fine-tuned YOLOv5 model or similar pretrained on bird species
    model = torch.hub.load('ultralytics/yolov5', 'custom', path='best_birds.pt')  # path to your trained model

    # 3. Process video frames
    cap = cv2.VideoCapture(temp_video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_interval = int(fps * 1)  # 1 frame per second

    max_species_count = 0
    max_species_frame_time = 0
    species_at_max = []

    frame_idx = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        if frame_idx % frame_interval == 0:
            # Run detection
            results = model(frame)
            detected_species = set()
            for *box, conf, cls in results.xyxy[0]:
                species_name = model.names[int(cls)]
                detected_species.add(species_name)

            if len(detected_species) > max_species_count:
                max_species_count = len(detected_species)
                max_species_frame_time = int(cap.get(cv2.CAP_PROP_POS_MSEC)) // 1000
                species_at_max = list(detected_species)

        frame_idx += 1

    cap.release()
    os.remove(temp_video_path)

    return {
        "max_species_count": max_species_count,
        "timestamp": f"{max_species_frame_time}s",
        "species_list": species_at_max
    }