src/ui.py

import streamlit as st
import time
import math
import uuid
import datetime
import os
import re

# 引用后端模块
from config import CUSTOM_CSS, PROVIDER_MAP, SIZE_PRESETS, QUALITY_PRESETS, ITEMS_PER_PAGE, LOCAL_CACHE_DIR, CONTAINER_DISK_QUOTA_GB, THUMBNAIL_DIR
from utils import get_disk_info, PSUTIL_AVAILABLE
from services import TokenManager, AuthManager, WebDAVManager
from logic import CoreGenerator, StorageManager, AutoTaskManager
from utils import PromptLoader, LocalUpscaler, SocialPosterGenerator

if PSUTIL_AVAILABLE:
    import psutil

# 依赖注入容器
@st.cache_resource(show_spinner=False)
def get_managers():
    return {
        "pl": PromptLoader(),
        "tm": TokenManager(),
        "cg": CoreGenerator(),
        "sm": StorageManager(),
        "wm": WebDAVManager(),
        "atm": AutoTaskManager(),
        "up": LocalUpscaler(),
        "pg": SocialPosterGenerator(),
        "am": AuthManager()
    }

# --- 底层容器监控工具类 (V30.19 New) ---
class ContainerMonitor:
    @staticmethod
    def _read_file(path):
        """安全读取文件内容，失败返回 None"""
        if not os.path.exists(path): return None
        try:
            with open(path, 'r') as f: return f.read().strip()
        except: return None

    @staticmethod
    def get_memory_stats():
        """获取容器内存使用情况 (适配 v1/v2)"""
        usage = 0
        limit = 0
        
        # 1. 尝试 Cgroup V2
        if os.path.exists('/sys/fs/cgroup/memory.current'):
            usage = int(ContainerMonitor._read_file('/sys/fs/cgroup/memory.current') or 0)
            limit_str = ContainerMonitor._read_file('/sys/fs/cgroup/memory.max')
            if limit_str == "max": limit = 10**15 # 无限制
            else: limit = int(limit_str or 0)
            
        # 2. 尝试 Cgroup V1
        elif os.path.exists('/sys/fs/cgroup/memory/memory.usage_in_bytes'):
            usage = int(ContainerMonitor._read_file('/sys/fs/cgroup/memory/memory.usage_in_bytes') or 0)
            limit = int(ContainerMonitor._read_file('/sys/fs/cgroup/memory/memory.limit_in_bytes') or 0)
        
        # Fallback: 如果读不到或无限制，尝试用 psutil 读物理内存作为上限
        if limit == 0 or limit > 10**14: 
            if PSUTIL_AVAILABLE: limit = psutil.virtual_memory().total
            else: limit = 16 * 1024**3 # 假定 16G
            
        return usage, limit

    @staticmethod
    def get_cpu_counter():
        """获取 CPU 累计使用时间 (纳秒)"""
        # 1. Cgroup V2
        if os.path.exists('/sys/fs/cgroup/cpu.stat'):
            content = ContainerMonitor._read_file('/sys/fs/cgroup/cpu.stat')
            if content:
                # 提取 usage_usec 并转为纳秒 (*1000)
                match = re.search(r'usage_usec\s+(\d+)', content)
                if match: return int(match.group(1)) * 1000
        
        # 2. Cgroup V1
        path_v1 = '/sys/fs/cgroup/cpu,cpuacct/cpuacct.usage' # 有些系统挂载点不同
        if os.path.exists(path_v1):
            return int(ContainerMonitor._read_file(path_v1) or 0)
            
        # Fallback: 虽然不准确，但总比没有好
        if PSUTIL_AVAILABLE: return int(time.time() * 1e9) 
        return 0

    @staticmethod
    def get_cpu_quota():
        """获取容器 CPU 配额 (核数)"""
        # 1. Cgroup V2
        if os.path.exists('/sys/fs/cgroup/cpu.max'):
            content = ContainerMonitor._read_file('/sys/fs/cgroup/cpu.max')
            if content:
                parts = content.split()
                if parts[0] == 'max': return os.cpu_count() or 1
                try: return int(parts[0]) / int(parts[1])
                except: pass
                
        # 2. Cgroup V1
        path_quota = '/sys/fs/cgroup/cpu/cpu.cfs_quota_us'
        path_period = '/sys/fs/cgroup/cpu/cpu.cfs_period_us'
        if os.path.exists(path_quota):
            quota = int(ContainerMonitor._read_file(path_quota) or -1)
            period = int(ContainerMonitor._read_file(path_period) or 100000)
            if quota == -1: return os.cpu_count() or 1
            return quota / period
            
        return os.cpu_count() or 1

# --- 通用弹窗组件 ---

@st.dialog("作品详情", width="large")
def view_image_dialog(image_bytes, record):
    if image_bytes:
        st.image(image_bytes, width="stretch")
        st.markdown("### 📝 提示词信息")
        st.info(record.get('prompt_zh', '无中文描述'))
        with st.expander("查看英文原文 (Prompt)"): 
            st.code(record.get('prompt_en', 'N/A'))
        
        st.markdown("### ℹ️ 元数据")
        c1, c2, c3 = st.columns(3)
        c1.markdown(f"<p style='font-size:14px; color:gray; margin-bottom:0;'>时间</p><p style='font-size:16px; font-weight:bold;'>{record['timestamp'].split(' ')[0]}</p>", unsafe_allow_html=True)
        c2.markdown(f"<p style='font-size:14px; color:gray; margin-bottom:0;'>模型</p><p style='font-size:16px; font-weight:bold;'>{record.get('model', 'Unknown')}</p>", unsafe_allow_html=True)
        c3.markdown(f"<p style='font-size:14px; color:gray; margin-bottom:0;'>风格</p><p style='font-size:16px; font-weight:bold;'>{record.get('style', 'None')}</p>", unsafe_allow_html=True)

@st.dialog("编辑 WebDAV")
def edit_webdav_dialog(conf):
    mgrs = get_managers()
    wm = mgrs["wm"]
    opts = conf["options"]
    with st.form("edit"):
        n = st.text_input("备注名称", value=conf["name"])
        h = st.text_input("服务器地址 (URL)", value=opts["webdav_hostname"])
        u = st.text_input("账号", value=opts["webdav_login"])
        p = st.text_input("密码", value=opts["webdav_password"], type="password")
        r = st.text_input("根路径 (可选)", value=opts.get("webdav_root", "/"))
        is_enabled = st.checkbox("✅ 启用此节点", value=conf.get("enabled", True))
        
        c1, c2 = st.columns(2)
        if c1.form_submit_button("测试连接", width="stretch"):
            ok, msg = wm.test_connection(h, u, p, r)
            if ok: st.success(f"✅ {msg}") 
            else: st.error(msg)
                
        if c2.form_submit_button("保存修改", type="primary", width="stretch"):
            ok, msg = wm.test_connection(h, u, p, r)
            if ok: 
                wm.update_connection(conf["id"], n, h, u, p, r, is_enabled)
                st.success("✅ 更新成功")
                time.sleep(1)
                st.rerun()
            else: 
                st.error(f"连接测试失败: {msg}")

def render_pagination(total_pages, key_prefix):
    c_space_l, c_first, c_prev, c_info, c_next, c_last, c_input, c_jump, c_space_r = st.columns([3, 0.7, 0.7, 1.2, 0.7, 0.7, 1.0, 0.7, 3])
    current = st.session_state.gallery_page
    
    with c_first:
        if st.button("首页", key=f"{key_prefix}_first", disabled=current==1, width="stretch"):
            st.session_state.gallery_page = 1
            st.rerun()
    with c_prev:
        if st.button("上一页", key=f"{key_prefix}_prev", disabled=current==1, width="stretch"):
            st.session_state.gallery_page = max(1, current - 1)
            st.rerun()
    with c_info:
        st.markdown(f"<div class='pagination-info'>{current} / {total_pages}</div>", unsafe_allow_html=True)
    with c_next:
        if st.button("下一页", key=f"{key_prefix}_next", disabled=current==total_pages, width="stretch"):
            st.session_state.gallery_page = min(total_pages, current + 1)
            st.rerun()
    with c_last:
        if st.button("尾页", key=f"{key_prefix}_last", disabled=current==total_pages, width="stretch"):
            st.session_state.gallery_page = total_pages
            st.rerun()
    with c_input:
        jump_val = st.number_input("页码", min_value=1, max_value=total_pages, value=current, key=f"{key_prefix}_input", label_visibility="collapsed")
    with c_jump:
        if st.button("跳转", key=f"{key_prefix}_jump", width="stretch"):
            if jump_val != current:
                st.session_state.gallery_page = jump_val
                st.rerun()

# --- Fragment 组件 ---

@st.fragment(run_every=2) # 缩短刷新间隔以获得实时的CPU数据
def render_header():
    """
    [V30.19] 容器原生资源监控版
    使用 cgroup 文件系统读取真实的容器资源使用率
    """
    mgrs = get_managers()
    sm, tm = mgrs["sm"], mgrs["tm"]
    atm = mgrs["atm"]
    
    current_mtime = sm.get_file_mtime()
    last_known_mtime = st.session_state.get("last_mtime", 0.0)
    
    if current_mtime > last_known_mtime:
        sm.refresh_cache()
        st.session_state.last_mtime = current_mtime
        st.rerun()
    
    hist_len = len(sm.cache)
    active_tasks = len(atm.active_tasks)
    
    # ------------------ 资源计算区 ------------------
    now = time.time()
    
    # 1. 内存计算 (Snapshot)
    mem_used, mem_limit = ContainerMonitor.get_memory_stats()
    mem_percent = (mem_used / mem_limit) * 100 if mem_limit > 0 else 0
    mem_text = f"RAM: {mem_used/1024**3:.1f}G / {mem_limit/1024**3:.1f}G ({mem_percent:.0f}%)"
    
    # 2. CPU 计算 (Differential)
    # 需要保存上一次的计数器状态来计算增量
    cpu_usage_ns = ContainerMonitor.get_cpu_counter()
    cpu_cores = ContainerMonitor.get_cpu_quota()
    
    if "mon_last_time" not in st.session_state:
        st.session_state.mon_last_time = now
        st.session_state.mon_last_cpu = cpu_usage_ns
        st.session_state.mon_cpu_text = "CPU: Calc..."
        st.session_state.mon_cpu_val = 0.0
    
    time_delta = now - st.session_state.mon_last_time
    
    # 至少间隔 1s 计算一次，否则数据抖动太大
    if time_delta > 1.0:
        cpu_delta = cpu_usage_ns - st.session_state.mon_last_cpu
        # CPU Usage = (Delta NS) / (Delta Time * 1e9 * Cores) * 100
        # 如果是 2核，1秒内跑满是 2e9 ns
        cpu_p = (cpu_delta / (time_delta * 1e9 * cpu_cores)) * 100
        
        st.session_state.mon_cpu_text = f"CPU: {cpu_p:.1f}% ({cpu_cores:.1f}C)"
        st.session_state.mon_cpu_val = min(cpu_p / 100, 1.0)
        
        # 更新状态
        st.session_state.mon_last_time = now
        st.session_state.mon_last_cpu = cpu_usage_ns

    # 3. 磁盘计算 (Snapshot)
    used_gb, free_gb_phys, total_phys, percent_phys = get_disk_info(LOCAL_CACHE_DIR)
    quota_gb = CONTAINER_DISK_QUOTA_GB
    disk_percent = min((used_gb / quota_gb) * 100, 100.0)
    disk_text = f"Disk: {used_gb:.1f}G / {quota_gb}G ({disk_percent:.0f}%)"

    # ------------------ UI 渲染 ------------------
    c_left, c_right = st.columns([6, 1])
    with c_left:
        st.markdown(f"""
            <div class="header-wrapper">
                <span class="header-title">🎨 AI 艺术 | {st.session_state.username}</span>
                <span class="header-badge">🖼️ {hist_len}</span>
                <span class="header-badge {'active' if active_tasks > 0 else ''}">{'🟢 自动运行中' if active_tasks > 0 else '⚪ 自动任务待机'}</span>
            </div>
        """, unsafe_allow_html=True)
        
        # 资源进度条 (3列)
        col_cpu, col_mem, col_disk = st.columns(3)
        
        with col_cpu:
            st.progress(st.session_state.mon_cpu_val, text=st.session_state.mon_cpu_text)
        
        with col_mem:
            st.progress(min(mem_percent/100, 1.0), text=mem_text)
        
        with col_disk:
            st.progress(min(disk_percent/100, 1.0), text=disk_text)

    with c_right:
        if st.button("🚪 退出", width="stretch"): 
            st.session_state.logged_in=False
            st.rerun()

@st.fragment
def render_automation_dashboard():
    """自动化任务面板"""
    mgrs = get_managers()
    sm, tm, wm, atm, cg = mgrs["sm"], mgrs["tm"], mgrs["wm"], mgrs["atm"], mgrs["cg"]
    
    col_auto_l, col_auto_r = st.columns([1, 2], gap="medium")
    
    with col_auto_l:
        st.subheader("⚙️ 任务控制台")
        task_options = ["➕ 新建任务"] + [t["name"] for t in st.session_state.saved_tasks]
        
        current_index = st.session_state.selected_task_index
        if current_index >= len(task_options): 
            current_index = 0
            st.session_state.selected_task_index = 0
            
        selected_task_name = st.selectbox("选择任务方案", task_options, index=current_index)
        
        current_task_data = {}
        if selected_task_name != "➕ 新建任务":
            found = [t for t in st.session_state.saved_tasks if t["name"] == selected_task_name]
            if found: current_task_data = found[0]
            
            try:
                idx = task_options.index(selected_task_name)
                if st.session_state.selected_task_index != idx:
                    st.session_state.selected_task_index = idx
                    st.rerun()
            except: pass

        # [V30.17] 移除表单隔离，使按钮可以实时响应
        task_name_input = st.text_input("任务名称", value=current_task_data.get("name", "我的新任务"))
        seeds_input = st.text_area("灵感种子池 (一行一个)", value=current_task_data.get("seeds", "Cyberpunk\nWasteland\nFuture City"), height=100)
        
        c_f1, c_f2 = st.columns(2)
        try: default_idx = list(PROVIDER_MAP.keys()).index(current_task_data.get("provider", "Pollinations AI (推荐)"))
        except: default_idx = 0
        
        iv = c_f1.number_input("周期(s)", min_value=60, value=max(60, current_task_data.get("interval", 3600)))
        pv = c_f1.selectbox("供应商", list(PROVIDER_MAP.keys()), index=default_idx)
        
        try: default_qual_idx = list(QUALITY_PRESETS.keys()).index(current_task_data.get("quality", "⚡ 标准 (Standard)"))
        except: default_qual_idx = 0
        qual = c_f2.selectbox("画质等级", list(QUALITY_PRESETS.keys()), index=default_qual_idx)
        
        c_b1, c_b2 = st.columns([1, 1])
        if c_b1.button("💾 保存/更新", type="primary", width="stretch"):
            if not task_name_input: 
                st.error("名称不能为空")
            else:
                new_task = {
                    "id": current_task_data.get("id", str(uuid.uuid4())),
                    "name": task_name_input,
                    "seeds": seeds_input,
                    "interval": iv,
                    "provider": pv,
                    "quality": qual
                }
                new_list = [t for t in st.session_state.saved_tasks if t["id"] != new_task["id"]]
                new_list.append(new_task)
                st.session_state.saved_tasks = new_list
                sm.save_tasks(new_list)
                st.success("已保存")
                time.sleep(1)
                st.rerun()
        
        if c_b2.button("🗑️ 删除", type="secondary", width="stretch"):
            if selected_task_name == "➕ 新建任务": st.error("无法删除")
            else:
                new_list = [t for t in st.session_state.saved_tasks if t["id"] != current_task_data["id"]]
                st.session_state.saved_tasks = new_list
                st.session_state.selected_task_index = 0
                sm.save_tasks(new_list)
                st.rerun()

        st.divider()
        
        current_task_id = current_task_data.get("id")
        if selected_task_name == "➕ 新建任务":
            if st.button("🚀 启动此临时任务 (Run Once)", type="primary", width="stretch"):
                config = {"seeds": [s.strip() for s in seeds_input.split('\n') if s.strip()], "interval": iv, "provider": pv, "quality": qual}
                temp_id = str(uuid.uuid4())
                ok, msg = atm.start(config, temp_id, task_name_input, sm, tm, wm, cg)
                if ok: st.success("已启动"); time.sleep(1); st.rerun()
                else: st.error(msg)
        else:
            is_running_this = current_task_id in atm.active_tasks
            if is_running_this:
                if st.button("🛑 停止此任务", type="primary", width="stretch"):
                    atm.stop(current_task_id)
                    st.rerun()
            else:
                is_full = len(atm.active_tasks) >= atm.MAX_CONCURRENT_TASKS
                btn_label = "⚠️ 任务槽已满" if is_full else f"🚀 启动任务: {task_name_input}"
                if st.button(btn_label, type="primary", disabled=is_full, width="stretch"):
                    config = {"seeds": [s.strip() for s in seeds_input.split('\n') if s.strip()], "interval": iv, "provider": pv, "quality": qual}
                    ok, msg = atm.start(config, current_task_id, task_name_input, sm, tm, wm, cg)
                    if ok: st.success("已启动"); time.sleep(1); st.rerun()
                    else: st.error(msg)

    with col_auto_r:
        st.subheader("📊 运行监控中心")
        c_m1, c_m2 = st.columns(2)
        active_items = list(atm.active_tasks.items())
        
        with c_m1:
            with st.container(border=True):
                if len(active_items) > 0:
                    tid, tdata = active_items[0]
                    st.markdown(f"**🟢 {tdata['name']}**")
                    st.caption(f"Status: {tdata['status']}")
                    if st.button("🛑 停止", key=f"stop_{tid}", width="stretch"):
                        atm.stop(tid)
                        st.rerun()
                else:
                    st.markdown(f"**⚪ 通道 1 空闲**")
                    st.markdown("<br>", unsafe_allow_html=True)
                    st.markdown("<br>", unsafe_allow_html=True)
        
        with c_m2:
            with st.container(border=True):
                if len(active_items) > 1:
                    tid, tdata = active_items[1]
                    st.markdown(f"**🟢 {tdata['name']}**")
                    st.caption(f"Status: {tdata['status']}")
                    if st.button("🛑 停止", key=f"stop_{tid}", width="stretch"):
                        atm.stop(tid)
                        st.rerun()
                else:
                    st.markdown(f"**⚪ 通道 2 空闲**")
                    st.markdown("<br>", unsafe_allow_html=True)
                    st.markdown("<br>", unsafe_allow_html=True)

        st.divider()
        st.markdown("### 📜 实时运行日志")
        log_container = st.container(height=500, border=True)
        if not atm.logs:
            log_container.info("暂无日志记录...")
        else:
            for log in atm.logs:
                color = "green" if log["level"]=="success" else "red" if log["level"]=="error" else "orange" if log["level"]=="warning" else "blue"
                log_container.markdown(f":{color}[{log['ts']}] {log['msg']}")

def login_page():
    st.markdown("<br><br><br>", unsafe_allow_html=True)
    st.markdown("<h1 style='text-align: center;'>🔐 AI Studio</h1>", unsafe_allow_html=True)
    c1, c2, c3 = st.columns([1, 1.2, 1])
    with c2:
        with st.form("login"):
            u = st.text_input("用户")
            p = st.text_input("密码", type="password")
            if st.form_submit_button("登录", type="primary", width="stretch"):
                am = get_managers()["am"]
                ok, must = am.login(u, p)
                if ok:
                    st.session_state.logged_in = True
                    st.session_state.username = u
                    st.session_state.must_change_pw = must
                    st.rerun()
                else: st.error("错误")

def change_pw_page():
    st.title("首次登录需修改密码")
    with st.form("force_pw"):
        p1 = st.text_input("新密码", type="password")
        p2 = st.text_input("确认", type="password")
        if st.form_submit_button("修改"):
            if p1==p2:
                get_managers()["am"].update_password(st.session_state.username, p1)
                st.session_state.must_change_pw = False
                st.rerun()
            else: st.error("不一致")