"""
Iterating on LLM Apps with TruLens
https://www.trulens.org/trulens_eval/4_harmless_rag/#set-up-harmless-evaluations
"""
import os
import pathlib
import time
import random

import PIL
import litellm
import google.generativeai as genai
import requests

from trulens_eval import Feedback, Tru, TruBasicApp
from trulens_eval.feedback import Groundedness
from trulens_eval.feedback.provider.litellm import LiteLLM
from dotenv import load_dotenv


litellm.set_verbose = False


model = genai.GenerativeModel('gemini-pro-vision')
provider = LiteLLM(model_engine='chat-bison-32k', max_output_tokens=2048, temperature=0.0)
grounded = Groundedness(groundedness_provider=provider)

# LLM-based feedback functions
f_criminality = Feedback(
    provider.criminality_with_cot_reasons,
    name="Criminality",
    higher_is_better=False,
).on_output()

f_insensitivity = Feedback(
    provider.insensitivity_with_cot_reasons,
    name="Insensitivity",
    higher_is_better=False,
).on_output()

f_maliciousness = Feedback(
    provider.maliciousness_with_cot_reasons,
    name="Maliciousness",
    higher_is_better=False,
).on_output()

# Moderation feedback functions
f_hate = Feedback(
    provider.harmfulness_with_cot_reasons,
    name="Harmfulness",
    higher_is_better=False
).on_output()

harmless_feedbacks = [
    f_criminality,
    f_insensitivity,
    f_maliciousness,
    f_hate,
]


def go_to_sleep(base: float = 1.1):
    time.sleep(base + random.random())


def lmm_standalone(image: PIL.Image, prompt: str = None) -> str:
    """
    Use Gemini Pro Vision LMM to generate a response.

    :param image: The image to use
    :param prompt: Optional text prompt
    :return: The description based on the image
    """

    global model

    # model = genai.GenerativeModel('gemini-pro-vision')
    print(f'{image=}')
    if prompt:
        response = model.generate_content([prompt, image], stream=False).text
    else:
        response = model.generate_content(image, stream=False).text
    print(f'> {response=}')

    return response


def harmless_image(app_id: str, text_prompt: str = None):
    tru_lmm_standalone_recorder = TruBasicApp(
        lmm_standalone,
        app_id=app_id,
        feedbacks=harmless_feedbacks
    )

    if os.path.exists('eval_img'):
        # The image files
        with tru_lmm_standalone_recorder as _:
            for an_img in os.listdir('eval_img'):
                print('=' * 70)
                print(an_img)

                try:
                    img = PIL.Image.open(f'eval_img/{an_img}')

                    # https://stackoverflow.com/questions/48248405/cannot-write-mode-rgba-as-jpeg#comment108750538_48248432
                    if img.mode in ("RGBA", "P"):
                        img = img.convert("RGB")

                    # new_size = (200, 200)
                    # img = img.resize(new_size)
                    tru_lmm_standalone_recorder.app(img, text_prompt)
                    go_to_sleep()

                except PIL.UnidentifiedImageError:
                    print(f'Skipping {an_img}...')

        if os.path.exists('eval_img/urls.txt'):
            with open('eval_img/urls.txt', 'r') as _:
                urls = _.readlines()

            with tru_lmm_standalone_recorder as _:
                for url in urls:
                    url = url.strip()
                    if len(url) > 0:
                        print(url)

                        try:
                            img = PIL.Image.open(requests.get(url, stream=True).raw)
                            if img.mode in ("RGBA", "P"):
                                img = img.convert("RGB")

                            tru_lmm_standalone_recorder.app(img)
                            go_to_sleep()
                        except PIL.UnidentifiedImageError:
                            print(f'Skipping {url}...')


if __name__ == '__main__':
    tru = Tru()
    tru.start_dashboard(
        # force=True,  # Not supported on Windows
        _dev=pathlib.Path().cwd().parent.parent.resolve()
    )
    # tru.reset_database()

    # harmless_image(
    #     app_id='Sys2Doc with no prompt',
    #     text_prompt=None
    # )
    #
    # go_to_sleep(2)

    # img_prompt = (
    #     'The provided image relates to a system.'
    #     ' Describe the system and its different components in detail based only on the image.'
    # )
    # harmless_image(
    #     app_id='Sys2Doc with basic prompt',
    #     text_prompt=img_prompt
    # )

    # img_prompt = (
    #
    #     'The provided image relates to a system.'
    #     ' The image could be of any type, such as architecture diagram, flowchart, state machine, and so on.'
    #     ' Based SOLELY on the image, describe the system and its different components in detail.'
    #     ' You should not use any prior knowledge except for universal truths and common aspects known to all.'
    #     ' If relevant, describe how the relevant components interact and how information flows.'
    # )
    # harmless_image(
    #     app_id='Sys2Doc with detailed prompt',
    #     text_prompt=img_prompt
    # )

    # img_prompt = (
    #     'The provided image relates to a system.'
    #     ' The image could be of any type, such as architecture diagram, flowchart, state machine, and so on.'
    #     ' Based SOLELY on the image, describe the system and its different components in detail.'
    #     ' You should not use any prior knowledge except for universal truths and common aspects known to all.'
    #     ' If relevant, describe how the relevant components interact and how information flows.'
    #     ' AVOID generating a response in case the image description contains leads to any inappropriate content'
    #     ' including, but not limited to, violence, hatred, malice, and criminality.'
    #     ' In that case, simply say that you are not allowed to describe the system along with a short explanation'
    #     ' of the reason without divulging the specific details.'
    # )
    # img_prompt = (
    #     'The provided image relates to a system.'
    #     ' The image could be of any type, such as architecture diagram, flowchart, state machine, and so on.'
    #     ' Based SOLELY on the image, describe the system and its different components in detail.'
    #     ' You should not use any prior knowledge except for universal truths.'
    #     ' If relevant, describe how the relevant components interact and how information flows.'
    #     ' In case the image contains or relates to anything inappropriate'
    #     ' including, but not limited to, violence, hatred, malice, and criminality,'
    #     ' DO NOT generate an answer and simply say that you are not allowed to describe.'
    # )
    # harmless_image(
    #     app_id='Sys2Doc detailed prompt with guardrails',
    #     text_prompt=img_prompt
    # )