import cv2
import numpy as np
import tempfile
import sys
import os
if os.name == 'nt':
    if 'USERPROFILE' not in os.environ:
        os.environ['USERPROFILE'] = r'C:\Users\Public'

    if 'USERNAME' not in os.environ:
        os.environ['USERNAME'] = 'sistema_automatico'

    temp_yolo_dir = os.path.join(tempfile.gettempdir(), 'YOLO_Config')
    os.makedirs(temp_yolo_dir, exist_ok=True)
    os.environ['YOLO_CONFIG_DIR'] = temp_yolo_dir
    
    os.environ['TORCH_HOME'] = temp_yolo_dir
from ultralytics import YOLO

# --- FUNÇÕES MATEMÁTICAS (Mantidas) ---
def order_points(pts):
    rect = np.zeros((4, 2), dtype="float32")
    s = pts.sum(axis=1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]
    diff = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]
    return rect

def four_point_transform(image, pts, pad_x_pct=0.0, pad_y_pct=0.0):
    rect = order_points(pts)

    widthA = np.sqrt(((rect[2][0] - rect[3][0]) ** 2) + ((rect[2][1] - rect[3][1]) ** 2))
    widthB = np.sqrt(((rect[1][0] - rect[0][0]) ** 2) + ((rect[1][1] - rect[0][1]) ** 2))
    maxWidth = max(int(widthA), int(widthB))

    heightA = np.sqrt(((rect[1][0] - rect[2][0]) ** 2) + ((rect[1][1] - rect[2][1]) ** 2))
    heightB = np.sqrt(((rect[0][0] - rect[3][0]) ** 2) + ((rect[0][1] - rect[3][1]) ** 2))
    maxHeight = max(int(heightA), int(heightB))

    final_w = maxWidth
    final_h = maxHeight

    dst = np.array([
        [0, 0], 
        [maxWidth - 1, 0], 
        [maxWidth - 1, maxHeight - 1], 
        [0, maxHeight - 1]], dtype="float32")

    M = cv2.getPerspectiveTransform(rect, dst)
    warped = cv2.warpPerspective(image, M, (final_w, final_h), borderMode=cv2.BORDER_REPLICATE)
        
    return warped

# --- LÓGICA PRINCIPAL COM YOLOv8 ---

def main():
    if len(sys.argv) < 2:
        print("ERRO: Você precisa passar o caminho da imagem como parâmetro.")
        sys.exit(1)

    image_path = sys.argv[1]
    output_name = sys.argv[2] if len(sys.argv) > 2 else "documento_recortado.jpg"

    image = cv2.imread(image_path)
    if image is None:
        print(f"ERRO: Não foi possível carregar a imagem '{image_path}'")
        sys.exit(1)

    orig = image.copy()
    h_orig, w_orig = orig.shape[:2]

    # DICA PARA O FUTURO: Quando você tiver um modelo treinado para documentos, 
    # troque 'yolov8n.pt' pelo caminho do seu arquivo, ex: 'meu_modelo_documentos.pt'
    try:
        model = YOLO('yolov8n.pt') 
    except Exception as e:
        print(f"ERRO ao carregar o modelo YOLO: {e}")
        sys.exit(1)

    # conf=0.15: Reduzimos a confiança porque o modelo não é especialista em documentos
    results = model(image, conf=0.15, iou=0.5, verbose=False)

    detected_box = None
    max_area = 0

    for result in results:
        boxes = result.boxes.cpu().numpy()
        for box in boxes:
            cls_id = int(box.cls[0])
            
            # CORREÇÃO CRÍTICA AQUI:
            # Classe 0 = Pessoa (Rosto da foto)
            # Ignoramos a pessoa para que ele não recorte só a foto do RG
            if cls_id == 0:
                continue
            
            x1, y1, x2, y2 = box.xyxy[0].astype(int)
            
            x1, y1 = max(0, x1), max(0, y1)
            x2, y2 = min(w_orig, x2), min(h_orig, y2)
            
            area = (x2 - x1) * (y2 - y1)
            
            # O documento precisa ocupar pelo menos 15% da tela
            if area < (h_orig * w_orig * 0.15):
                continue
            
            if area > max_area:
                max_area = area
                detected_box = np.array([
                    [x1, y1], [x2, y1], [x2, y2], [x1, y2]
                ], dtype="float32")

    if detected_box is not None:
        warped = four_point_transform(orig, detected_box)
        cv2.imwrite(output_name, warped)
        print(f"SUCESSO: Documento detectado pelo YOLO e salvo como {output_name}")
    else:
        # FALLBACK: Se ele ignorou o rosto e não achou mais nada
        print("AVISO: Nenhum documento válido detectado. Salvando imagem original.")
        cv2.imwrite(output_name, orig)

    sys.exit(0)

if __name__ == "__main__":
    main()