import os

from langchain.chains import LLMChain
from langchain.chat_models import ChatOpenAI
from langchain_core.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate, \
    PromptTemplate

from app.core.config import OPENAI_MODEL, OPENAI_API_KEY

# os.environ["http_proxy"] = "http://127.0.0.1:7890"
# os.environ["https_proxy"] = "http://127.0.0.1:7890"


llm = ChatOpenAI(model_name=OPENAI_MODEL,
                 openai_api_key=OPENAI_API_KEY,
                 temperature=0)


def translate_to_en(text):
    template = (
        """You are a translation expert, proficient in various languages.
        And can translate various languages into English.
        Please translate to grammatically correct English regardless of the input language.
        If the input is in English or numbers, check for grammatical errors. If there are no errors, output the input directly.
        If there are grammatical errors, correct them and then output the sentence."""
    )
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)

    # 待翻译文本由 Human 角色输入
    human_template = "User input : {text}"
    human_message_prompt = HumanMessagePromptTemplate.from_template(input_variables=["text"], template=human_template)

    # 使用 System 和 Human 角色的提示模板构造 ChatPromptTemplate
    chat_prompt_template = ChatPromptTemplate.from_messages(
        [system_message_prompt, human_message_prompt]
    )
    translate_chain = LLMChain(llm=llm, prompt=chat_prompt_template)

    template = (
        """
        Input sentence:
        {translate}
        1. Based on the input,adjust the input sentence to make it more suitable for prompts for generating images,
        ensuring all key nouns or adjectives related to the image are retained.
        2. Simplify complex sentence structures and clarify ambiguous expressions.
        3. Only Output the adjusted English sentence.

        Output :
        """
    )
    # "Based on the input sentence, extract key adjectives and nouns.Only Output extracted key words."
    # 1. Check if the input sentence contains any grammatical errors. If there are errors, please correct them before proceeding.

    prompt_template = PromptTemplate(input_variables=["translate"], template=template)
    prompt_chain = LLMChain(llm=llm, prompt=prompt_template)

    from langchain.chains import SimpleSequentialChain
    overall_chain = SimpleSequentialChain(chains=[translate_chain, prompt_chain], verbose=True)

    response = overall_chain.run(text)
    return response


def main():
    """Main function"""
    translate_to_en("生成一件运动风格的夹克，带有拉链和口袋，适合休闲穿着")


if __name__ == "__main__":
    main()