# main.py
import json
import pandas as pd
from PIL import Image
import cv2
import numpy as np
from aiohttp import ClientSession
from io import BytesIO
import asyncio
import logging
from Data.config import *
from Utils.segment import CropModel
from Utils.sorting import SortModel
from Data.config import frameModel
from Data.config import blankModel
import matplotlib.pyplot as plt

# Set up logging
logging.basicConfig(level=logging.INFO)

class ImageFetcher:
    @staticmethod
    async def fetch_image(url, session):
        try:
            async with session.get(url) as response:
                if response.status == 200:
                    img_data = await response.read()
                    return Image.open(BytesIO(img_data))
                else:
                    logging.error(f"Failed to fetch image from {url}, status code: {response.status}")
                    return None
        except Exception as e:
            logging.error(f"Exception during image fetching from {url}: {e}")
            return None


class ImageProcessor:
    def __init__(self):
        self.frameModel = frameModel
        self.blankModel = blankModel

    async def process_image(self, img_url):
        async with ClientSession() as session:
            image = await ImageFetcher.fetch_image(img_url, session)
            if image is None:
                return {"error": "Failed to fetch image"}

            image = np.array(image)  # Convert PIL image to NumPy array if needed
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image = cv2.resize(image, (frame_shape[0], frame_shape[1]))  # Resize to match the frame shape

            try:
                crop_model = CropModel(model=frameModel)  # Initialize with your model
                padded_image, other_class_name = crop_model.get_predicted_warped_image(image)

                if other_class_name:
                    print(f"Other class name: {other_class_name}")
                else:
                    print("No other class detected.")

                # plt.imshow(padded_image)
                # plt.axis('off')  # Hide axes
                # plt.show() 

                sort_model = SortModel(model=blankModel, classes_to_delete=classes_to_delete, conf_threshold=det_conf, expected_segments=expected_segments) 
                sequence_str = sort_model.process_image(padded_image, other_class_name)
                
                return {"result": sequence_str}
            except Exception as e:
                logging.error(f"Exception during image processing: {e}")
                return {"error": "Image processing failed."}