from fastapi import UploadFile
from sqlalchemy.ext.asyncio import AsyncSession
from src.components.input_images import SavingInputImage
from src.components.comparison import ImageComparator
from src.components.image_comparison import ImageComparison
candidate_labels = [
    "a stack of wooden logs",
    "a concrete mixer machine",
    "metal rebar rods",
    "construction scaffolding",
    "a construction door",
    "a pile of bricks",
    "metal shuttering panels"
]


class PredictionPipeline:
    def __init__(self, db: AsyncSession, id: int):
        self.db = db
        self.id = id

    async def run_pipeline(self, before_image: UploadFile, after_image: UploadFile):
        # Step 1: Save input images
        saver = SavingInputImage(db=self.db, id = self.id)
        save_result = await saver.saving_input_images(before_image, after_image)

        prediction_id = save_result["prediction_id"]
        before_path = save_result["before_path"]  # Add this to return from saving component
        after_path = save_result["after_path"]    # Add this to return from saving component

        # Step 2: Run comparison logic
        comparator = ImageComparator(db=self.db)
        comparison_result = await comparator.run_comparison(
        prediction_id=prediction_id,
        before_path=before_path,
        after_path=after_path
        )


        comaprator2 = ImageComparison(db=self.db)
        comparison_result2 = await comaprator2.compare_and_annotate(
            before_path=before_path,
            after_path=after_path,
            candidate_labels=candidate_labels,
            prediction_id = prediction_id,
            iou_threshold=0.2,
            score_threshold=0.3
        )



        return {
            "message": "Pipeline executed successfully",
            "prediction_id": prediction_id,
            "before_url": save_result["before_url"],
            "after_url": save_result["after_url"],
            "text_results": comparison_result["comparison_report"],
            "status": "completed",
            "output_image_url": comparison_result2
        }