from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import List, Optional

from tqdm import tqdm

import instalabel.client
from instalabel.client.models.coco_dataset import CocoDataset
from instalabel.client.models.coco_upload_request import CocoUploadRequest
from instalabel.client.models.project_response import ProjectResponse
from instalabel.client.models.project_update_request import ProjectUpdateRequest
from instalabel.config import INSTALABEL_CLIENT_CONFIGURATION as configuration
from instalabel.core.image import Image


class Project:
    def __init__(self, project: ProjectResponse):
        self.id = project.project_id
        self.name = project.project_name
        self.description = project.project_description
        self.workspace_id = project.workspace_id
        self.task = project.task
        self.created_at = project.created_at
        self.updated_at = project.updated_at
        self.image_count = project.image_count
        self.ontology = project.ontology

    def upload_single_image(self, file_path: str) -> Optional[Image]:
        """Upload a single image to the project"""
        supported_extensions = (".jpg", ".jpeg", ".png")

        # Check if the file format is supported
        if not str(file_path).lower().endswith(supported_extensions):
            print(f"Unsupported file format for {file_path}. Skipping upload.")
            return None

        with instalabel.client.ApiClient(configuration) as api_client:
            api_instance = instalabel.client.ImagesApi(api_client)

            try:
                api_response = api_instance.upload_image_api_v1_images_post(
                    project_id=self.id, file=file_path  # Pass the file path directly
                )
                return Image(api_response)
            except Exception as e:
                print(f"Exception when uploading image: {e}")
                return None

    def upload_images(self, image_paths: List[str]) -> None:
        """Upload multiple images in parallel"""

        def upload(image_path: str):
            self.upload_single_image(image_path)

        # Using ThreadPoolExecutor for parallel uploads
        with ThreadPoolExecutor() as executor:
            list(
                tqdm(
                    executor.map(upload, image_paths),
                    total=len(image_paths),
                    desc="Uploading images",
                )
            )

    def upload_directory(self, dataset_path: str) -> None:
        """Upload all images from a directory"""
        dataset_path = Path(dataset_path)

        # Check if the dataset_path exists
        if not dataset_path.exists():
            print(f"Dataset path {dataset_path} does not exist.")
            return

        # Get all image files from the dataset directory
        supported_extensions = (".jpg", ".jpeg", ".png")
        image_paths = [
            str(file_path)
            for file_path in dataset_path.rglob("*")
            if file_path.suffix.lower() in supported_extensions
        ]

        if not image_paths:
            print(f"No supported images found in {dataset_path}")
            return

        print(f"Found {len(image_paths)} images in the dataset.")
        # Upload all images found in the dataset
        self.upload_images(image_paths)

    def update(
        self,
        project_name: Optional[str] = None,
        project_description: Optional[str] = None,
    ) -> bool:
        """Update project details"""
        if project_name is None and project_description is None:
            print("Please provide either project_name or project_description")
            return False

        with instalabel.client.ApiClient(configuration) as api_client:
            api_instance = instalabel.client.ProjectsApi(api_client)
            update_request = instalabel.client.ProjectUpdateRequest(
                project_name=project_name, project_description=project_description
            )

            try:
                api_response = (
                    api_instance.update_project_api_v1_projects_project_id_put(
                        project_id=self.id, project_update_request=update_request
                    )
                )
                # Update local object with response data
                self.__init__(api_response)
                return True
            except Exception as e:
                print(f"Exception when updating project: {e}")
                return False

    def get_images(self, limit: Optional[int] = None) -> List[Image]:
        """Get all images in the project with pagination support

        Args:
            limit: Optional maximum number of images to return per page. If None, returns all images.

        Returns:
            List[Image]: List of images in the project
        """
        with instalabel.client.ApiClient(configuration) as api_client:
            api_instance = instalabel.client.ImagesApi(api_client)
            images = []
            last_evaluated_key = None

            try:
                # Initialize progress bar
                pbar = tqdm(
                    desc="Fetching images",
                    unit=" images",
                    initial=0,
                    total=limit if limit else self.image_count,
                )

                while True:
                    api_response = api_instance.get_project_images_api_v1_images_get(
                        project_id=self.id,
                        limit=limit,
                        last_evaluated_key=last_evaluated_key,
                    )

                    # Get new images from current page
                    new_images = [Image(image) for image in api_response.images]
                    images.extend(new_images)

                    # Update progress bar with new images
                    pbar.update(len(new_images))

                    # Update description with current count
                    pbar.set_description(f"Retrieved {len(images)} images")

                    # Check if we've reached the last page
                    if not api_response.next_key:
                        break

                    last_evaluated_key = api_response.next_key

                    # If limit is specified, we only want one page
                    if limit is not None:
                        break

                pbar.close()
                return images

            except Exception as e:
                print(f"Exception when getting project images: {e}")
                return []

    def upload_coco_from_s3(self, integration_id: str, coco_data: dict) -> List[Image]:
        """Upload images from an S3 bucket using COCO format annotations.

        Args:
            integration_id (str): AWS integration ID for accessing the S3 bucket
            coco_data (CocoDataset): COCO format annotation data containing image paths and annotations

        Returns:
            List[Image]: List of created images with their details

        Raises:
            Exception: If there's an error during the upload process
        """

        coco_data_dict = CocoDataset(**coco_data).model_dump()

        with instalabel.client.ApiClient(configuration) as api_client:
            api_instance = instalabel.client.ImagesApi(api_client)

            try:
                # Create the upload request
                upload_request = CocoUploadRequest(
                    integration_id=integration_id, coco_data=coco_data_dict
                )

                # Make the API call
                api_response = api_instance.upload_images_with_coco_api_v1_images_external_s3_coco_upload_post(
                    project_id=self.id, coco_upload_request=upload_request
                )

                # Convert API response to Image objects
                return [Image(image) for image in api_response]

            except Exception as e:
                print(f"Exception when uploading COCO dataset: {e}")
                raise

    def __str__(self):
        return (
            f"Project(ID: {self.id}, Name: {self.name}, Task: {self.task.value}, "
            f"Image Count: {self.image_count}, Created At: {self.created_at})"
        )
