import streamlit as st
import json
import pandas as pd
from groq import Groq
from prompt_engine import PromptGenerator
from utils import load_techniques, load_styles, parse_prompt_result

# モデル特性に基づくカラーテーマ設定
model_themes = {
    "deepseek-r1-distill-llama-70b": {"primary": "#3366cc", "secondary": "#b3c6ff"},
    "llama-3.3-70b-versatile": {"primary": "#4b367c", "secondary": "#d1c4e9"},
    "meta-llama/llama-4-maverick-17b-128e-instruct": {"primary": "#00897b", "secondary": "#b2dfdb"},
    "mistral-saba-24b": {"primary": "#0277bd", "secondary": "#b3e5fc"},
    "qwen-qwq-32b": {"primary": "#6a1b9a", "secondary": "#e1bee7"}
}

# 選択されたモデル（初回読み込み時のデフォルト値）
current_model = "llama-3.3-70b-versatile"
theme = model_themes.get(current_model, {"primary": "#4b367c", "secondary": "#d1c4e9"})

# ページ設定
st.set_page_config(
    page_title="AI Art Prompt Generator",
    page_icon="🎨",
    layout="wide",
    initial_sidebar_state="expanded"
)

# カスタムCSSの適用
st.markdown(f"""
<style>
    .stApp {{
        background-color: {theme["secondary"]}20;  /* 20%の透明度 */
    }}
    .stButton>button {{
        background-color: {theme["primary"]};
        color: white;
    }}
    .stTabs [data-baseweb="tab-list"] {{
        border-bottom-color: {theme["primary"]}40;
    }}
    .stTabs [aria-selected="true"] {{
        color: {theme["primary"]} !important;
        border-bottom-color: {theme["primary"]} !important;
    }}
</style>
""", unsafe_allow_html=True)

# Groq APIクライアント設定
@st.cache_resource
def get_groq_client():
    return Groq(api_key=st.secrets["GROQ_API_KEY"])

# データロード
@st.cache_data
def load_data():
    techniques = load_techniques()
    styles = load_styles()
    return techniques, styles

# 初期化
client = get_groq_client()
techniques, styles = load_data()
prompt_generator = PromptGenerator()

# タイトルとイントロ
st.title("🎨 究極のMidJourney/nijiJourneyプロンプトジェネレーター")
st.markdown("""
このツールは、MidJourneyとnijiJourneyで使用する高度なプロンプトを簡単に生成します。
基本的な説明から始めて、AIが最先端のテクニックを応用した洗練されたプロンプトに変換します。
""")

# サイドバー
with st.sidebar:
    st.header("⚙️ 基本設定")
    
    mode = st.radio(
        "モード選択",
        ["MidJourney", "nijiJourney"],
        help="生成するプロンプトの対象となるAIツールを選択します"
    )
    
    complexity = st.select_slider(
        "複雑さレベル",
        options=["シンプル", "バランス", "詳細", "実験的"],
        value="バランス",
        help="生成されるプロンプトの複雑さを調整します"
    )
    
    # モデル選択セクションを追加
    st.markdown("---")
    st.header("🧠 AIモデル")
    
    model = st.selectbox(
        "Groqモデルを選択",
        [
            "deepseek-r1-distill-llama-70b",
            "llama-3.3-70b-versatile",
            "meta-llama/llama-4-maverick-17b-128e-instruct",
            "mistral-saba-24b", 
            "qwen-qwq-32b"
        ],
        index=1,  # デフォルトは llama-3.3-70b-versatile
        help="使用するGroqのAIモデルを選択します。各モデルによって生成結果の特性が異なります"
    )
    
    # モデル選択時に変数を更新
    if model != current_model:
        current_model = model
        theme = model_themes.get(current_model, {"primary": "#4b367c", "secondary": "#d1c4e9"})
        
        # テーマの更新（実際には再読み込みが必要）
        st.markdown(f"""
        <style>
            .stApp {{
                background-color: {theme["secondary"]}20;
            }}
            .stButton>button {{
                background-color: {theme["primary"]};
                color: white;
            }}
        </style>
        """, unsafe_allow_html=True)
    
    # モデル情報の表示
    model_info = {
        "deepseek-r1-distill-llama-70b": "DeepSeekのLLaMA 70Bベースモデル。知識と推論能力に優れています。",
        "llama-3.3-70b-versatile": "Meta AIの最新LLaMA 3.3。バランスの取れた高性能モデルで、多くのタスクに適しています。",
        "meta-llama/llama-4-maverick-17b-128e-instruct": "LLaMA 4 Maverickの指示チューニング版。小型ながら高性能です。",
        "mistral-saba-24b": "Mistral AIのSaba 24Bモデル。創造性と一貫性のバランスに優れています。",
        "qwen-qwq-32b": "Qwen QWQモデル。創造的な表現力が高く、芸術的プロンプトに適しています。"
    }
    
    st.caption(model_info.get(model, ""))
    
    st.markdown("---")
    
    st.header("🧙‍♂️ テクニック設定")
    auto_select = st.checkbox("テクニックを自動選択", value=True, 
                             help="オフにすると手動でテクニックを選択できます")
    
    if not auto_select:
        selected_techniques = st.multiselect(
            "使用するテクニック",
            options=list(techniques.keys()),
            default=[],
            help="プロンプトに適用する特定のテクニックを選択します"
        )
    else:
        selected_techniques = []

# メインインプット
st.header("🎯 あなたのイメージを教えてください")
user_input = st.text_area(
    "描きたいイメージを自由に入力してください",
    height=100,
    help="日本語で自然に入力できます。AIが最適な英語のプロンプトに変換します。"
)

# 詳細設定タブ
tab1, tab2, tab3, tab4 = st.tabs(["基本設定", "高度なテクニック", "専門家向け設定", "実験的機能"])

with tab1:
    col1, col2 = st.columns(2)
    
    with col1:
        aspect_ratio = st.selectbox(
            "アスペクト比",
            ["1:1 (正方形)", "16:9 (横長)", "9:16 (縦長)", "3:2 (標準)", "4:3", "2:3", "カスタム"],
            help="画像の縦横比を設定します"
        )
        
        if aspect_ratio == "カスタム":
            custom_ratio = st.text_input("カスタム比率 (例: 5:4)")
            aspect_ratio = custom_ratio if custom_ratio else "1:1"
        
        quality = st.slider(
            "クオリティ (--q)", 
            min_value=0.25, 
            max_value=5.0, 
            value=2.0, 
            step=0.25,
            help="高いほど詳細な画像が生成されますが、処理時間が長くなります"
        )
    
    with col2:
        style_category = st.selectbox(
            "スタイルカテゴリ",
            list(styles.keys()),
            help="生成画像の基本的なスタイル方向性を選択します"
        )
        
        style_options = styles[style_category]
        style = st.selectbox(
            "スタイル",
            style_options,
            help="具体的なスタイルを選択します"
        )
        
        stylize = st.slider(
            "スタイル強度 (--s)", 
            min_value=0, 
            max_value=1000, 
            value=100, 
            step=50,
            help="AIの創造性レベルを調整します。高いほど創造的な結果になります"
        )

with tab2:
    if auto_select:
        st.info("「テクニックを自動選択」がオンになっています。AIが最適なテクニックを選択します。")
    else:
        st.markdown("### 選択したテクニック")
        
        if selected_techniques:
            for tech in selected_techniques:
                st.markdown(f"**{tech}**: {techniques[tech]['description']}")
        else:
            st.warning("テクニックが選択されていません。サイドバーから選択してください。")
    
    st.markdown("### テクニック解説")
    tech_df = pd.DataFrame({
        "テクニック": list(techniques.keys()),
        "説明": [t["description"] for t in techniques.values()],
        "効果": [t["effect"] for t in techniques.values()],
    })
    st.dataframe(tech_df, use_container_width=True)

with tab3:
    st.markdown("### 専門家向け設定")
    
    col1, col2 = st.columns(2)
    
    with col1:
        camera_angle = st.selectbox(
            "カメラアングル",
            ["自動選択", "eye level", "low angle", "high angle", "dutch angle", "aerial view", "worm's eye view"],
            help="撮影視点の角度を指定します"
        )
        
        lens_type = st.selectbox(
            "レンズタイプ",
            ["自動選択", "wide angle", "telephoto", "macro", "fisheye", "tilt-shift", "normal"],
            help="カメラレンズの種類を指定します"
        )
    
    with col2:
        lighting = st.selectbox(
            "照明設定",
            ["自動選択", "natural light", "golden hour", "blue hour", "dramatic lighting", "backlight", "studio lighting"],
            help="画像の照明条件を指定します"
        )
        
        composition = st.selectbox(
            "構図",
            ["自動選択", "rule of thirds", "golden ratio", "symmetrical", "asymmetrical", "leading lines", "frame within frame"],
            help="画像の構図テクニックを指定します"
        )
    
    # MidJourney特有のパラメータセクションを追加
    st.markdown("### MidJourney/nijiJourney特有のパラメータ")
    
    # タブでパラメータをカテゴリ分け
    param_tab1, param_tab2, param_tab3 = st.tabs(["スタイル参照", "パーソナライズ", "その他"])
    
    with param_tab1:
        # スタイルリファレンス (--sref) 設定
        use_style_reference = st.checkbox("スタイルリファレンス (--sref) を使用", 
                                         help="別画像のスタイルを参照し、色調や質感を新規生成に適用します")
        
        if use_style_reference:
            sref_type = st.radio("参照方法", ["ランダム", "画像URL", "コード指定"], horizontal=True)
            
            if sref_type == "ランダム":
                st.info("ランダムに選ばれたスタイルが適用されます (--sref random)")
                sref_value = "random"
            elif sref_type == "画像URL":
                sref_value = st.text_input("画像URL", help="スタイル参照する画像のURLを入力してください")
            else:
                sref_value = st.text_input("スタイルコード", help="以前生成されたsrefコードを入力 (例: 1234567)")
            
            # スタイルウェイト
            style_weight = st.slider("スタイルウェイト (--sw)", 0, 1000, 100, 
                                    help="参照画像の影響度を調整します。高いほど強く反映されます")
            
            # スタイルバージョン
            style_version = st.radio("スタイルバージョン (--sv)", [1, 2, 3, 4], index=3, horizontal=True,
                                    help="Style Referenceアルゴリズムのバージョンを選択します")
    
    with param_tab2:
        # パーソナライゼーション (--p) 設定
        use_personalization = st.checkbox("パーソナライゼーション (--p) を使用", 
                                         help="自身の好みに学習させたプロファイルを適用します")
        
        if use_personalization:
            personalization_type = st.radio("適用方法", ["デフォルト", "プロファイルID", "コード指定"], horizontal=True)
            
            if personalization_type == "デフォルト":
                st.info("設定済みのデフォルトプロファイルが適用されます (--p)")
                personalization_value = ""
            elif personalization_type == "プロファイルID":
                personalization_value = st.text_input("プロファイルID", help="使用するプロファイルIDを入力")
            else:
                personalization_value = st.text_input("パーソナライズコード", help="以前生成されたコードを入力")
            
            # パーソナライズ強度は --stylize を共用
            st.caption("パーソナライズの適用度はスタイル強度 (--s) パラメータで調整されます")
    
    with param_tab3:
        # その他のパラメータ
        col1, col2 = st.columns(2)
        
        with col1:
            use_character_reference = st.checkbox("キャラクター参照 (--cref)", 
                                               help="複数シーンで同一キャラクターを維持します")
            
            if use_character_reference:
                cref_value = st.text_input("キャラクター参照URL", help="キャラクターが写っている画像のURLを入力")
            
            use_repeat = st.checkbox("繰り返し生成 (--repeat / --r)", 
                                   help="同一プロンプトで複数グリッドを生成します")
            
            if use_repeat:
                repeat_count = st.number_input("繰り返し回数", min_value=1, max_value=10, value=2)
        
        with col2:
            use_image_weight = st.checkbox("画像プロンプトの影響度 (--iw)", 
                                         help="画像プロンプト（イメージURL）の影響度を調整します")
            
            if use_image_weight:
                image_weight = st.slider("画像ウェイト", 0.0, 2.0, 1.0, 0.1)
            
            generation_mode = st.radio("生成モード", 
                                      ["デフォルト", "fast", "relax", "turbo", "draft"], 
                                      index=0, horizontal=True,
                                      help="GPU速度やコストを最適化するモードを選択")
            
            visibility = st.radio("共有設定", 
                                ["デフォルト", "public", "stealth"], 
                                index=0, horizontal=True,
                                help="MidJourneyサイトへの公開/非公開を制御")
    
    # カスタムパラメータセクション
    st.markdown("### カスタムパラメータ")
    advanced_params = st.text_area(
        "カスタムパラメータ (JSONフォーマット)", 
        "{}",
        help="上級者向け: 追加のカスタムパラメータをJSON形式で指定できます"
    )

with tab4:
    st.markdown("### 実験的機能")
    st.warning("これらの機能は予測不能な結果をもたらす可能性があります")
    
    col1, col2 = st.columns(2)
    
    with col1:
        use_pattern_interrupt = st.checkbox(
            "パターン中断", 
            help="対照的な概念を組み合わせて意外性のある結果を生成します"
        )
        
        use_logical_paradox = st.checkbox(
            "論理パラドックス", 
            help="矛盾する要素を組み合わせて創造的な表現を引き出します"
        )
    
    with col2:
        use_magic_words = st.checkbox(
            "魔法の単語", 
            help="特定の「魔法の単語」を使って特殊な効果を引き出します"
        )
        
        use_emotion_matrix = st.checkbox(
            "感情マトリックス", 
            help="複雑な感情表現を階層的に組み込みます"
        )
    
    if use_pattern_interrupt:
        pattern_concepts = st.text_input(
            "対照的な概念を入力 (コンマ区切り)", 
            help="例: 「エーテル,サイバーパンク」「古代,未来」"
        )
    
    if use_emotion_matrix:
        primary_emotion = st.selectbox(
            "主要感情", 
            ["joy", "fear", "sadness", "anger", "surprise", "contemplative", "peaceful", "mysterious"]
        )
        secondary_emotion = st.selectbox(
            "二次感情", 
            ["hopeful", "melancholic", "anxious", "determined", "curious", "nostalgic", "dreamy", "tense"]
        )

# 生成ボタン
if st.button("🚀 プロンプトを生成", type="primary", use_container_width=True):
    if not user_input:
        st.error("描きたいイメージを入力してください")
    else:
        with st.spinner("AIがプロンプトを生成中..."):
            # パラメータの準備
            params = {
                "input": user_input,
                "mode": mode.lower(),
                "complexity": complexity,
                "model": model,
                "aspect_ratio": aspect_ratio,
                "quality": quality,
                "style": style,
                "stylize": stylize,
                "advanced": {
                    "camera_angle": None if camera_angle == "自動選択" else camera_angle,
                    "lens_type": None if lens_type == "自動選択" else lens_type,
                    "lighting": None if lighting == "自動選択" else lighting,
                    "composition": None if composition == "自動選択" else composition,
                    "advanced_params": json.loads(advanced_params) if advanced_params else {}
                },
                "experimental": {
                    "pattern_interrupt": use_pattern_interrupt,
                    "logical_paradox": use_logical_paradox,
                    "magic_words": use_magic_words,
                    "emotion_matrix": use_emotion_matrix
                }
            }
            
            # MidJourney特有のパラメータを追加
            
            # スタイルリファレンス
            if 'use_style_reference' in locals() and use_style_reference:
                params["advanced"]["style_reference"] = {
                    "use": True,
                    "value": sref_value,
                    "style_weight": style_weight,
                    "style_version": style_version
                }
            
            # パーソナライゼーション
            if 'use_personalization' in locals() and use_personalization:
                params["advanced"]["personalization"] = {
                    "use": True,
                    "type": personalization_type,
                    "value": personalization_value
                }
            
            # キャラクター参照
            if 'use_character_reference' in locals() and use_character_reference:
                params["advanced"]["character_reference"] = {
                    "use": True,
                    "value": cref_value
                }
            
            # 繰り返し生成
            if 'use_repeat' in locals() and use_repeat:
                params["advanced"]["repeat"] = {
                    "use": True,
                    "count": repeat_count
                }
            
            # 画像プロンプトの影響度
            if 'use_image_weight' in locals() and use_image_weight:
                params["advanced"]["image_weight"] = {
                    "use": True,
                    "value": image_weight
                }
            
            # 生成モード
            if 'generation_mode' in locals() and generation_mode != "デフォルト":
                params["advanced"]["generation_mode"] = generation_mode
            
            # 共有設定
            if 'visibility' in locals() and visibility != "デフォルト":
                params["advanced"]["visibility"] = visibility
            
            # 実験的機能の詳細
            if use_pattern_interrupt and 'pattern_concepts' in locals() and pattern_concepts:
                params["experimental"]["pattern_concepts"] = [c.strip() for c in pattern_concepts.split(",")]
            
            if use_emotion_matrix and 'primary_emotion' in locals() and 'secondary_emotion' in locals():
                params["experimental"]["emotions"] = {
                    "primary": primary_emotion,
                    "secondary": secondary_emotion
                }
            
            # テクニック
            if not auto_select and selected_techniques:
                params["techniques"] = selected_techniques
            
            # プロンプト生成呼び出し
            result = prompt_generator.generate_prompt(client, params)
            
            # 結果の解析
            prompt, explanation = parse_prompt_result(result)
            
            # 結果表示
            st.success("プロンプトを生成しました！")
            
            st.markdown("### 生成されたプロンプト")
            
            # まずプロンプトを表示
            st.code(prompt, language=None)
            
            # JavaScriptを使ったコピー機能の実装
            js_code = f"""
            <script>
            function copyToClipboard() {{
                const text = `{prompt.replace("`", "\\`").replace("'", "\\'")}`;
                navigator.clipboard.writeText(text).then(function() {{
                    document.getElementById('copy-status').innerHTML = "コピーしました！";
                    setTimeout(function() {{
                        document.getElementById('copy-status').innerHTML = "";
                    }}, 2000);
                }})
                .catch(function(error) {{
                    document.getElementById('copy-status').innerHTML = "コピーできませんでした: " + error;
                }});
            }}
            </script>
            <button 
                onclick="copyToClipboard()" 
                style="background-color:{theme['primary']}; color:white; border:none; padding:8px 16px; border-radius:4px; cursor:pointer;"
            >
                📋 クリップボードにコピー
            </button>
            <span id="copy-status" style="margin-left:10px; color:green;"></span>
            """
            st.markdown(js_code, unsafe_allow_html=True)
            
            st.caption("MidJourney/nijiJourneyにこのプロンプトを貼り付けて使用してください")
            
            # 使用したモデルの表示を追加
            st.caption(f"使用モデル: **{model}**")
            
            if explanation:
                st.markdown("### プロンプト解説")
                st.markdown(explanation)

# フッター
st.markdown("---")
st.markdown("© 2025 AI Art Prompt Generator")