from fastapi import FastAPI
from nicegui import ui, events, run
import base64
import io
import numpy as np
import cv2
import traceback
import logging
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s [%(levelname)s] %(message)s', handlers=[logging.StreamHandler()])
from cv_processor import process

from pydantic import BaseModel

class Item(BaseModel):
    image: str # Base64

app = FastAPI()

@app.post('/process')
def process_image(item: Item):
    logging.info("POST /process")
    try:
        image_data = item.image
        if (image_data is None):
            return {"error": "No image data provided"}

        # Decode base64 string
        image_bytes = base64.b64decode(image_data)
        image_stream = io.BytesIO(image_bytes)
        # Convert bytes to NumPy array
        img_array = np.frombuffer(image_stream.getvalue(), dtype=np.uint8)
        # Decode image using OpenCV
        img_bgr = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
        # Valid image
        assert(img_bgr is not None)

        # Process the image
        results = process(img_bgr)

        # Encode processed image to base64
        _, buffer = cv2.imencode('.jpg', results.get("image"), [cv2.IMWRITE_JPEG_QUALITY, 100])
        processed_image_base64 = base64.b64encode(buffer).decode('utf-8')

        # Update image with base64 encoded
        results["image_b64"] = processed_image_base64
        # Pop image (not serializable)
        results.pop("image")
        return results
    
    except Exception as e:
        logging.warning("Exception: {}".format(traceback.format_exc()))
        return {"error": traceback.format_exc()}

# Define the NiceGUI UI components
@ui.page("/")
def main_page():

    async def handle_upload(e: events.UploadEventArguments) -> None:
        ui.notify('Processing...')
        # Read content -> image
        nparr = np.frombuffer(e.content.read(), np.uint8)
        img_np_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        # Async process
        results = await run.io_bound(process, img_np_bgr)

        # Display
        with ui.dialog() as dialog:
            # Encode
            retval, buffer = cv2.imencode('.png', results.get("image"))
            img_buffer_encoded = base64.b64encode(buffer).decode('utf-8')
            img_encoded = "data:image/png;base64,{}".format(img_buffer_encoded)
            content = ui.image(img_encoded).props('fit=scale-down')
        dialog.open()

    ui.upload(on_upload=handle_upload, auto_upload=True, on_rejected=lambda: ui.notify('Rejected!')).props('accept=image').classes('max-w-full')

if __name__ == '__main__':
    ui.run_with(app, title="CV")