from __future__ import annotations
import os
# import openai
from langchain.prompts import PromptTemplate
from langchain.chat_models import ChatOpenAI
from typing import Any
from langchain.base_language import BaseLanguageModel
from langchain.chains.llm import LLMChain
import gradio as gr
# import streamlit as st

# from google.cloud import aiplatform
import vertexai
from vertexai.language_models import TextGenerationModel, TextEmbeddingModel

# import vertexai

# PROJECT_ID = "[your-project-id]"  # @param {type:"string"}
# vertexai.init(project=PROJECT_ID, location="us-central1")
# from vertexai.language_models import TextGenerationModel

# OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
# OPENAI_API_KEY='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'

PROJECT_ID = "agileai-poc"

vertexai.init(project=PROJECT_ID, location="us-central1")
parameters = {
    "max_output_tokens": 256,
    "temperature": 0.2,
    "top_p": 0.8,
}

generation_model = TextGenerationModel.from_pretrained("text-bison@001")
print("model called")
response = generation_model.predict(
    "Generate a product description that is creative and SEO compliant. Emojis should be added to make product description look appealing. Begin!", **parameters)
print(response)
embedding_model = TextEmbeddingModel.from_pretrained("textembedding-gecko@001")
print("model2 called")
prompt_file = "prompt_template.txt"
print(generation_model.predict("describe", **parameters))


class ProductDescGen(LLMChain):
    """LLM Chain specifically for generating multi paragraph rich text product description using emojis."""

    @classmethod
    def from_llm(
        cls, llm: BaseLanguageModel, prompt: str, **kwargs: Any
    ) -> ProductDescGen:
        """Load ProductDescGen Chain from LLM."""
        return cls(llm=llm, prompt=prompt, **kwargs)


def product_desc_generator(product_name, keywords):
    with open(prompt_file, "r") as file:
        prompt_template = file.read()
# # llm = ChatOpenAI(
#     #     model_name="gpt-3.5-turbo",
#     #     temperature=0.7,
#     #     openai_api_key=OPENAI_API_KEY,
#     # )
    llm = vertexai(max_output_tokens=256, temperature=0.2, top_p=0.8)
    print("runned")
    PROMPT = PromptTemplate(
        input_variables=["product_name", "keywords"], template=prompt_template
    )
#     # llm = ChatOpenAI(
#     #     model_name="gpt-3.5-turbo",
#     #     temperature=0.7,
#     #     openai_api_key=OPENAI_API_KEY,
#     # )

    # llm2 = vertexai(max_output_tokens=256, temperature=0.2, top_p=0.8)

    ProductDescGen_chain = ProductDescGen.from_llm(llm=llm, prompt=PROMPT)
    ProductDescGen_query = ProductDescGen_chain.apply_and_parse(
        [{"product_name": product_name, "keywords": keywords}]
    )
#     response = generation_model.predict(
#         "Generate a product description that is creative and SEO compliant. Emojis should be added to make product description look appealing. Begin!", **llm)

    return ProductDescGen_query[0]["text"]  # , {response.text}


with gr.Blocks() as demo:
    gr.HTML("""<h1>Welcome to Product Description Generator</h1>""")
    gr.Markdown(
        "Generate Product Description for your products instantly!<br>"
        "Provide product name and keywords related to that product. Click on 'Generate Description' button and multi-paragraph rich text product description will be genrated instantly.<br>"
        "Note: Generated product description is SEO compliant and can be used to populate product information."
    )

    with gr.Tab("Generate Product Description!"):
        product_name = gr.Textbox(
            label="Product Name",
            placeholder="Nike Shoes",
        )
        keywords = gr.Textbox(
            label="Keywords (separated by commas)",
            placeholder="black shoes, leather shoes for men, water resistant",
        )
        product_description = gr.Textbox(label="Product Description")
        click_button = gr.Button(value="Generate Description!")
        click_button.click(
            product_desc_generator, [
                product_name, keywords], product_description
        )

demo.launch(share=True)