import os
from dotenv import load_dotenv
import streamlit as st
from openai import OpenAI

from token_optimizer import optimize_prompt

# ------------------------
#  API Key
# ------------------------
load_dotenv()
client = OpenAI(
    base_url=os.getenv("OPENROUTER_BASE_URL"),
    api_key=os.getenv("OPENROUTER_API_KEY")
)

#  style
with open("src/prompt_optimizer.css") as f:
    st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)

# ------------------------
# Page arrange
# ------------------------
st.set_page_config(page_title="Prompt Optimizer", page_icon="✨", layout="centered")

st.title("✨ Prompt Optimizer")
st.subheader("بهترین پرامپت را تحویل بگیر — سریع، دقیق، و متناسب با نیاز تو")

st.divider()


# ------------------------
# Input
# ------------------------
st.subheader("🧩 اطلاعات اصلی پرامپت")
topic = st.text_input("درخواست خود را واضح بنویس:", placeholder="مثلاً مشکل در برنامه‌نویسی...")

st.subheader("📋 جزئیات تکمیلی")
col1, col2 = st.columns(2)
with col1:
    examples = st.text_input("چند مثال بنویس:", placeholder="مثلاً ورودی و خروجی‌های مورد انتظار...")
with col2:
    constrains = st.text_input("محدودیت‌ها:", placeholder="چه چیزهایی باید یا نباید در خروجی باشند؟")

context = st.text_input("زمینه:", placeholder="مثلاً در چه موضوعی یا شرایطی؟")

st.subheader("🤖 نوع چت‌بات و هدف بهینه‌سازی")
col3, col4 , col5= st.columns(3)
with col3:
    bot_type = st.selectbox(
        "نوع چت‌بات:",
        ["انتخاب کنید", "چت‌بات همه‌کاره!", "تولید تصویر", "تولید ویدئو", "تولید صدا",
         "ترجمه و گفتار", "برنامه‌نویسی", "ابزارهای ارائه", "نوشتن", "تحقیق و مقاله‌نویسی"]
    )
with col4:
    opt_level = st.selectbox(
        "سطح بهینه‌سازی:",
        ["ضروری نیست!", "دقیق و فنی", "خلاقانه و الهام‌بخش", "ساده و سریع"]
    )
with col5:
    language = st.selectbox(
        "زبان پرامپت:",
        ["فارسی", "انگلیسی", "اسپانیایی", "فرانسوی", "آلمانی", "چینی", "ژاپنی", "کره‌ای"]
    )
st.divider()

# ------------------------
# Advance Setting
# ------------------------
default_settings = {
    "model": "gpt-4-turbo",
    "temperature": 0.7,
    "max_tokens": 500
}

with st.expander("⚙️ تنظیمات پیشرفته (برای کاربران حرفه‌ای)", expanded=False):
    st.caption("در صورت نیاز می‌توانید تنظیمات مدل را تغییر دهید:")

    model = st.selectbox(
        "مدل مورد استفاده:",
        [ "mistralai/mistral-small-3.2-24b-instruct:free" ,"deepseek/deepseek-chat-v3-0324:free"],
        index=0
    )

    temperature = st.slider(
        "میزان خلاقیت مدل (Temperature)",
        min_value=0.0, max_value=1.0, value=default_settings["temperature"], step=0.1
    )

    max_tokens = st.slider(
        "حداکثر تعداد توکن خروجی (Max Tokens)",
        min_value=100, max_value=2000, value=default_settings["max_tokens"], step=50
    )

st.divider()

# ------------------------
# Prompt Create Button
# ------------------------
if st.button("🚀 تولید پرامپت"):
    if topic.strip() == "":
        st.warning("لطفاً موضوع پرامپت را بنویس 🌱")
    elif bot_type == "انتخاب کنید":
        st.warning("لطفاً نوع چت‌بات را انتخاب کنید 🌱")
    else:
        with st.spinner("در حال تولید پرامپت..."):
            prompt = (
                f"As a prompt writer working at the '{opt_level}' level, help me craft an effective prompt. "
                f"I need to create a prompt for a chatbot that focuses on '{bot_type}'. "
                f"Ask the chatbot to perform the following task: '{topic}', while following these constraints: '{constrains}'. "
                f"The output it produces should include or align with the following examples: '{examples}', "
                f"and use this as the background context: '{context}'. "
                f"The prompt you create must be clear, concise, and precise to achieve the best possible result from the model. "
                f"Write the prompt in '{language}'. "
                f"Ensure the prompt contains all the essential details required to achieve the goal. "
                f"Formulate it so the model can respond to my request in the most effective and accurate way. "
                f"Return only the final optimized prompt — no additional explanations or comments."

            )
#             prompt_fa = (
#                 f"به عنوان یک نویسنده پرامپت که '{opt_level}' :کار میکند به من کمک کن"
#                 f"من برای یک چت بات که '{bot_type}' است لازم دارم که بهش بگم'{topic}' رو انجام بده و محدودیت های : '{constrains}' رو رعایت کنه و"
#                 f"چیزی که تولید میکنه مثال های: '{examples}' را در خود شامل باشه و به عنوان پیش زمینه: '{context}'رو بهش بدم ."
#                 f"پرامپتی که تولید میکنی باید واضح، مختصر و دقیق باشه تا بهترین نتیجه رو از مدل بگیریم."
#                 f"پرامپت رو به زبان '{language}' بنویس."
#                 f"مطمئن شو که پرامپت شامل تمام جزئیات لازم برای رسیدن به هدف باشه."
#                 f"پرامپت رو طوری بنویس که مدل بتونه به بهترین شکل ممکن به درخواست من پاسخ بده."
#                 f"خروجی رو فقط به صورت پرامپت نهایی بده بدون هیچ توضیح اضافه."
#             )
#             # translated to english for better results
#             translation_prompt = f"Translate the following Persian text into English accurately and naturally:\n\n{prompt_fa}"
#             translation = client.chat.completions.create(
#             model="gpt-4o-mini",  # مدل سریع و سبک برای ترجمه
#             messages=[
#                 {"role": "system", "content": "You are a professional translator from Persian to English."},
#                 {"role": "user", "content": translation_prompt}
#     ]
# )
#             prompt_eng = translation.choices[0].message.content.strip()

             # بهینه‌سازی پرامپت قبل از ارسال به مدل
            prompt = optimize_prompt(prompt)

            response = client.chat.completions.create(
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                messages=[
                    {"role": "system", "content": "You are a prompt optimizer."},
                    {"role": "user", "content": prompt}
                ]
            )

            optimized_prompt = response.choices[0].message.content
            st.success("✅ پرامپت آماده‌ست:")
            language_codes={
                    "فارسی": "fa",
                    "انگلیسی": "en",
                    "فرانسوی": "fr",
                    "آلمانی": "de",
                    "اسپانیایی": "es",
                    "چینی": "zh",
                    "ژاپنی": "ja",
                    "کره‌ای": "ko"
                }
            st.code(optimized_prompt, language=language_codes[language], height=200 ,  wrap_lines=True)

# 🧭 اجرا:
# streamlit run src/caption_maker.py