import streamlit as st
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from loguru import logger

from ui.Component import side_bar_links

st.set_page_config(
    page_title='工具箱',
    page_icon='🔨',
    layout='centered',
)

st.title("学术中英互译")


def clean_screen():
    st.session_state.ac_translate_history = []
    st.session_state.ac_code_text = ''


with st.sidebar:
    side_bar_links()

    st.toggle('翻译模式', key='is_en2zh')
    if st.session_state.get('is_en2zh'):
        st.caption('en → zh')
    else:
        st.caption('zh → en')

    st.button("清屏", on_click=clean_screen)

if 'ac_translate_history' not in st.session_state:
    st.session_state.ac_translate_history = []

if 'ac_code_text' not in st.session_state:
    st.session_state.ac_code_text = ''

chat_container = st.container(height=600, border=False)

with chat_container:
    for message in st.session_state.ac_translate_history:
        icon = 'logo.png' if message['role'] != 'user' else None
        with st.chat_message(message['role']):
            st.markdown(message['content'])

if st.session_state.ac_code_text:
    with st.container(height=120, border=False):
        st.code(st.session_state.ac_code_text, 'markdown')


def ac_translate(original_text: str, is_en2zh: bool):
    if is_en2zh:
        _prompt = ChatPromptTemplate.from_messages(
            [
                SystemMessage(content="你的任务是将用户给出的英文文本翻译为**语句通顺的**中文，同时务必注意不能漏翻或错翻。"
                                      "对于专有名词或公式等，若没有合适的译文可以保留原本英文文本不做翻译。"
                                      "对于原句子中被markdown代码块或latex公式包裹的内容，不做改动原样返回。"),
                ("human", "这是你需要翻译的句子:\n{original_text}\n\n注意你的回答只需要给出翻译后的英文文本，不要包含其他东西")
            ]
        )
    else:
        _prompt = ChatPromptTemplate.from_messages(
            [
                SystemMessage(content="你的任务是以**学术论文的风格**将用户给出的中文文本翻译成英文，务必注意不能漏翻或错翻。"),
                ("human", "这是你需要翻译的句子:\n{original_text}\n\n注意你的回答只需要给出翻译后的英文文本，不要包含其他东西")
            ]
        )

    llm = ChatOpenAI(
        model_name="gpt-4o-mini",
        temperature=0,
        openai_api_key=st.secrets['gpt_key'],
        streaming=True
    )

    chain = _prompt | llm

    llm_result = chain.stream({"original_text": original_text})

    return llm_result


if prompt := st.chat_input():
    logger.info(f'[translate]: {prompt}')

    chat_container.chat_message("human").write(prompt)
    st.session_state.ac_translate_history.append({'role': 'user', 'content': prompt})

    response = ac_translate(prompt, st.session_state.get('is_en2zh'))
    translate_result = chat_container.chat_message("ai").write_stream(response)
    st.session_state.ac_translate_history.append({'role': 'assistant', 'content': translate_result})
    st.session_state.ac_code_text = translate_result

    st.rerun()