import streamlit as st
import time
import math
import uuid
import datetime
import json
import os
import random
import requests
import urllib.parse

# 引用配置与工具
from config import CUSTOM_CSS, PROVIDER_MAP, SIZE_PRESETS, QUALITY_PRESETS, STYLE_PRESETS, ITEMS_PER_PAGE, WARMUP_COUNT, LOCAL_CACHE_DIR, THUMBNAIL_DIR
from utils import ensure_png_bytes, auto_translate

# 引用拆分出的 UI 组件 (V30.11/12 新架构)
from ui import (
    get_managers, render_header, render_pagination, render_automation_dashboard, 
    view_image_dialog, edit_webdav_dialog, login_page, change_pw_page
)

# 页面配置
st.set_page_config(
    page_title="AI 艺术工作室 (旗舰版)",
    page_icon="🎨",
    layout="wide",
    initial_sidebar_state="collapsed"
)
st.markdown(CUSTOM_CSS, unsafe_allow_html=True)

# Session State 初始化
keys = ["current_record", "logged_in", "username", "must_change_pw", "prompt_input", "global_error", "upscale_error", "prompt_seed", "history_loaded", "poster_bytes", "auto_task_config", "style_preset", "gallery_page", "batch_upload_logs", "saved_tasks", "tokens_loaded", "selected_task_index", "last_mtime"]
for k in keys:
    if k not in st.session_state:
        if "batch_upload_logs" == k: st.session_state[k] = []
        elif "saved_tasks" == k: st.session_state[k] = []
        elif "gallery_page" == k: st.session_state[k] = 1
        # [Fix] 关键修复：last_mtime 必须初始化为 0.0，否则与浮点数比较会报错
        elif "last_mtime" == k: st.session_state[k] = 0.0
        else: st.session_state[k] = None

# 初始化管理器
mgrs = get_managers()
pl, tm, cg, sm, wm, atm, up, pg, am = mgrs["pl"], mgrs["tm"], mgrs["cg"], mgrs["sm"], mgrs["wm"], mgrs["atm"], mgrs["up"], mgrs["pg"], mgrs["am"]

# 启动后台服务
if not st.session_state.history_loaded:
    sm.start_background_cacher(WARMUP_COUNT)
    st.session_state.history_loaded = True

# --- 登录校验 ---
if not st.session_state.get("logged_in", False):
    login_page()
    st.stop()

if st.session_state.get("must_change_pw", False):
    change_pw_page()
    st.stop()

# --- 主界面逻辑 ---
render_header()

# 顶部 Tab 导航
tab_main, tab_gallery, tab_batch, tab_settings = st.tabs(["🎨 创作工坊", "🖼️ 画廊归档", "🤖 自动化任务", "⚙️ 系统设置"])

# ==================== Tab 1: 创作工坊 ====================
with tab_main:
    c_side, c_canvas = st.columns([1, 2], gap="medium")
    
    with c_side:
        st.subheader("💡 灵感控制台")
        
        # [V30.40 Fix] 修复 get_best_token 调用错误，改为 get_active_tokens
        def on_random_click():
            try:
                # 获取可用 Token 列表
                active_tokens = tm.get_active_tokens()
                token = active_tokens[0] if active_tokens else None
                
                seed_text = random.choice(["Cyberpunk City", "Peaceful Garden", "Future Tech", "Magic Forest", "Portrait of a Lady"])
                
                # 尝试调用文本生成接口
                if token:
                    u = f"https://gen.pollinations.ai/text/{urllib.parse.quote(f'Imagine: {seed_text}')}?model=openai&seed={random.randint(0,999)}&private=true&key={token}"
                    try:
                        resp = requests.get(u, timeout=5)
                        if resp.status_code == 200: seed_text = resp.text
                    except: pass
                
                st.session_state.prompt_input = seed_text
            except Exception as e:
                st.session_state.global_error = f"灵感生成失败: {str(e)}"

        # 提示词输入区
        prompt = st.text_area("提示词 (Prompt)", value=st.session_state.get("prompt_input", ""), height=120, key="prompt_area")
        
        c_btn1, c_btn2 = st.columns([1, 1])
        if c_btn1.button("🎲 随机灵感", use_container_width=True):
            on_random_click()
            st.rerun()
            
        if c_btn2.button("🧹 清空", use_container_width=True):
            st.session_state.prompt_input = ""
            st.rerun()

        # 参数设置
        with st.expander("🛠️ 参数设置", expanded=True):
            c_p1, c_p2 = st.columns(2)
            
            # 尺寸选择
            size_name = c_p1.selectbox("画幅比例", list(SIZE_PRESETS.keys()), index=1)
            width, height = SIZE_PRESETS[size_name]
            if size_name == "Custom (自定义)":
                c_cust1, c_cust2 = st.columns(2)
                width = c_cust1.number_input("宽", 64, 2048, 1024, step=64)
                height = c_cust2.number_input("高", 64, 2048, 1024, step=64)

            # 模型与风格
            prov = c_p2.selectbox("模型供应商", list(PROVIDER_MAP.keys()), index=0)
            models = PROVIDER_MAP.get(prov, {})
            model_key = st.selectbox("选择模型", list(models.keys()), format_func=lambda x: models[x])
            
            # 风格与画质
            style = st.selectbox("艺术风格", list(STYLE_PRESETS.keys()), index=0)
            quality = st.selectbox("画质增强", list(QUALITY_PRESETS.keys()), index=0)
            
            negative = st.text_input("反向提示词 (Negative)", value="")

    with c_canvas:
        # 错误提示
        if st.session_state.global_error:
            st.error(f"❌ {st.session_state.global_error}")
            if st.button("清除错误"): 
                st.session_state.global_error = None
                st.rerun()

        # 生成按钮区
        c_gen_main, c_gen_status = st.columns([1, 3])
        is_generating = False
        
        with c_gen_main:
            if st.button("🚀 开始生成", type="primary", use_container_width=True):
                if not prompt:
                    st.warning("请先输入提示词")
                else:
                    is_generating = True
                    st.session_state.global_error = None

        # 预览与生成逻辑
        preview_container = st.container(border=True)
        with preview_container:
            if is_generating:
                with st.status("🎨 正在绘制中...", expanded=True) as status:
                    start_time = time.time()
                    try:
                        status.write("🔍 正在优化提示词...")
                        # 自动翻译逻辑
                        final_prompt = prompt
                        if any("\u4e00" <= char <= "\u9fff" for char in prompt):
                             final_prompt = auto_translate(prompt)
                             status.write(f"🌐 已翻译: {final_prompt[:30]}...")
                        
                        status.write("🖌️ 正在请求绘图引擎...")
                        
                        # [V30.38 Fix] 调用 generate 时传递 token_manager，不传具体 Key
                        # 手动模式策略：balance (余额优先轮询)
                        img_bytes, final_p, src_info = cg.generate(
                            final_prompt, width, height, prov, model_key, 
                            style, quality, negative, 
                            token_manager=tm, 
                            strategy="balance" 
                        )
                        
                        status.write("💾 正在保存作品...")
                        record = {
                            "id": uuid.uuid4().hex,
                            "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M"),
                            "prompt_zh": prompt,
                            "prompt_en": final_p,
                            "model": f"{src_info}",
                            "style": style
                        }
                        sm.save_record(record, img_bytes)
                        st.session_state.current_record = record
                        st.session_state.poster_bytes = None # 重置海报
                        
                        # 自动上传 WebDAV
                        if wm.configs:
                            status.write("☁️ 同步到 WebDAV...")
                            wm.upload_to_all_nodes(img_bytes, f"manual_{record['id']}.png")
                            
                        status.update(label=f"✅ 完成! (耗时 {time.time()-start_time:.1f}s)", state="complete", expanded=False)
                        st.rerun()
                        
                    except Exception as e:
                        st.session_state.global_error = str(e)
                        status.update(label="❌ 生成失败", state="error")

            # 显示当前作品
            if st.session_state.current_record:
                rec = st.session_state.current_record
                local_path = os.path.join(LOCAL_CACHE_DIR, rec["img_path_in_repo"])
                if os.path.exists(local_path):
                    st.image(local_path, use_container_width=True)
                    
                    # 操作栏
                    c_act1, c_act2, c_act3 = st.columns(3)
                    if c_act1.button("🔍 查看详情", use_container_width=True):
                        with open(local_path, "rb") as f: 
                            view_image_dialog(f.read(), rec)
                            
                    if c_act2.button("🖼️ 生成海报", use_container_width=True):
                        with open(local_path, "rb") as f:
                            st.session_state.poster_bytes = pg.generate_social_poster(f.read(), rec)
                    
                    if c_act3.button("✨ 超分放大", use_container_width=True):
                        with st.spinner("正在放大 2x..."):
                            try:
                                with open(local_path, "rb") as f:
                                    big_bytes = up.upscale_image(f.read())
                                    new_rec = rec.copy()
                                    new_rec["id"] = uuid.uuid4().hex
                                    new_rec["timestamp"] = datetime.datetime.now().strftime("%Y-%m-%d %H:%M")
                                    new_rec["model"] += " (Upscaled)"
                                    sm.save_record(new_rec, big_bytes)
                                    st.session_state.current_record = new_rec
                                    st.success("放大成功，已保存为新副本")
                                    time.sleep(1)
                                    st.rerun()
                            except Exception as e:
                                st.error(f"放大失败: {e}")

                    # 海报展示区
                    if st.session_state.poster_bytes:
                        st.divider()
                        st.subheader("📱 社交海报预览")
                        st.image(st.session_state.poster_bytes, width=360)
                        st.download_button("⬇️ 下载海报", st.session_state.poster_bytes, file_name=f"poster_{rec['id']}.jpg", mime="image/jpeg")

            else:
                st.markdown("""
                <div class="empty-state-container">
                    <div style="font-size: 48px;">🎨</div>
                    <div>在左侧输入提示词，开始您的创作之旅</div>
                </div>
                """, unsafe_allow_html=True)

# ==================== Tab 2: 画廊归档 ====================
with tab_gallery:
    col_gal_head, col_gal_search = st.columns([3, 1])
    with col_gal_head: st.subheader("📂 历史作品库")
    
    # 获取数据
    all_history = sm.load_index()
    if not all_history:
        st.info("暂无历史记录")
    else:
        # 分页逻辑
        total_items = len(all_history)
        total_pages = math.ceil(total_items / ITEMS_PER_PAGE)
        current_page = st.session_state.gallery_page
        start_idx = (current_page - 1) * ITEMS_PER_PAGE
        end_idx = start_idx + ITEMS_PER_PAGE
        page_items = all_history[start_idx:end_idx]
        
        # 渲染九宫格
        rows = [page_items[i:i+3] for i in range(0, len(page_items), 3)]
        for row in rows:
            cols = st.columns(3)
            for idx, item in enumerate(row):
                with cols[idx]:
                    img_path = item.get("img_path_in_repo", f"images/{item['id']}.png")
                    
                    # 优先加载本地缩略图 > 本地原图 > 远程
                    thumb_path = os.path.join(THUMBNAIL_DIR, f"{item['id']}.jpg")
                    local_full = os.path.join(LOCAL_CACHE_DIR, img_path)
                    
                    display_img = None
                    if os.path.exists(thumb_path): display_img = thumb_path
                    elif os.path.exists(local_full): display_img = local_full
                    
                    if display_img:
                        st.markdown(f'<div class="gallery-img-container">', unsafe_allow_html=True)
                        st.image(display_img, use_container_width=True)
                        st.markdown('</div>', unsafe_allow_html=True)
                        
                        b1, b2 = st.columns([3, 1])
                        if b1.button(f"👁️ {item['prompt_zh'][:10]}...", key=f"v_{item['id']}"):
                            # 点击查看详情时，确保原图存在
                            full_data = sm.download_image(img_path)
                            if full_data:
                                view_image_dialog(full_data, item)
                            else:
                                st.toast("正在从云端下载原图，请稍后...", icon="☁️")
                        
                        if b2.button("🗑️", key=f"d_{item['id']}"):
                            # 暂未实现物理删除，仅演示
                            st.toast("删除功能开发中", icon="🚧")
                    else:
                        st.markdown("""
                        <div class="skeleton-loader" style="height:200px; border-radius:8px;">
                            ☁️ 图片同步中...
                        </div>
                        """, unsafe_allow_html=True)
                        # 触发后台下载
                        sm.ensure_image_on_disk(img_path, item['id'])

        render_pagination(total_pages, "gal_main")

# ==================== Tab 3: 自动化任务 ====================
with tab_batch:
    render_automation_dashboard()

# ==================== Tab 4: 系统设置 ====================
with tab_settings:
    st.subheader("🔑 Token 管理 (Pollinations)")
    
    # Token 列表管理
    if st.session_state.get("tokens_loaded") != tm.last_save_time:
        st.session_state.tokens_loaded = tm.last_save_time
        
    with st.container(border=True):
        c_add1, c_add2, c_add3 = st.columns([2, 4, 1])
        new_alias = c_add1.text_input("备注名", placeholder="例如: MyKey_01")
        new_key = c_add2.text_input("Token Key", placeholder="pollinations_sk_...", type="password")
        if c_add3.button("➕ 添加", use_container_width=True):
            if new_alias and new_key:
                tm.add_token(new_alias, new_key)
                st.success("添加成功")
                time.sleep(1)
                st.rerun()

        st.divider()
        
        # Token 状态展示
        for idx, t in enumerate(tm.tokens):
            c1, c2, c3, c4 = st.columns([2, 2, 2, 1])
            c1.markdown(f"**{t['alias']}**")
            
            # 状态指示灯
            status_map = {
                "active": "🟢 正常",
                "exhausted": "🔴 余额不足",
                "invalid": "🚫 无效/被封"
            }
            c2.markdown(status_map.get(t['status'], "⚪ 未知"))
            c3.markdown(f"💰 `${t.get('balance', 0.0):.4f}`")
            
            if c4.button("🗑️", key=f"del_t_{idx}"):
                tm.delete_token(idx)
                st.rerun()

    st.subheader("☁️ WebDAV 备份节点")
    with st.expander("添加新节点", expanded=False):
        with st.form("add_webdav"):
            n = st.text_input("节点名称")
            h = st.text_input("服务器地址 (URL)")
            u = st.text_input("账号")
            p = st.text_input("密码", type="password")
            r = st.text_input("根路径", value="/")
            en = st.checkbox("立即启用", value=True)
            if st.form_submit_button("保存"):
                wm.add_connection(n, h, u, p, r, en)
                st.rerun()
    
    if wm.configs:
        for c in wm.configs:
            with st.container(border=True):
                c1, c2, c3 = st.columns([0.5, 2.5, 1])
                is_enabled = c.get("enabled", True)
                c1.markdown("🟢" if is_enabled is True else "⚪")
                c2.markdown(f"**{c['name']}**")
                cc1, cc2 = c3.columns(2)
                if cc1.button("✏️", key=f"ed_{c['id']}"): edit_webdav_dialog(c)
                if cc2.button("🗑️", key=f"dl_{c['id']}"): 
                    wm.delete_connection(c['id'])
                    st.rerun()

    st.divider()
    st.subheader("🔐 安全")
    with st.expander("修改密码"):
        with st.form("pw"):
            p1 = st.text_input("新密码", type="password")
            p2 = st.text_input("确认", type="password")
            if st.form_submit_button("修改"):
                if p1==p2: 
                    mgrs["am"].update_password(st.session_state.username, p1)
                    st.success("成功")
                    time.sleep(1)
                    st.rerun()
                else: st.error("不一致")