import os
import requests
from PIL import Image, ImageDraw
import numpy as np
import json

# Função para enviar requisição para o ChatGPT e obter a resposta
def get_gpt_response(image_path):
    # Dados do formulário
    form = {
        "model": "gpt-3.5-turbo",
        "messages": [
            {
                "role": "user",
                "content": "Identify the weed in this image, returning just label in yolov7 format",
            }
        ]
    }
    formData = json.dumps(form)

    # Abrir a imagem como arquivo binário
    with open(image_path, 'rb') as f:
        image_data = f.read()

    # Configuração da requisição
    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer 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",
    }
    url = "https://api.openai.com/v1/chat/completions"

    # Fazendo a requisição POST
    response = requests.post(url, headers=headers, data=formData)
    
    if response.status_code == 200:
        return response.json()['choices'][0]['message']['content']
    else:
        print(f"Erro ao obter resposta do ChatGPT: {response.status_code}")
        try:
            error_data = response.json()
            print(error_data)
        except ValueError:
            print("Resposta inválida do servidor")
        return None
    
# Função para gerar arquivo de rótulos no formato YOLOv7
def generate_yolov7_label(label_file, class_name, bbox_coords, image_size):
    x_center = bbox_coords[0] + bbox_coords[2] / 2.0
    y_center = bbox_coords[1] + bbox_coords[3] / 2.0
    width = bbox_coords[2]
    height = bbox_coords[3]
    
    x_center /= image_size[0]
    y_center /= image_size[1]
    width /= image_size[0]
    height /= image_size[1]
    
    with open(label_file, 'w') as f:
        f.write(f"{class_name} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")

# Função para desenhar o bbox na imagem
def draw_bbox(image_path, bbox_coords, output_image_path):
    image = Image.open(image_path)
    draw = ImageDraw.Draw(image)
    draw.rectangle(bbox_coords, outline='red', width=3)
    image.save(output_image_path)

# Pasta de entrada e saída
input_folder = 'imagens/'
output_folder = 'saida/'

# Iterar sobre todas as imagens na pasta de entrada
for filename in os.listdir(input_folder):
    if filename.endswith('.jpg') or filename.endswith('.png'):
        image_path = os.path.join(input_folder, filename)
        output_label_file = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.txt")
        output_image_path = os.path.join(output_folder, filename)
        
        # Obter rótulo do ChatGPT
        label = get_gpt_response(image_path)

        print(label)
        
        # Suponha que o ChatGPT retorna algo como "weed: dandelion, bbox: [x, y, width, height]"
        if label.startswith("weed:"):
            parts = label.split(',')
            class_name = parts[0].split(':')[1].strip()
            bbox_coords = list(map(int, parts[1].split(':')[1].strip()[1:-1].split()))
            
            # Gerar arquivo de rótulo YOLOv7
            image_size = Image.open(image_path).size
            generate_yolov7_label(output_label_file, class_name, bbox_coords, image_size)
            
            # Desenhar o bbox na imagem
            draw_bbox(image_path, bbox_coords, output_image_path)
            
            print(f"Processado: {filename}")
        else:
            print(f"Não foi possível identificar a erva em: {filename}")