from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from qwen_vl_utils import process_vision_info
from smolagents import tool
import torch


@tool
def video_reasoner(file_path : str, query : str) -> str: 
    """
    This tool performs requested visual reasoning task on the provided video and returns the generated output.

    Args:
        file_path: Path of a local video file on which visual reasoning is to be done.
        query: visual reasoning that is to be done.
    """
    try:
        # default: Load the model on the available device(s)
        model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            "Qwen/Qwen2.5-VL-7B-Instruct", torch_dtype="auto", device_map="auto"
        )

        # default processer
        processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")

        messages = [
        {
        "role": "user",
        "content": [
            {
                "type": "video",
                "video": file_path,
                "max_pixels": 360 * 360,
                "fps": 0.3,
            },
            {"type": "text", "text": f"{query}\n\nAdditional instruction: Treat the two types of penguins as distinct species e.g. Adelie and Emperor Penguins are considered two different species of birds."},
            ],
        }      
        ]

        # Preparation for inference
        text = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        image_inputs, video_inputs = process_vision_info(messages)
        inputs = processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        inputs = inputs.to("cuda")

        # Inference: Generation of the output
        generated_ids = model.generate(**inputs, max_new_tokens=512)
        generated_ids_trimmed = [
            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]
        output_text = processor.batch_decode(
            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
        )

        import gc
        
        # After inference
        del image_inputs
        del video_inputs
        del inputs
        del model
        del processor
        gc.collect()                 # Force Python garbage collection
        torch.cuda.empty_cache()     # Clear cached memory
        
        return output_text
    
    except Exception as e:
        return f'error occured: {e}'
    

@tool
def image_reasoner(file_path : str, query : str) -> str: 
    """
    This tool performs requested visual reasoning task on the provided image and returns the generated output.

    Args:
        file_path: Path of a local image file on which visual reasoning is to be done.
        query: visual reasoning that is to be done.
    """
    try:
        # default: Load the model on the available device(s)
        model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            "Qwen/Qwen2.5-VL-7B-Instruct", torch_dtype="auto", device_map="auto"
        )

        # default processer
        processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")

        messages = [
        {
        "role": "user",
        "content": [
            {
                "type": "image",
                "image": file_path,
            },
            {"type": "text", "text": f"{query}\n\nAdditional instruction: Review your answer for correctness."},
            ],
        }      
        ]

        # Preparation for inference
        text = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        image_inputs, video_inputs = process_vision_info(messages)
        inputs = processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        inputs = inputs.to("cuda")

        # Inference: Generation of the output
        generated_ids = model.generate(**inputs, max_new_tokens=512)
        generated_ids_trimmed = [
            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]
        output_text = processor.batch_decode(
            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
        )

        import gc
        
        # After inference
        del image_inputs
        del video_inputs
        del inputs
        del model
        del processor
        gc.collect()                 # Force Python garbage collection
        torch.cuda.empty_cache()     # Clear cached memory
        
        return output_text
    
    except Exception as e:
        return f'error occured: {e}'