xhs_factory/services/connection.py

"""
services/connection.py
LLM 提供商管理、SD 连接、MCP 连接、XHS 登录等服务函数
"""
import os
import re
import logging

import gradio as gr
from PIL import Image

from .config_manager import ConfigManager
from .llm_service import LLMService
from .sd_service import SDService, get_model_profile_info
from .mcp_client import get_mcp_client

logger = logging.getLogger("autobot")
cfg = ConfigManager()

def _get_llm_config() -> tuple[str, str, str]:
    """获取当前激活 LLM 的 (api_key, base_url, model)"""
    p = cfg.get_active_llm()
    if p:
        return p["api_key"], p["base_url"], cfg.get("model", "")
    return "", "", ""


def connect_llm(provider_name):
    """连接选中的 LLM 提供商并获取模型列表"""
    if not provider_name:
        return gr.update(choices=[], value=None), "⚠️ 请先选择或添加 LLM 提供商"
    cfg.set_active_llm(provider_name)
    p = cfg.get_active_llm()
    if not p:
        return gr.update(choices=[], value=None), "❌ 未找到该提供商配置"
    try:
        svc = LLMService(p["api_key"], p["base_url"])
        models = svc.get_models()
        if models:
            return (
                gr.update(choices=models, value=models[0]),
                f"✅ 已连接「{provider_name}」，加载 {len(models)} 个模型",
            )
        else:
            # API 无法获取模型列表，保留手动输入
            current_model = cfg.get("model", "")
            return (
                gr.update(choices=[current_model] if current_model else [], value=current_model or None),
                f"⚠️ 已连接「{provider_name}」，但未获取到模型列表，请手动输入模型名",
            )
    except Exception as e:
        logger.error("LLM 连接失败: %s", e)
        current_model = cfg.get("model", "")
        return (
            gr.update(choices=[current_model] if current_model else [], value=current_model or None),
            f"❌ 连接「{provider_name}」失败: {e}",
        )


def add_llm_provider(name, api_key, base_url):
    """添加新的 LLM 提供商"""
    msg = cfg.add_llm_provider(name, api_key, base_url)
    names = cfg.get_llm_provider_names()
    active = cfg.get("active_llm", "")
    return (
        gr.update(choices=names, value=active),
        msg,
    )


def remove_llm_provider(provider_name):
    """删除 LLM 提供商"""
    if not provider_name:
        return gr.update(choices=cfg.get_llm_provider_names(), value=cfg.get("active_llm", "")), "⚠️ 请先选择要删除的提供商"
    msg = cfg.remove_llm_provider(provider_name)
    names = cfg.get_llm_provider_names()
    active = cfg.get("active_llm", "")
    return (
        gr.update(choices=names, value=active),
        msg,
    )


def on_provider_selected(provider_name):
    """切换 LLM 提供商时更新显示信息"""
    if not provider_name:
        return "未选择提供商"
    for p in cfg.get_llm_providers():
        if p["name"] == provider_name:
            cfg.set_active_llm(provider_name)
            masked_key = p["api_key"][:8] + "***" if len(p["api_key"]) > 8 else "***"
            return f"**{provider_name}**  \nAPI Key: `{masked_key}`  \nBase URL: `{p['base_url']}`"
    return "未找到该提供商"


# ==================================================
#                  Tab 1: 内容创作
# ==================================================


def connect_sd(sd_url):
    """连接 SD 并获取模型列表"""
    try:
        svc = SDService(sd_url)
        ok, msg = svc.check_connection()
        if ok:
            models = svc.get_models()
            cfg.set("sd_url", sd_url)
            first = models[0] if models else None
            info = get_model_profile_info(first) if first else "未检测到模型"
            return gr.update(choices=models, value=first), f"✅ {msg}", info
        return gr.update(choices=[]), f"❌ {msg}", ""
    except Exception as e:
        logger.error("SD 连接失败: %s", e)
        return gr.update(choices=[]), f"❌ SD 连接失败: {e}", ""


def on_sd_model_change(model_name):
    """SD 模型切换时显示模型档案信息"""
    if not model_name:
        return "未选择模型"
    return get_model_profile_info(model_name)


def check_mcp_status(mcp_url):
    """检查 MCP 连接状态"""
    try:
        client = get_mcp_client(mcp_url)
        ok, msg = client.check_connection()
        if ok:
            cfg.set("mcp_url", mcp_url)
            return f"✅ MCP 服务正常 - {msg}"
        return f"❌ {msg}"
    except Exception as e:
        return f"❌ MCP 连接失败: {e}"


# ==================================================
#              小红书账号登录
# ==================================================


def get_login_qrcode(mcp_url):
    """获取小红书登录二维码"""
    try:
        client = get_mcp_client(mcp_url)
        result = client.get_login_qrcode()
        if "error" in result:
            return None, f"❌ 获取二维码失败: {result['error']}"
        qr_image = result.get("qr_image")
        msg = result.get("text", "")
        if qr_image:
            return qr_image, f"✅ 二维码已生成，请用小红书 App 扫码\n{msg}"
        return None, f"⚠️ 未获取到二维码图片，MCP 返回:\n{msg}"
    except Exception as e:
        logger.error("获取登录二维码失败: %s", e)
        return None, f"❌ 获取二维码失败: {e}"


def logout_xhs(mcp_url):
    """退出登录：清除 cookies 并重置本地 token"""
    try:
        client = get_mcp_client(mcp_url)
        result = client.delete_cookies()
        if "error" in result:
            return f"❌ 退出失败: {result['error']}"
        cfg.set("xsec_token", "")
        client._reset()
        return "✅ 已退出登录，可以重新扫码登录"
    except Exception as e:
        logger.error("退出登录失败: %s", e)
        return f"❌ 退出失败: {e}"


def _auto_fetch_xsec_token(mcp_url) -> str:
    """从推荐列表自动获取一个有效的 xsec_token"""
    try:
        client = get_mcp_client(mcp_url)
        entries = client.list_feeds_parsed()
        for e in entries:
            token = e.get("xsec_token", "")
            if token:
                return token
    except Exception as e:
        logger.warning("自动获取 xsec_token 失败: %s", e)
    return ""


def check_login(mcp_url):
    """检查登录状态，登录成功后自动获取 xsec_token 并保存"""
    try:
        client = get_mcp_client(mcp_url)
        result = client.check_login_status()
        if "error" in result:
            return f"❌ {result['error']}", gr.update(), gr.update()
        text = result.get("text", "")
        if "未登录" in text:
            return f"🔴 {text}", gr.update(), gr.update()

        # 登录成功 → 自动获取 xsec_token
        token = _auto_fetch_xsec_token(mcp_url)
        if token:
            cfg.set("xsec_token", token)
            logger.info("自动获取 xsec_token 成功")
            return (
                f"🟢 {text}\n\n✅ xsec_token 已自动获取并保存",
                gr.update(value=cfg.get("my_user_id", "")),
                gr.update(value=token),
            )
        return f"🟢 {text}\n\n⚠️ 自动获取 xsec_token 失败，请手动刷新", gr.update(), gr.update()
    except Exception as e:
        return f"❌ 检查登录状态失败: {e}", gr.update(), gr.update()


def save_my_user_id(user_id_input):
    """保存用户 ID (验证 24 位十六进制格式)"""
    uid = (user_id_input or "").strip()
    if not uid:
        cfg.set("my_user_id", "")
        return "⚠️ 已清除用户 ID"
    if not re.match(r'^[0-9a-fA-F]{24}$', uid):
        return (
            "❌ 格式错误！用户 ID 应为 24 位十六进制字符串\n"
            f"你输入的: `{uid}` ({len(uid)} 位)\n\n"
            "💡 如果你输入的是小红书号 (纯数字如 18688457507)，那不是 userId。"
        )
    cfg.set("my_user_id", uid)
    return f"✅ 用户 ID 已保存: `{uid}`"


# ================= 头像/换脸管理 =================

def upload_face_image(img):
    """上传并保存头像图片"""
    if img is None:
        return None, "❌ 请上传头像图片"
    try:
        if isinstance(img, str) and os.path.isfile(img):
            img = Image.open(img).convert("RGB")
        elif not isinstance(img, Image.Image):
            return None, "❌ 无法识别图片格式"
        path = SDService.save_face_image(img)
        return img, f"✅ 头像已保存至 {os.path.basename(path)}"
    except Exception as e:
        return None, f"❌ 保存失败: {e}"


def load_saved_face_image():
    """加载已保存的头像"""
    img = SDService.load_face_image()
    if img:
        return img, "✅ 已加载保存的头像"
    return None, "ℹ️ 尚未设置头像"