from PIL import Image, ImageEnhance
from escpos.printer import Usb
from picamera2 import Picamera2
import RPi.GPIO as GPIO
import cv2
import numpy as np
import os
import signal
import threading
import time
import io

shots_dir = './shots'
last_image = None
os.makedirs(shots_dir, exist_ok=True)

camera_busy = False
reprint_busy = False

PRINTER_VENDOR_ID = 0x28e9
PRINTER_PRODUCT_ID = 0x0289

GPIO.cleanup()

POWER_LED_PIN = 23
CAMERA_LED_PIN = 4
POWER_BUTTON_PIN = 22
CAMERA_BUTTON_PIN = 27
REPRINT_BUTTON_PIN = 26

GPIO.setmode(GPIO.BCM)
GPIO.setup(POWER_LED_PIN, GPIO.OUT)
GPIO.setup(CAMERA_LED_PIN, GPIO.OUT)
GPIO.setup(POWER_BUTTON_PIN, GPIO.IN, pull_up_down=GPIO.PUD_UP)
GPIO.setup(CAMERA_BUTTON_PIN, GPIO.IN, pull_up_down=GPIO.PUD_UP)
GPIO.setup(REPRINT_BUTTON_PIN, GPIO.IN, pull_up_down=GPIO.PUD_UP)

picam = Picamera2()

camera_config = picam.create_still_configuration({"size": (2000, 2000)})
picam.align_configuration(camera_config)
picam.configure(camera_config)
picam.set_controls({"FrameRate": 1})
picam.start()

GPIO.output(CAMERA_LED_PIN, GPIO.LOW)
GPIO.output(POWER_LED_PIN, GPIO.HIGH)

def get_image_brightness(image):
    grayscale = image.convert('L')
    histogram = grayscale.histogram()
    pixels = sum(histogram)
    brightness = sum(i * count for i, count in enumerate(histogram)) / pixels
    return brightness

def histogram_equalization(image):
    img_gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    equalize_hist = cv2.equalizeHist(img_gray)
    equalized_image = Image.fromarray(equalize_hist)
    return equalized_image

def gamma_correction(image, gamma=0.7):
    img_array = np.array(image) / 255.0
    corrected_array = np.power(img_array, gamma)
    corrected_image = Image.fromarray((corrected_array * 255).astype('uint8'))
    return corrected_image

def apply_clahe(image):
    img_gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    clahe_applied = clahe.apply(img_gray)
    equalized_image = Image.fromarray(clahe_applied)
    return equalized_image

def contrast_stretch(image, low=50, high=200):
    img_array = np.array(image)
    min_val = np.min(img_array)
    max_val = np.max(img_array)
    
    stretched = (img_array - min_val) * ((high - low) / (max_val - min_val)) + low
    stretched = np.clip(stretched, 0, 255)
    stretched_image = Image.fromarray(stretched.astype('uint8'))
    return stretched_image

def camera_button_callback(channel):
    global last_image, camera_busy

    if camera_busy:
        return

    camera_busy = True

    def job():
        global last_image, camera_busy

        try:
            GPIO.output(CAMERA_LED_PIN, GPIO.HIGH)
            timestamp = int(time.time())
            filename = f"shots/image_{timestamp}.jpg"

            picam.capture_file(filename)

            GPIO.output(CAMERA_LED_PIN, GPIO.LOW)

            image = Image.open(filename)
            image = image.resize((576, int(image.height * 576 / image.width)))

            brightness = get_image_brightness(image)

            if brightness < 60:
                image = histogram_equalization(image)
            elif brightness < 90:
                image = gamma_correction(image)
            elif brightness < 110:
                image = apply_clahe(image)
            elif brightness < 130:
                image = gamma_correction(image)
            else:
                image = contrast_stretch(image)

            converted_filename = f"shots/image_{timestamp}_converted.jpg"
            image.save(converted_filename)

            last_image = io.BytesIO()
            image.save(last_image, format='JPEG')
            last_image.seek(0)

            printer = Usb(PRINTER_VENDOR_ID, PRINTER_PRODUCT_ID) 
            printer.image(converted_filename)
            printer.textln()
            printer.textln()
            printer.close()

            os.remove(converted_filename)
        except Exception as e:
            print(f"Error during capture/print: {e}")

        finally:
            camera_busy = False

    threading.Thread(target=job, daemon=True).start()

def power_button_callback(channel):
    os.system('shutdown now')

def reprint_button_callback(channel):
    global last_image, reprint_busy

    if reprint_busy:
        return

    reprint_busy = True

    def job():
        global last_image, reprint_busy

        try:
            if last_image is None:
                return
            
            GPIO.output(CAMERA_LED_PIN, GPIO.HIGH)

            printer = Usb(PRINTER_VENDOR_ID, PRINTER_PRODUCT_ID) 
            printer.image(Image.open(last_image))
            printer.textln()
            printer.textln()
            printer.close()

            GPIO.output(CAMERA_LED_PIN, GPIO.LOW)
        finally:
            reprint_busy = False

    threading.Thread(target=job, daemon=True).start()

def cleanup():
    GPIO.output(CAMERA_LED_PIN, GPIO.HIGH)
    GPIO.cleanup()
    picam.close()

GPIO.add_event_detect(CAMERA_BUTTON_PIN, GPIO.FALLING, callback=camera_button_callback, bouncetime=200)
GPIO.add_event_detect(POWER_BUTTON_PIN, GPIO.FALLING, callback=power_button_callback, bouncetime=200)
GPIO.add_event_detect(REPRINT_BUTTON_PIN, GPIO.FALLING, callback=reprint_button_callback, bouncetime=200)

stop_event = threading.Event()

def worker():
    while not stop_event.is_set():
        time.sleep(1)

def handle_signal(signal_num, frame):
    cleanup()
    stop_event.set()

signal.signal(signal.SIGTERM, handle_signal)
signal.signal(signal.SIGINT, handle_signal)

thread = threading.Thread(target=worker, daemon=True)
thread.start()

stop_event.wait()