from PIL import Image
import io
import os
import requests
from crewai.tools import tool
from langchain_community.tools import DuckDuckGoSearchRun

class SupportingTools:
    def __init__(self,API_URL=None):
        self.API_URL = API_URL
        self.hf_api_key = os.getenv("HUGGINGFACE_API_KEY")
        self.headers = {"Authorization": f"Bearer {self.hf_api_key }"}
    
    # API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
    # hugginface_api_key = os.getenv("HUGGINGFACE_API_KEY")
    # headers = {"Authorization": f"Bearer {hugginface_api_key}"}
    def query(self,payload):
        response = requests.post(self.API_URL, headers=self.headers, json=payload)
        return response.content
    
    def generate_image(self,input: str):
        """Create an image based on input"""
        image_bytes = self.query({
            "inputs": input,
        })
        #return image_bytes
        # You can access the image with PIL.Image for example

        image = Image.open(io.BytesIO(image_bytes))
        name = input.split(" ")[0]
        image.save(f"./images/generate_{name}.jpg")
        return image

    @tool("Image Generate Tool")
    def image_generate(prompt: str):
        """
        Args:
           prompt (str): the prompt string return from the classification agent is the argument prompt in the functiom generat image tool.
        """
        API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
        hugginface_api_key = os.getenv("HUGGINGFACE_API_KEY")
        headers = {"Authorization": f"Bearer {hugginface_api_key}"}

        def query(payload):
            response = requests.post(API_URL, headers=headers, json=payload)
            return response.content
        image_bytes = query({
            "inputs": prompt,
        })
        name = prompt.split(" ")[0]
        image = Image.open(io.BytesIO(image_bytes)).resize((1024, 1024))
        image.save(f"./outputs/generate_{name}.jpg")
        import streamlit as st 
        st.write("************ # IMAGE GEN # ********")
        st.image(image)
        return image
    
    
    def duckduckgosearchtool(self):
        search_tool = DuckDuckGoSearchRun()
        return search_tool