import streamlit as st
from streamlit_option_menu import option_menu

from dataclasses import dataclass

import google.generativeai as genai

from langdetect import detect
import pycountry
from deep_translator import GoogleTranslator


import base64

from PIL import Image

import requests

from dotenv import load_dotenv
import os

# Charger les variables d'environnement
load_dotenv()

HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACE_API_KEY")
GOOGLE_GENAI_API_KEY = os.getenv("GOOGLE_GENAI_API_KEY")

headers = {"Authorization": f"Bearer {HUGGINGFACE_API_KEY}"}

genai.configure(api_key=GOOGLE_GENAI_API_KEY)

# Configuration de la page
st.set_page_config(page_title="Application IA", layout="wide")


st.markdown(
    """
    <style>
        .content{
            border-radius: 8px;
            box-shadow: 0 1px 3px rgba(65,65,65,.492);
            background: #fbfafa;
            padding : 0 10px 0 20px
        }
        .stChatMessage .markdown-text-container {
            word-wrap: break-word;
        }
    </style>
    """,
    unsafe_allow_html=True,
)

# Options Menu
with st.sidebar:
    menu = option_menu(
        "Onyx AI",
        ["Acceuil", "Documentation"],
        icons=["play-btn", "info-circle"],
        menu_icon="intersect",
        default_index=0,
        styles={
            "icon": {"color": "#374151", "font-size": "17px"},
            "nav-link": {
                "font-size": "17px",
                "text-align": "left",
                "margin": "0px",
                "--hover-color": "#eee",
            },
            "nav-link-selected": {"background-color": "#519ba7"},
        },
    )


# Intro Page
if menu == "Acceuil":
    # Header

    # # Contenu de la navbar avec des icônes appropriées
    selected = option_menu(
        menu_title=None,  # Aucun titre pour le menu
        options=[
            "Chatbot",
            "TextToSpeech",
            "Classification",
            "Image to text",
        ],
        icons=[
            "robot",
            "volume-up",
            "camera",
            "file-text",
        ],  # Icônes pour chaque option
        menu_icon="cast",  # Icône du menu
        default_index=0,  # Option par défaut
        orientation="horizontal",
        styles={
            "container": {"background-color": "#fafafa"},
            "icon": {"color": "#374151", "font-size": "17px"},
            "nav-link": {
                "font-size": "17px",
                "text-align": "left",
                "--hover-color": "#eee",
            },
            "nav-link-selected": {"background-color": "#519ba7"},
        },
    )  # Orientation du menu

    # Affichage du contenu basé sur le choix de la navbar avec st.markdown pour une meilleure intégration
    if selected == "Chatbot":

        # Affichage du titre et du message de bienvenue
        st.markdown(
            """
                <div class='content'>
                    <h1>Chatbot intelligent</h1>
                </div>
                <br/>
                 <p>Bienvenue dans la section Chatbot intelligent.</p>
            """,
            unsafe_allow_html=True,
        )

        model = genai.GenerativeModel("gemini-pro")

        @dataclass
        class Message:
            actor: str
            payload: str

        USER = "user"
        ASSISTANT = "ai"
        MESSAGES = "messages"

        if MESSAGES not in st.session_state:
            st.session_state[MESSAGES] = [
                Message(
                    actor=ASSISTANT, payload="Salut! Comment puis-je vous aider ? 😎"
                )
            ]

        msg: Message
        for msg in st.session_state[MESSAGES]:
            st.chat_message(msg.actor).write(msg.payload)

        prompt: str = st.chat_input("Enter a prompt here")

        if prompt:
            st.session_state[MESSAGES].append(Message(actor=USER, payload=prompt))
            st.chat_message(USER).write(prompt)
            generate = model.generate_content(prompt)
            response: str = generate.text
            st.session_state[MESSAGES].append(
                Message(actor=ASSISTANT, payload=response)
            )
            st.chat_message(ASSISTANT).write(response)

    elif selected == "TextToSpeech":

        # Affichage du titre et du message de bienvenue
        st.markdown(
            """
                <div class='content'>
                    <h1>Text to speech</h1>
                </div>
                <br/>
                <p>Bienvenue dans la section Text to speech en espagnol.</p>

            """,
            unsafe_allow_html=True,
        )

        API_URL = "https://api-inference.huggingface.co/models/facebook/mms-tts-spa"

        @dataclass
        class Message:
            actor: str
            payload: str

        USER = "user"
        ASSISTANT = "ai"
        MESSAGES = "messages_audio"

        if MESSAGES not in st.session_state:
            st.session_state[MESSAGES] = [
                Message(
                    actor=ASSISTANT,
                    payload="Salut! Quelle phrase (peut importe la langue) souhaitez-vous prononcer en espagnol ? 😎",
                )
            ]

        def query(payload):
            response = requests.post(API_URL, headers=headers, json=payload)
            return response.content

        # msg: Message
        for msg in st.session_state[MESSAGES]:
            st.chat_message(msg.actor).write(msg.payload)

        prompt: str = st.chat_input("Entrez le texte à prononcer :")

        if prompt:
            st.session_state[MESSAGES].append(Message(actor=USER, payload=prompt))
            st.chat_message(USER).write(prompt)

            lang = pycountry.languages.get(alpha_2=detect(prompt))

            lang_name = lang.name if lang else "Langue inconnue"

            translated = GoogleTranslator(source="auto", target="es").translate(prompt)

            st.chat_message(ASSISTANT).write(
                f"Langue détecté {lang_name} \n \n {translated} \n \n Génération de l'audio ..."
            )

            parameters = {"inputs": translated}

            audio = query(parameters)

            st.chat_message(ASSISTANT).audio(data=audio)

    elif selected == "Classification":
        st.markdown(
            """
                    <div class='content'>
                        <h1>Zero shot image classification</h1>
                    </div>
                    <br/>
                    <p>Découvrez la classification d'images avec zerp shot.</p>

                """,
            unsafe_allow_html=True,
        )

        API_URL = "https://api-inference.huggingface.co/models/openai/clip-vit-large-patch14-336"

        def query(image_bytes, parameters):
            # Encode the image bytes in base64
            encoded_image = base64.b64encode(image_bytes).decode("utf-8")
            payload = {"parameters": parameters, "inputs": encoded_image}
            response = requests.post(API_URL, headers=headers, json=payload)
            return response.json()

        @dataclass
        class Message:
            actor: str
            payload: str

        USER = "user"
        ASSISTANT = "ai"

        if "default_label" not in st.session_state:
            st.session_state["default_label"] = (
                "voiture,ananas,pastèque,tomate,bouteille"
            )

        # Interface utilisateur pour l'entrée de l'image et du libellé
        uploaded_image = st.file_uploader(
            "Upload an image", type=["jpg", "png", "jpeg"]
        )

        st.markdown(
            "Veuillez entrer les labels séparés par des virgules (exemple : voiture,ananas,pastèque,tomate,bouteille):"
        )

        label = st.chat_input("Entrez le label séparé des virgules")

        # Bouton pour envoyer l'image et le libellé
        if label:
            if uploaded_image is not None and label:
                # Lire et stocker l'image temporairement
                image = Image.open(uploaded_image)

                image_bytes = uploaded_image.getvalue()

                st.chat_message(USER).image(image)

                st.chat_message(ASSISTANT).text("Classification en cours...")

                convert_label_to_english = GoogleTranslator(
                    source="auto", target="en"
                ).translate(label)

                labels = convert_label_to_english.split(
                    ","
                )  # Splitting the input label string into a list

                print(labels)

                parameters = {"candidate_labels": labels}

                response = query(image_bytes, parameters)

                translated = GoogleTranslator(source="auto", target="fr").translate(
                    f"The highest score is {response[0]['score']:.3f} for the label {response[0]['label']}"
                )

                st.chat_message(ASSISTANT).text(translated)

            else:
                st.write("Please upload an image and enter a label.")

    elif selected == "Image to text":
        st.markdown(
            """
            <div class='content'>
                <h1>Image to text</h1>
            </div>
            <br/>
            <p>Convertissez des images en texte grâce à notre outil.</p>

                        """,
            unsafe_allow_html=True,
        )

        model_pro_vision = genai.GenerativeModel("gemini-pro-vision")

        @dataclass
        class Message:
            actor: str
            payload: str

        USER = "user"
        ASSISTANT = "ai"

        # Interface utilisateur pour l'entrée de l'image et du libellé
        uploaded_image_ImgTText = st.file_uploader(
            "Upload an image", type=["jpg", "png", "jpeg"]
        )

        # Bouton pour envoyer l'image et le libellé
        if st.button("Send"):
            if uploaded_image_ImgTText is not None:
                # Lire et stocker l'image temporairement
                image = Image.open(uploaded_image_ImgTText)

                st.chat_message(USER).image(image)

                st.chat_message(ASSISTANT).text("Génération en cours...")

                response = model_pro_vision.generate_content(
                    ["Write a short description of the image", image],
                    stream=True,
                )

                response.resolve()

                print(response.text)

                translated = GoogleTranslator(source="auto", target="fr").translate(
                    response.text
                )

                st.chat_message(ASSISTANT).text(translated)

            else:
                st.write("Please upload an image and enter a label.")


elif menu == "Documentation":

    # Lorsque "About" est sélectionné, afficher le contenu À propos dans le corps principal
    st.markdown(
        """
        # À propos de cette application

        Cette application Streamlit démontre l'utilisation avancée de l'intelligence artificielle pour fournir une variété de services, y compris un chatbot intelligent, la conversion de texte en parole, la classification d'images sans exemples préalables, et la conversion d'images en texte. Cette application est un exemple de comment les technologies de pointe peuvent être intégrées dans des applications web modernes.

        ## Créateur

        - [James Olongo](https://jamesolongo.vercel.app/)

        ## Technologies utilisées

        - **Streamlit**: Un framework pour la création rapide d'applications web pour la science des données et le machine learning.
        - **PyTorch & Transformers**: Utilisés pour les opérations de traitement du langage naturel et de génération de texte.
        - **torchaudio & gTTS**: Pour la génération de parole et le traitement audio.
        - **Langdetect & Pycountry**: Pour la détection de la langue et le traitement basé sur la langue.
        - **Deep Translator**: Un puissant outil de traduction, utilisé en remplacement de Googletrans.
        - **Google Generative AI**: Pour le chatbot intelligent, utilisant le modèle "gemini-pro" pour générer des réponses textuelles.
        - **VITS**: Un modèle de synthèse vocale de pointe de Facebook pour la conversion de texte en parole en espagnol.
        - **OpenAI CLIP**: Pour la classification d'images sans exemples, utilisant le modèle "clip-vit-large-patch14-336" pour la classification d'image en zero-shot.

        ## Détails des Applications

        ### Chatbot Intelligent
        Utilise le modèle "gemini-pro" de Google Generative AI pour répondre aux requêtes des utilisateurs. Ce modèle est capable de comprendre et de générer des réponses à une grande variété de questions.
        ```python
        import google.generativeai as genai

        genai.configure(api_key="API_KEY")

        model = genai.GenerativeModel("gemini-pro")
        
        generate = model.generate_content("Comment déclarer une variable python")

        response = generate.text   

        ```


        ### Texte en Parole
        Emploie le modèle "facebook/mms-tts-spa" pour convertir du texte en parole en espagnol, démontrant la capacité de transformer le texte en un fichier audio naturel.
        ```python
        model = VitsModel.from_pretrained("facebook/mms-tts-spa")
        
        tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-spa")

        inputs = tokenizer(translated, return_tensors="pt")

        with torch.no_grad():
            output = model(**inputs).waveform
        ```


        ### Zero shot image classification
        Utilise le modèle "openai/clip-vit-large-patch14-336" d'OpenAI pour classifier des images sans exemples préalables, illustrant comment une IA peut comprendre et catégoriser le contenu visuel.
        ```python


        model_name_zero_shot_classification = "openai/clip-vit-large-patch14-336"
        
        classifier_zero_shot_classification = pipeline(
            "zero-shot-image-classification", model=model_name_zero_shot_classification
        )
        
        scores = classifier_zero_shot_classification(
            image_to_classify, candidate_labels="YOUR_LABEL"
        )

        print f"The highest score is {scores[0]['score']:.3f} for the label {scores[0]['label']}"

        ```

        ### Conversion d'Images en Texte



        """,
        unsafe_allow_html=True,
    )