import os
from typing import List, Optional
import gradio as gr
from dotenv import load_dotenv
import instructor
from pydantic import BaseModel
from openai import OpenAI
from groq import Groq

load_dotenv()

groq_client = instructor.from_groq(Groq(), mode=instructor.Mode.JSON)

QUERY = """
    extract all singular ingredients from the provided ingredient string. 
    While doing so, shorten the ingredient name to only the essential part.
    Do not translate the names.
    """
STIMULI = "This is very important to my career."


class Ingredient(BaseModel):
    """
    Represents an ingredient with a name and an optional amount.

    Attributes:
        name (str): The name of the ingredient.
        amount (Optional[float]): The amount of the ingredient. This can be None if the amount is not specified.
    """

    name: str
    amount: Optional[float]
    is_allergen: Optional[bool]


class Ingredients(BaseModel):
    """
    Ingredients model that contains a list of Ingredient objects.

    Attributes:
        contains (List[Ingredient]): A list of Ingredient objects.
    """

    contains: List[Ingredient]


# client = instructor.from_openai(OpenAI(), mode=instructor.Mode.JSON)


def predict(ingredients: str, openai_key: str = "") -> Ingredients:
    """
    Predicts the ingredients using a GPT-3.5-turbo model.

    Args:
        ingredients (str): A string containing the ingredients to be parsed.

    Returns:
        Ingredients: The parsed ingredients in the form of an Ingredients object.
    """

    if openai_key == "":
        return groq_client.chat.completions.create(
            model="llama-3.1-70b-versatile",
            messages=[
                {
                    "role": "user",
                    "content": f"""
                {QUERY}
                {ingredients}
                {STIMULI}
                """,
                },
            ],
            response_model=Ingredients,
            temperature=0.0,
        ).model_dump_json(exclude_unset=True, exclude_none=True)

    client = OpenAI(api_key=openai_key)

    return (
        client.beta.chat.completions.parse(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "user",
                    "content": f"""
                {QUERY}
                {ingredients}
                {STIMULI}
                """,
                },
            ],
            response_format=Ingredients,
            temperature=0.0,
        )
        .choices[0]
        .message.parsed.model_dump_json(exclude_unset=True, exclude_none=True)
    )


demo = gr.Interface(
    fn=predict,
    inputs=["text", gr.Text(label="OpenAI API Key", type="password")],
    outputs="json",
    examples=[
        [
            "97,2 % DINKELVOLLKORNMEHL, GERSTENMALZMEHL, Salz",
            "",
        ],
        [
            "BIO Grüne Linsen. Grüne Linsen aus kontrolliert biologischem Anbau. Kühl bei unter +30°C lagern. Nicht-EU-Landwirtschaft.",
            "",
        ],
        [
            "Milchschokoladenkuvertüre, Kakao: 40,5% mindestens. Zucker, Kakaobutter, VOLLMILCHPULVER, Kakaomasse (Ghana), Emulgator: SOJALECITHIN, natürliches Vanillepulver. Trocken und kühl bei +17°C bis +18°C lagern. Eigenschaften: Eiweiß aus tierischer Milch.",
            "",
        ],
    ],
)
demo.launch()