"""Image extraction application on modal."""

from concurrent.futures import ThreadPoolExecutor
import time
import modal
from openai import OpenAI
from openai.types.chat import ChatCompletionMessageParam, ChatCompletionContentPartParam
from suno_utils.worker.modal_base import get_modal_base_image
from suno_utils.worker.settings import s3_client
from suno_utils.worker.schema import QueueItem
import json
from suno_utils.utils.image_extraction import (
    does_image_pass_moderation,
    ImageModerationFailure,
    ImageExtractionError,
)

# setup Modal stub (i.e. App)
DEPLOYMENT_TYPE = "dev"  # priority, staging, dev, prod, msft
APP_NAME = f"image-extraction-{DEPLOYMENT_TYPE}"
assert APP_NAME.endswith(DEPLOYMENT_TYPE)

# define image here...
SECRETS = [
    modal.Secret.from_name("openai-secret"),
    modal.Secret.from_name("studio-aws"),
    modal.Secret.from_name("hive-secret"),
]
baseImage = get_modal_base_image().add_local_python_source("suno_utils", copy=False)
app = modal.App(APP_NAME, image=baseImage)


@app.cls(
    secrets=SECRETS,
    min_containers=1,
    retries=modal.Retries(
        max_retries=2,
        backoff_coefficient=2.0,
        initial_delay=5.0,  # retries is why you need to return errors quickly
    ),
    cloud="aws",
    region="us-east",
)
@modal.concurrent(max_inputs=4)
class ImageExtractionStub:
    """Get text from image using OpenAI's vision library"""

    def __init__(self):
        """Set up ImageExtractionStub."""
        self.openai_client = OpenAI()
        self.system_prompt = """You are an artist, curiously examining weird and interesting details that pop out at you from the image.

            If the image is a cartoon or a screenshot of text (e.g. conversation, meme, tweet, cartoon, etc.): Place **Transcription:** at the beginning of your response, and then transcribe all text, with no other additions. 

            Else, be sure to include the following in no more than 100 words:
                1. Main subject and setting, including transcribing any text and user context (1-2 sentences)
                2. Mood and unique elements (1-2 sentences)
                3. Any pop culture references (be specific, 1-2 sentences)
        """
        self.user_prompt = """Analyze the image and provide a concise description[INSERT USER_CONTEXT HERE].
        
            If the image is a cartoon or a screenshot of text (e.g. conversation, meme, tweet, cartoon, etc.):
            1. Transcribe the text content.

            Else, describe:
            1. Main subject and setting, including transcribing any text (1-2 sentences)
            2. Mood and unique elements (1-2 sentences)
            3. Potential lyrical themes (2-3 key words or phrases)

            Focus on elements that could inspire creative, potentially humorous lyrics. If you catch your self describing an image as a screenshot, stop describing the image and just use OCR to transcribe the full text immediately.
            
            Number of images: [INSERT NUMBER]
            """
        self.multiple_image_suffix = """After describing each image individually, 
            provide a brief summary of any overall themes, narrative, or emotional journey presented by the full set of images. 
            Highlight any contrasts, similarities, or progression across the images."""

        self.single_image_user_prompt = """Analyze the image and provide a concise description[INSERT USER_CONTEXT HERE].
        
            If the image is a cartoon or a screenshot of text (e.g. conversation, meme, tweet, cartoon, etc.):
            1. Transcribe the text content.

            Else, describe:
            1. Main subject and setting, including transcribing any text (1-2 sentences)
            2. Mood and unique elements (1-2 sentences)
            3. Potential lyrical themes (2-3 key words or phrases)

            Focus on elements that could inspire creative, potentially humorous lyrics. If you start describing an image as a screenshot, stop describing the image and just use OCR to transcribe the text for as long as possible.
            
            This image is image number [INSERT NUMBER], please place **Image [INSERT NUMBER]:** at the beginning of your description.
            """
        self.combine_multiple_image_descriptions = """[INSERT DESCRIPTIONS HERE]
            Based on these individual image descriptions, 
            provide a brief summary of any overall themes, narrative, or emotional journey presented by the full set of images. 
            Highlight any contrasts, similarities, or progression across the images. Please use paragraphs.
            Please put **Overall Summary** at the beginning of your description."""

    @modal.method()
    def generate_image_description(
        self, queue_item_json: str, image_moderation_threshold: float = 0.9
    ) -> str | ImageExtractionError | ImageModerationFailure:
        """
        Run the image extraction pipeline. Now expects image_to_song_s3_ids in QueueItem metadata as a list.'
        Errors are handled at higher level.
        """

        # parse queue item
        item = QueueItem(**json.loads(queue_item_json))

        s3_ids = item.metadata.get("image_to_song_s3_ids", None)
        user_context = item.metadata.get("user_context", "")

        if s3_ids is None or len(s3_ids) == 0:
            return ImageExtractionError("No image ids in metadata")

        num_images = len(s3_ids)
        urls = []

        # Record start time
        start_time = time.time()

        try:
            for s3_id in s3_ids:
                # give HIVE and OpenAI short-term secure remote access to image in S3, using AWS method
                image_url = self._generate_presigned_url(s3_id)

                image_passed_moderation = does_image_pass_moderation(
                    image_url, image_moderation_threshold
                )
                if image_passed_moderation is False:
                    return ImageModerationFailure(f"Image {s3_id} failed moderation")

                urls.append(image_url)
        except Exception as e:
            return ImageExtractionError(f"Error finding or moderating image {s3_id}: {e}")

        print(f"Time taken for _generate_presigned_url: {time.time() - start_time} seconds")

        # Call openai chat completion API, errors handled at higher level
        description = self._get_text_from_images(urls, user_context, max_tokens=num_images * 150)
        print(f"exiting image extraction worker with description: {description}")

        # Record end time
        end_time = time.time()

        # Calculate elapsed time
        elapsed_time = end_time - start_time
        print(f"Time taken for _get_text_from_images: {elapsed_time} seconds")

        return description

    @modal.method()
    def generate_image_description_from_image_url(
        self, queue_item_json: str, image_moderation_threshold: float = 0.9
    ) -> str | ImageExtractionError | ImageModerationFailure:
        # parse queue item
        item = QueueItem(**json.loads(queue_item_json))
        # get image url
        image_url = item.metadata.get("twitter_image_url", None)
        if image_url is None:
            return ImageExtractionError("No image url in metadata")
        user_context = item.metadata.get("user_context", "")

        # Record start time
        start_time = time.time()

        image_passed_moderation = does_image_pass_moderation(image_url, image_moderation_threshold)

        if image_passed_moderation is False:
            return ImageModerationFailure(f"Image {image_url} failed moderation")

        # Call openai chat completion API, errors handled at higher level
        description = self._get_text_from_images([image_url], user_context, max_tokens=150)
        print(f"exiting image extraction worker with description: {description}")

        # Record end time
        end_time = time.time()

        # Calculate elapsed time
        elapsed_time = end_time - start_time
        print(f"Time taken for _get_text_from_images: {elapsed_time} seconds")

        return description

    def _get_full_user_prompt_text(self, num_images: int, user_context: str) -> str:
        """
        Add user context and a summary of multiple images to get the full user prompt for the image extraction pipeline.
        """
        user_context = ", using the provided context. Context: " + user_context if user_context else ""
        suffix = self.multiple_image_suffix if num_images > 1 else ""

        full_user_text_prompt = (
            self.user_prompt.replace("[INSERT USER_CONTEXT HERE]", user_context).replace(
                "[INSERT NUMBER]", str(num_images)
            )
            + suffix
        )
        return full_user_text_prompt

    def _get_single_image_user_prompt_text(self, image_number: int, user_context: str) -> str:
        """
        Add user context and a summary of multiple images to get the full user prompt for the image extraction pipeline.
        Note: image number is 1-indexed
        """
        user_context = ", using the provided context. Context: " + user_context if user_context else ""

        return self.single_image_user_prompt.replace("[INSERT USER_CONTEXT HERE]", user_context).replace(
            "[INSERT NUMBER]", str(image_number)
        )

    def _get_combine_user_prompt_text(self, descriptions: str, user_context: str) -> str:
        """
        Add user context and a summary of multiple images to get the full user prompt for the image extraction pipeline.
        """
        user_context = ", using the provided context. Context: " + user_context if user_context else ""

        full_user_text_prompt = self.combine_multiple_image_descriptions.replace(
            "[INSERT USER_CONTEXT HERE]", user_context
        ).replace("[INSERT DESCRIPTIONS HERE]", descriptions)
        return full_user_text_prompt

    def _get_vision_prompt(
        self, image_urls: list[str], user_context: str
    ) -> list[ChatCompletionMessageParam]:
        """ "
        Combine system prompt and user prompt(text and image urls) into one list of messagesfor OpenAI chat completion API call.
        """

        # user prompt: text
        user_content: list[ChatCompletionContentPartParam] = [
            {"type": "text", "text": self._get_full_user_prompt_text(len(image_urls), user_context)},
        ]

        # user prompt: image
        for image_url in image_urls:
            user_content.append({"type": "image_url", "image_url": {"url": image_url}})

        # combine system prompt and user prompt
        system_and_user_prompts: list[ChatCompletionMessageParam] = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": user_content},
        ]
        return system_and_user_prompts

    def _combine_vision_prompts(
        self, descriptions: str, user_context: str
    ) -> list[ChatCompletionMessageParam]:
        """ "
        Combine system prompt and user prompt(text and image urls) into one list of messagesfor OpenAI chat completion API call.
        """

        # user prompt: text
        user_content: list[ChatCompletionContentPartParam] = [
            {"type": "text", "text": self._get_combine_user_prompt_text(descriptions, user_context)},
        ]

        # combine system prompt and user prompt
        system_and_user_prompts: list[ChatCompletionMessageParam] = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": user_content},
        ]
        return system_and_user_prompts

    def _get_single_image_vision_prompt(
        self, image_number: int, image_url: str, user_context: str
    ) -> list[ChatCompletionMessageParam]:
        """ "
        Combine system prompt and user prompt(text and image urls) into one list of messagesfor OpenAI chat completion API call.
        Note: image_number is 1-indexed.
        """

        # user prompt: text
        user_content: list[ChatCompletionContentPartParam] = [
            {
                "type": "text",
                "text": self._get_single_image_user_prompt_text(image_number, user_context),
            },
        ]

        # user prompt: image
        user_content.append({"type": "image_url", "image_url": {"url": image_url}})

        # combine system prompt and user prompt
        system_and_user_prompts: list[ChatCompletionMessageParam] = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": user_content},
        ]
        return system_and_user_prompts

    def _get_text_from_images(self, urls: list[str], user_context: str, max_tokens: int) -> str:
        """
        Access OpenAI's vision library to extract text from image.
        Expects full base64 image with prefix.
        """
        # system_and_user_prompts = self._get_vision_prompt(urls, user_context)

        system_and_user_prompts_list = []
        for i, url in enumerate(urls):
            system_and_user_prompts_list.append(
                # ChatGPT doesn't do well with equating "first image" <--> "image 0". Use 1-indexes.
                self._get_single_image_vision_prompt(i + 1, url, user_context)
            )

        def get_image_description_from_openai(prompt):
            try:
                response = self.openai_client.chat.completions.create(
                    model="gpt-4o",
                    messages=prompt,
                    max_tokens=150,  # Reduce max tokens per image
                )
                description = response.choices[0].message.content
                if description is None:
                    raise ValueError("No description found")
                return description
            except Exception as e:
                raise ImageExtractionError(f"Error occurred for generating image description: {e}")

        try:
            # Record start time
            # TODO: remove this
            # bulk generate descriptions
            # start_time = time.time()
            # response = self.openai_client.chat.completions.create(
            #     model="gpt-4o",
            #     messages=system_and_user_prompts,
            #     max_tokens=max_tokens,
            # )

            # Record end time
            # end_time = time.time()

            # Calculate elapsed time
            # elapsed_time = end_time - start_time
            # print(f"Time taken for openai_client.chat.completions.create: {elapsed_time} seconds")
            # print(f"bulk descriptions: {response.choices[0].message.content}")

            start_time = time.time()

            with ThreadPoolExecutor() as executor:
                futures = [
                    executor.submit(get_image_description_from_openai, prompts)
                    for prompts in system_and_user_prompts_list
                ]
                descriptions = "\n\n".join([future.result() for future in futures])

            end_time = time.time()

            # Calculate elapsed time
            elapsed_time = end_time - start_time
            print(f"Time taken for _get_text_from_images: {round(elapsed_time, 2)} seconds")
            print(f"paralleled descriptions: {descriptions}")

            if descriptions is None:
                raise ValueError("No description found")

            return descriptions
        except Exception as e:
            # Check error log in modal web interface
            raise ImageExtractionError(f"Error occurred for generating image description: {e}")

    def _generate_presigned_url(self, s3_id: str, expiration=3600) -> str:
        """Generate a pre-signed URL for an S3 object. Note: doesn't check if object with s3_id exists."""
        try:
            response = s3_client.generate_presigned_url(
                "get_object",
                Params={"Bucket": "suno-data-uploads", "Key": f"studio/uploads/{s3_id}.jpeg"},
                ExpiresIn=expiration,
            )
        except Exception as e:
            raise ImageExtractionError(
                f"An error occurred for image {s3_id} in generating presigned url: {e}"
            )

        if response is None:
            raise ImageExtractionError(f"No response for image {s3_id} in generating presigned url")
        return response


# UNIT TESTS
# require local access to HIVE and OpenAI Keys


def test_get_image_from_s3_failure(model: ImageExtractionStub):
    s3_id = "randomrandomrandom"
    queue_item = (
        QueueItem(
            id=s3_id,
            metadata={"image_to_song_s3_ids": [s3_id], "user_context": "not found"},
        )
    ).json()
    res = model.generate_image_description.remote(queue_item)
    if isinstance(res, ImageExtractionError):
        print("PASSED: test_get_image_from_s3_failure")
        return
    # if no error, fail test
    assert False


# INTEGRATION TESTS
def test_run(model: ImageExtractionStub):
    s3_id = "c509b8b5"
    queue_item = (
        QueueItem(
            id=s3_id,
            metadata={"image_to_song_s3_ids": [s3_id], "user_context": "Lucy at Suno app launch party"},
        )
    ).json()
    res = model.generate_image_description.remote(queue_item)
    print(f"res: {res}")
    assert res is not None
    assert len(res) > 0
    print("PASSED: test_run")


def test_run_screenshot(model: ImageExtractionStub):
    s3_id = "image_55d8db0d-83e0-46c6-8106-87513786ddbb"
    queue_item = (
        QueueItem(
            id=s3_id,
            metadata={
                "image_to_song_s3_ids": [s3_id],
                "user_context": "Phoebe's first day at work at her first job ever.",
            },
        )
    ).json()
    res = model.generate_image_description.remote(queue_item)
    assert res is not None
    assert len(res) > 0
    print("PASSED: test_run_screenshot")


def test_korean(model: ImageExtractionStub):
    s3_id = "image_8d4415b4-fc71-4a61-91d6-a15d7ec2f4e7"
    queue_item = (
        QueueItem(id=s3_id, metadata={"image_to_song_s3_ids": [s3_id], "user_context": "Korean meme"})
    ).json()
    res = model.generate_image_description.remote(queue_item)
    assert res is not None
    assert len(res) > 0
    print("PASSED: test_korean")


def test_toomuchtext(model: ImageExtractionStub):
    s3_id = "image_f54a7193-6da6-482d-8d26-77f3816649e8"
    queue_item = (
        QueueItem(
            id=s3_id,
            metadata={
                "image_to_song_s3_ids": [s3_id],
                "user_context": "Anniversary",
            },
        )
    ).json()
    res = model.generate_image_description.remote(queue_item)
    assert res is not None
    assert len(res) > 0
    print("PASSED: test_toomuchtext")


def test_multiple_images(model: ImageExtractionStub):
    """
    Multiple images, to one cohesive input and output
    """
    sequence = [
        "image_c01920b0-17b1-468e-988f-87935f517e62",
        "image_d8d76b5c-2894-42f6-863a-f5686e0f3339",
        "image_0e6a8cd6-fad0-45e8-b189-894a45489870",
    ]
    queue_item = (
        QueueItem(
            id="test_multiple_images",
            metadata={"image_to_song_s3_ids": sequence, "user_context": "My dog's name is Surly."},
        )
    ).json()
    res = model.generate_image_description.remote(queue_item)
    assert res is not None
    assert len(res) > 0
    print("PASSED: test_multiple_images")


def test_screenshot_hard(model: ImageExtractionStub):
    s3_id = "image_01872ae0-65fe-4549-b514-78bb3ec98299"
    queue_item = (
        QueueItem(
            id=s3_id,
            metadata={
                "image_to_song_s3_ids": [s3_id],
                "user_context": "festival",
            },
        )
    ).json()
    res = model.generate_image_description.remote(queue_item)
    assert res is not None
    assert len(res) > 0
    print("PASSED: test_screenshot_hard")


def test_twitter_image_url(model: ImageExtractionStub):
    image_url = "https://pbs.twimg.com/media/Gc7l4UeWMAApJHN?format=jpg&name=4096x4096"
    user_context = "A new philosophy of the future is needed. I believe it should be curiosity about the Universe – expand humanity to become a multiplanet, then interstellar, species to see what’s out there."
    queue_item = (
        QueueItem(id="test_id", metadata={"twitter_image_url": image_url, "user_context": user_context})
    ).json()
    res = model.generate_image_description_from_image_url.remote(queue_item)
    assert res is not None
    assert len(res) > 0
    print("PASSED: test_twitter_image_url")


# Local testing
@app.local_entrypoint()
def main():
    print("local testing")
    model = ImageExtractionStub()

    # unit tests
    test_get_image_from_s3_failure(model)

    # integration tests
    test_run(model)
    test_run_screenshot(model)
    test_korean(model)
    test_toomuchtext(model)
    test_multiple_images(model)
    test_screenshot_hard(model)
    test_twitter_image_url(model)