import tensorflow as tf
from tensorflow.keras.models import load_model
from imutils import paths
import numpy as np
import cv2
import pickle
from tratar_captcha import tratar_base64
import base64

tamanho_imagem = (35, 50)

# Posições fixas para o corte de cada caractere
posicoes_corte = [(6, 40), (56, 90), (111, 144), (161, 197), (214, 250), (265, 297)]

def quebrar_captcha_base64(base64_string):
    if base64_string.startswith('data:image/png;base64,'):
        base64_string = base64_string.replace('data:image/png;base64,', '')
    
    with open("rotulos_modelo.dat", "rb") as arquivo_tradutor:
        lb = pickle.load(arquivo_tradutor)
    
    modelo = load_model("modelo_treinado.keras")
    imagem_tratada = tratar_base64(base64_string)

    # Aplicar um limiar para binarizar a imagem
    _, imagem_bin = cv2.threshold(imagem_tratada, 128, 255, cv2.THRESH_BINARY_INV)

    imagem_final = cv2.merge([imagem_bin] * 3)  # Convertendo de volta para RGB para desenhar os retângulos coloridos

    # Processar cada região de corte definida
    previsao = []
    for x1, x2 in posicoes_corte:
        imagem_caractere = imagem_bin[20:-15, x1:x2]  # Cortar a região do caractere
        imagem_caractere = cv2.bitwise_not(imagem_caractere)
        
        imagem_caractere = cv2.resize(imagem_caractere, tamanho_imagem)

        imagem_caractere = np.expand_dims(imagem_caractere, axis=2)
        imagem_caractere = np.expand_dims(imagem_caractere, axis=0)

        letra_prevista = modelo.predict(imagem_caractere)
        letra_prevista = lb.inverse_transform(letra_prevista)[0]
        previsao.append(letra_prevista)

    texto_previsao = "".join(previsao)

    print(texto_previsao)

    return texto_previsao


# Exemplo de uso
if __name__ == "__main__":
    # Aqui você deve substituir por uma string base64 real
    base64_string = "data:image/png;base64, [encoded data, 7600 characters]"
    
    quebrar_captcha_base64(base64_string)
