import streamlit as st
from transformers import pipeline
import os

# ページ設定
st.set_page_config(page_title="お嬢様語変換アプリ", page_icon="💃", layout="wide")

# カスタムCSSでお嬢様らしいデザイン
st.markdown("""
    <style>
    .stApp {
        background-color: #FFF0F5;
    }
    h1 {
        color: #C71585;
        font-family: 'Georgia', serif;
    }
    .stTextInput > div > div > input {
        border: 2px solid #FFB6C1;
        border-radius: 10px;
    }
    .stButton > button {
        background-color: #FF69B4;
        color: white;
        border-radius: 10px;
        font-family: 'Georgia', serif;
    }
    </style>
""", unsafe_allow_html=True)

# タイトルと説明
st.title("💃 お嬢様語変換アプリ💃 ")
st.markdown("""
ごきげんよう！このアプリは、皆様の入力したテキストをエレガントなお嬢様語に変換いたしますわ！  
テキストを入力して「変換するわ！」ボタンを押してくださいませ。どんな言葉もお嬢様語に変換いたしますわ。おほほほほ！
""")

# サイドバーでエレガント度を設定
st.sidebar.header("設定")
elegance_level = st.sidebar.slider("エレガント度", min_value=1, max_value=3, value=2, step=1, 
                                  help="1: 控えめ, 2: 標準, 3: 超お嬢様")

# プロンプトをエレガント度に応じて設定（日本語特化）
prompt_templates = {
    1: "以下のテキストを丁寧で上品な日本語に変換してください。自然でエレガントな表現を使ってください：\n{}",
    2: "以下のテキストをお嬢様風のエレガントな日本語に変換してください。「～ですわ」「おほほ」などを適度に使い、華やかで優雅な口調にしてください。例: 「こんにちは」→「ごきげんよう、お元気ですわ？」：\n{}",
    3: "以下のテキストを非常に誇張されたお嬢様風の日本語に変換してください。「～ですわよ！」「おほほほほ！」を多用し、華麗でドラマチックな表現にしてください。例: 「こんにちは」→「ごきげんようでございますわ！なんともお元気でいらっしゃいますの？おほほほほ！」：\n{}"
}

# LLMのロード（初回起動時にモデルをキャッシュ）
@st.cache_resource
def load_model():
    try:
        # rinna/japanese-gpt2-small をロード
        pipe = pipeline("text-generation", model="rinna/japanese-gpt2-small", device=-1)  # CPU使用
        # トークナイザーのパディング設定
        if pipe.tokenizer.pad_token_id is None:
            pipe.tokenizer.pad_token_id = pipe.tokenizer.eos_token_id
        return pipe
    except Exception as e:
        st.error(f"モデルロードに失敗しましたわ！: {str(e)}")
        return None

generator = load_model()

# ユーザー入力
input_text = st.text_input("テキストを入力してくださいませ（例: こんにちは、元気？）", "なんでや、阪神関係ないやろ - 入力: なんでや、阪神関係ないやろ！")

# 変換ボタン
if st.button("変換するわ！"):
    if input_text:
        if generator is None:
            st.error("モデルがロードできておりませんわ！おほほ…")
        else:
            with st.spinner("エレガントに変換中ですわ…おほほ！"):
                # プロンプトを生成
                prompt = prompt_templates[elegance_level].format(input_text)
                
                # LLMでテキスト生成
                try:
                    output = generator(prompt, max_length=150, num_return_sequences=1, 
                                     do_sample=True, temperature=0.9, top_p=0.95, 
                                     pad_token_id=generator.tokenizer.pad_token_id,
                                     bos_token_id=generator.tokenizer.bos_token_id,
                                     eos_token_id=generator.tokenizer.eos_token_id)
                    converted_text = output[0]["generated_text"].split(prompt)[-1].strip()
                    
                    # 後処理：お嬢様らしさを強調
                    if elegance_level >= 2 and not converted_text.endswith("おほほ"):
                        converted_text += " おほほ！" if elegance_level == 2 else " おほほほほ！"
                    
                    # 結果表示
                    st.subheader("変換結果")
                    st.markdown(f"**入力**: {input_text}")
                    st.markdown(f"**お嬢様語**: {converted_text}")
                except Exception as e:
                    st.error(f"変換中にエラーが発生しましたわ！: {str(e)}")
    else:
        st.warning("テキストを入力してくださいませ！")

# フッター
st.markdown("---")
st.markdown("Created with 💖 by Streamlit & Hugging Face Spaces")