import cv2
import numpy as np
import sys
import json
from pathlib import Path
import os
from PIL import Image

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 process_document(image_path, doc_type, output_dir):
    try:
        # Carrega e verifica a imagem
        image = cv2.imread(str(image_path))
        if image is None:
            raise Exception("Não foi possível carregar a imagem")

        original = image.copy()
        
        # Pré-processamento
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        edges = cv2.Canny(blurred, 50, 150)

        # Encontra contornos
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if not contours:
            raise Exception("Nenhum contorno encontrado")

        # Ordena contornos por área
        contours = sorted(contours, key=cv2.contourArea, reverse=True)
        doc_corners = None

        # Procura o documento
        for contour in contours:
            peri = cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, 0.02 * peri, True)
            if len(approx) == 4:
                doc_corners = approx
                break

        if doc_corners is None:
            raise Exception("Documento não encontrado")

        # Processa cantos e aplica transformação
        doc_corners = doc_corners.reshape(4, 2)
        ordered_corners = order_points(doc_corners)

        (tl, tr, br, bl) = ordered_corners
        widthA = np.linalg.norm(br - bl)
        widthB = np.linalg.norm(tr - tl)
        maxWidth = max(int(widthA), int(widthB))

        heightA = np.linalg.norm(tr - br)
        heightB = np.linalg.norm(tl - bl)
        maxHeight = max(int(heightA), int(heightB))

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

        # Aplica transformação perspectiva
        matrix = cv2.getPerspectiveTransform(ordered_corners, dst)
        warped = cv2.warpPerspective(original, matrix, (maxWidth, maxHeight))

        # Salva resultado
        output_path = os.path.join(output_dir, f"{doc_type}_{Path(image_path).name}")
        cv2.imwrite(str(output_path), warped)
        
        return {
            "success": True,
            "message": "Processamento concluído com sucesso",
            "width": maxWidth,
            "height": maxHeight
        }

    except Exception as e:
        return {
            "success": False,
            "message": str(e)
        }

if __name__ == "__main__":
    if len(sys.argv) != 4:
        print(json.dumps({"success": False, "message": "Uso: python document_detector.py <input_path> <doc_type> <output_dir>"}))
        sys.exit(1)

    result = process_document(Path(sys.argv[1]), sys.argv[2], Path(sys.argv[3]))
    print(json.dumps(result))