import os
from datetime import datetime
from sqlalchemy.ext.asyncio import AsyncSession
from src.models.prediction import Prediction
from src.utils.files_utils import save_input_images_async
from src.common.logger import logger
from fastapi import UploadFile

class SavingInputImage:
    def __init__(self, db: AsyncSession, id: int):
        self.db = db
        self.id = id    
    async def saving_input_images(self, before_image: UploadFile, after_image: UploadFile):
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        folder_path = os.path.join("Input_Images", timestamp)
        try:
            os.makedirs(folder_path, exist_ok=True)
            logger.info(f"Created folder for input images: {folder_path}")

            before_path = os.path.join(folder_path, "before.jpg")
            after_path = os.path.join(folder_path, "after.jpg")

            try:
                with open(before_path, "wb") as bf:
                    bf.write(await before_image.read())
                logger.info(f"Saved 'before' image at: {before_path}")
            except Exception as e:
                logger.exception("Failed to save 'before' image.")
                raise Exception("Failed to save 'before' image.") from e

            try:
                with open(after_path, "wb") as af:
                    af.write(await after_image.read())
                logger.info(f"Saved 'after' image at: {after_path}")
            except Exception as e:
                logger.exception("Failed to save 'after' image.")
                raise Exception("Failed to save 'after' image.") from e

            try:
                before_url, after_url = await save_input_images_async(timestamp)
                logger.info(f"Generated URLs - Before: {before_url}, After: {after_url}")
            except Exception as e:
                logger.exception("Failed to generate image URLs.")
                raise Exception("Failed to generate image URLs.") from e
            id = self.id
            try:
                prediction = Prediction(
                    user_id = id,
                    before_image_url=before_url,
                    after_image_url=after_url,
                    text_results={},
                    output_image_url="",
                    status="pending"
                )
                self.db.add(prediction)
                await self.db.commit()
                await self.db.refresh(prediction)
                logger.info(f"Saved prediction with ID: {prediction.id}")
            except Exception as e:
                logger.exception("Failed to save prediction to the database.")
                raise Exception("Database operation failed.") from e

            return {
                "message": "Images uploaded successfully!",
                "prediction_id": prediction.id,
                "before_url": before_url,
                "after_url": after_url,
                "before_path": before_path,  
                "after_path": after_path     
            }


        except Exception as e:
            logger.exception("Error occurred while processing input images.")
            raise Exception(f"Error while saving images: {str(e)}")