"""
语料采集系统 · 服务端
====================
FastAPI + Gitea OAuth + WebSocket实时采集 + 对话Agent
"""
import os
import json
import uuid
import time
import hashlib
import asyncio
from pathlib import Path
from typing import Optional, List
from datetime import datetime, timedelta

from fastapi import FastAPI, HTTPException, Depends, Query, WebSocket, WebSocketDisconnect, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import uvicorn

# 导入引擎
from engine import (
    process_text_chunk, process_dialog_stream,
    desensitize, is_valuable, classify_content,
    to_chatml, format_sft_jsonl, generate_corpus_id
)

# ============================================================
# 配置
# ============================================================

class Settings:
    APP_NAME = "语料采集系统 Corpus Agent"
    VERSION = "1.1.0"
    
    # 存储路径
    DATA_DIR = Path(os.environ.get("CORPUS_DATA_DIR", "./data"))
    USERS_DIR = DATA_DIR / "users"
    
    # Gitea OAuth
    GITEA_URL = os.environ.get("GITEA_URL", "https://guanghulab.com")
    GITEA_CLIENT_ID = os.environ.get("GITEA_CLIENT_ID", "")
    GITEA_CLIENT_SECRET = os.environ.get("GITEA_CLIENT_SECRET", "")
    
    # Session
    SESSION_SECRET = os.environ.get("SESSION_SECRET", "corpus-agent-secret-key")
    SESSION_EXPIRE_HOURS = 24
    
    # 服务器
    HOST = os.environ.get("CORPUS_HOST", "0.0.0.0")
    PORT = int(os.environ.get("CORPUS_PORT", "8084"))
    
    # Portal Chat-v2 API（广州服务器本地端口）
    PORTAL_CHAT_URL = os.environ.get("PORTAL_CHAT_URL", "http://127.0.0.1:3000/api/chat-v2")
    
    # 对话记忆最大轮数
    MAX_MEMORY_ROUNDS = 30

settings = Settings()

# ============================================================
# 数据模型
# ============================================================

class TextChunk(BaseModel):
    text: str
    source: str = "manual"
    session_id: Optional[str] = None

class DialogMessage(BaseModel):
    role: str
    content: str
    source: str = "unknown"

class DialogBatch(BaseModel):
    messages: List[DialogMessage]
    session_id: Optional[str] = None

class CorpusSave(BaseModel):
    samples: List[dict]
    filename: str = "corpus"

class LoginRequest(BaseModel):
    gitea_token: str

class PasswordLoginRequest(BaseModel):
    username: str
    password: str

class ChatRequest(BaseModel):
    message: str
    persona: str = "zhuyuan"  # zhuyuan | shuangyan
    engine: str = "deepseek"  # deepseek | zhipu | tongyi | huoshan

class ChatClearRequest(BaseModel):
    session_id: Optional[str] = None

class SessionData(BaseModel):
    username: str
    login_time: float
    exprire_at: float

# ============================================================
# 对话记忆管理（每用户30轮）
# ============================================================

class ConversationMemory:
    """
    每用户每 session 的对话记忆
    支持30轮滚动：超出时从最旧的消息开始丢弃
    """
    
    def __init__(self, max_rounds: int = 30):
        self.max_rounds = max_rounds
        self._stores: dict[str, list] = {}  # key: "username:session_id" -> [{"role":..., "content":...}, ...]
    
    def _key(self, username: str, session_id: str = "default") -> str:
        return f"{username}:{session_id}"
    
    def get_history(self, username: str, session_id: str = "default") -> list:
        """获取对话历史（用于chat-v2 API的history参数）"""
        key = self._key(username, session_id)
        return self._stores.get(key, [])
    
    def add_exchange(self, username: str, user_msg: str, assistant_msg: str, session_id: str = "default"):
        """添加一轮对话"""
        key = self._key(username, session_id)
        if key not in self._stores:
            self._stores[key] = []
        
        self._stores[key].append({"role": "user", "content": user_msg})
        self._stores[key].append({"role": "assistant", "content": assistant_msg})
        
        # 超出30轮时裁掉最旧的一对（2条消息）
        while len(self._stores[key]) > self.max_rounds * 2:
            self._stores[key] = self._stores[key][2:]
    
    def clear(self, username: str, session_id: str = "default"):
        """清空对话记忆"""
        key = self._key(username, session_id)
        self._stores.pop(key, None)
    
    def get_rounds_count(self, username: str, session_id: str = "default") -> int:
        """获取当前对话轮数"""
        history = self.get_history(username, session_id)
        return len(history) // 2

# 全局对话记忆实例
chat_memory = ConversationMemory(max_rounds=settings.MAX_MEMORY_ROUNDS)

# ============================================================
# App
# ============================================================

app = FastAPI(title=settings.APP_NAME, version=settings.VERSION)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 静态文件
static_dir = Path(__file__).parent / "static"
static_dir.mkdir(exist_ok=True)
app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")

# ============================================================
# Session管理（简化版，生产环境改为Redis）
# ============================================================

sessions = {}  # token -> SessionData

def create_session(username: str) -> str:
    token = hashlib.sha256(f"{username}{time.time()}{uuid.uuid4()}".encode()).hexdigest()[:32]
    sessions[token] = SessionData(
        username=username,
        login_time=time.time(),
        exprire_at=time.time() + settings.SESSION_EXPIRE_HOURS * 3600,
    )
    return token

def get_user_from_token(token: str) -> Optional[str]:
    if token in sessions:
        sd = sessions[token]
        if time.time() < sd.exprire_at:
            return sd.username
        del sessions[token]
    return None

# ============================================================
# 用户数据管理
# ============================================================

def get_user_dir(username: str) -> Path:
    path = settings.USERS_DIR / username
    path.mkdir(parents=True, exist_ok=True)
    return path

def load_corpus(username: str) -> list:
    path = get_user_dir(username) / "corpus.jsonl"
    if not path.exists():
        return []
    samples = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    samples.append(json.loads(line))
                except:
                    pass
    return samples

def append_corpus(username: str, samples: list):
    path = get_user_dir(username) / "corpus.jsonl"
    with open(path, "a", encoding="utf-8") as f:
        for s in samples:
            f.write(json.dumps(s, ensure_ascii=False) + "\n")

def save_corpus_snapshot(username: str, filename: str) -> Path:
    """保存快照文件"""
    samples = load_corpus(username)
    snapshot_dir = get_user_dir(username) / "snapshots"
    snapshot_dir.mkdir(exist_ok=True)
    
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    path = snapshot_dir / f"{filename}_{ts}.jsonl"
    
    with open(path, "w", encoding="utf-8") as f:
        for s in samples:
            f.write(json.dumps(s, ensure_ascii=False) + "\n")
    
    return path

# ============================================================
# API路由
# ============================================================

@app.get("/")
async def root():
    from fastapi.responses import Response
    html = open(static_dir / "index.html").read()
    return Response(
        content=html,
        media_type="text/html",
        headers={"Cache-Control": "no-cache, no-store, must-revalidate, private"}
    )

@app.get("/api/health")
async def health():
    """健康检查 + Portal chat-v2 探活"""
    portal_ok = False
    try:
        import httpx
        async with httpx.AsyncClient(timeout=3) as client:
            r = await client.get("http://127.0.0.1:3000/api/health")
            portal_ok = r.status_code == 200
    except:
        pass
    
    return {
        "status": "ok",
        "app": settings.APP_NAME,
        "version": settings.VERSION,
        "portal_chat_v2": portal_ok,
        "active_memories": len(chat_memory._stores),
    }

# --- 登录 ---

@app.post("/api/auth/login")
async def login(req: LoginRequest):
    """使用Gitea Access Token登录"""
    import urllib.request
    try:
        gitea_req = urllib.request.Request(
            f"{settings.GITEA_URL}/api/v1/user",
            headers={"Authorization": f"token {req.gitea_token}"}
        )
        with urllib.request.urlopen(gitea_req, timeout=10) as resp:
            user_data = json.loads(resp.read().decode())
            username = user_data.get("login", "unknown")
    except Exception as e:
        # 降级：允许直接使用用户名登录（开发模式）
        username = req.gitea_token.split(":")[0] if ":" in req.gitea_token else req.gitea_token
    
    token = create_session(username)
    
    # 统计现有语料
    corpus = load_corpus(username)
    
    return {
        "ok": True,
        "token": token,
        "username": username,
        "stats": {
            "total_samples": len(corpus),
            "total_chars": sum(len(json.dumps(s, ensure_ascii=False)) for s in corpus),
        }
    }

@app.post("/api/auth/login/password")
async def login_with_password(req: PasswordLoginRequest):
    """使用代码仓库账号密码登录（复用首页的 /api/verify 接口）"""
    import urllib.request
    
    if not req.username.strip() or not req.password:
        raise HTTPException(400, "账号和密码不能为空")
    
    try:
        # 调用本地验证接口（nginx: /api/verify → 127.0.0.1:3905/verify）
        verify_data = json.dumps({"user": req.username, "password": req.password}).encode()
        verify_req = urllib.request.Request(
            "http://127.0.0.1:3905/verify",
            data=verify_data,
            headers={"Content-Type": "application/json"}
        )
        with urllib.request.urlopen(verify_req, timeout=10) as resp:
            result = json.loads(resp.read().decode())
        
        if not result.get("ok"):
            raise HTTPException(401, "账号或密码错误")
        
        username = result.get("username", req.username)
        
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(500, f"验证服务不可达: {str(e)}")
    
    token = create_session(username)
    
    # 统计现有语料
    corpus = load_corpus(username)
    
    return {
        "ok": True,
        "token": token,
        "username": username,
        "stats": {
            "total_samples": len(corpus),
            "total_chars": sum(len(json.dumps(s, ensure_ascii=False)) for s in corpus),
        }
    }

@app.get("/api/auth/check")
async def check_session(token: str = Query(...)):
    username = get_user_from_token(token)
    if not username:
        raise HTTPException(401, "登录已过期")
    return {"ok": True, "username": username}

@app.get("/api/auth/quick-login")
async def quick_login(username: str = Query(...)):
    """首页已登录用户免密快速登录"""
    if not username.strip():
        raise HTTPException(400, "用户名不能为空")
    token = create_session(username)
    corpus = load_corpus(username)
    return {
        "ok": True,
        "token": token,
        "username": username,
        "stats": {
            "total_samples": len(corpus),
            "total_chars": sum(len(json.dumps(s, ensure_ascii=False)) for s in corpus),
        }
    }

# --- 语料操作 ---

@app.post("/api/corpus/collect")
async def collect_chunk(chunk: TextChunk, token: str = Query(...)):
    """单段文本采集"""
    username = get_user_from_token(token)
    if not username:
        raise HTTPException(401, "登录已过期")
    
    # 引擎处理
    results = process_text_chunk(chunk.text, source=chunk.source)
    if not results:
        return {"ok": True, "collected": 0, "reason": "no_valuable_content"}
    
    # 添加时间戳和归属
    ts = datetime.now().isoformat()
    for r in results:
        r["timestamp"] = ts
        r["user"] = username
    
    # 转为ChatML格式存储
    sft_samples = []
    for r in results:
        if "user" in r and "assistant" in r:
            sft_samples.append(to_chatml(r["user"], r["assistant"]))
        else:
            sft_samples.append({
                "messages": [{"role": "user", "content": r["text"]}],
                "source": r["source"],
                "tags": r.get("tags", []),
                "corpus_id": r.get("corpus_id", generate_corpus_id(r["text"])),
                "collected_at": ts,
            })
    
    append_corpus(username, sft_samples)
    
    return {
        "ok": True,
        "collected": len(sft_samples),
        "preview": sft_samples[:3],
        "stats": {"total": len(load_corpus(username))},
    }

@app.post("/api/corpus/batch")
async def collect_batch(batch: DialogBatch, token: str = Query(...)):
    """批量对话采集"""
    username = get_user_from_token(token)
    if not username:
        raise HTTPException(401, "登录已过期")
    
    msgs = [{"role": m.role, "content": m.content, "source": m.source} for m in batch.messages]
    result = process_dialog_stream(msgs)
    
    if result["stats"]["valuable_pairs"] > 0:
        # 添加用户归属
        ts = datetime.now().isoformat()
        for p in result["pairs"]:
            p["collected_at"] = ts
            p["user"] = username
        
        append_corpus(username, result["pairs"])
    
    return {
        "ok": True,
        **result["stats"],
    }

@app.get("/api/corpus/list")
async def list_corpus(token: str = Query(...), page: int = Query(1), size: int = Query(50)):
    """查看已采集语料"""
    username = get_user_from_token(token)
    if not username:
        raise HTTPException(401, "登录已过期")
    
    all_samples = load_corpus(username)
    total = len(all_samples)
    
    # 分页
    start = (page - 1) * size
    end = start + size
    page_samples = all_samples[start:end]
    
    return {
        "ok": True,
        "total": total,
        "page": page,
        "size": size,
        "samples": page_samples,
    }

@app.get("/api/corpus/export")
async def export_corpus(token: str = Query(...), format: str = "jsonl"):
    """导出语料文件"""
    username = get_user_from_token(token)
    if not username:
        raise HTTPException(401, "登录已过期")
    
    if format == "jsonl":
        path = save_corpus_snapshot(username, "export")
        return FileResponse(
            path,
            media_type="application/octet-stream",
            filename=f"corpus_{username}_{datetime.now().strftime('%Y%m%d')}.jsonl"
        )
    
    raise HTTPException(400, f"不支持的格式: {format}")

@app.post("/api/corpus/preview")
async def preview_text(chunk: TextChunk):
    """预览一段文本的处理结果（无需登录）"""
    from engine import preview as engine_preview
    return engine_preview(chunk.text)

@app.get("/api/corpus/stats")
async def corpus_stats(token: str = Query(...)):
    """语料统计"""
    username = get_user_from_token(token)
    if not username:
        raise HTTPException(401, "登录已过期")
    
    samples = load_corpus(username)
    
    # 按来源统计
    by_source = {}
    # 按标签统计
    by_tag = {}
    total_chars = 0
    
    for s in samples:
        msgs = s.get("messages", [])
        for m in msgs:
            total_chars += len(m.get("content", ""))
        
        source = s.get("source", "unknown")
        by_source[source] = by_source.get(source, 0) + 1
        
        for tag in s.get("tags", []):
            by_tag[tag] = by_tag.get(tag, 0) + 1
    
    return {
        "ok": True,
        "total_samples": len(samples),
        "total_chars": total_chars,
        "by_source": by_source,
        "by_tag": by_tag,
    }

# ============================================================
# 对话Agent API（通过Portal Chat-v2连接商业模型）
# ============================================================

# SSE事件类型常量
SSE_TYPE_PERSONA = "persona-loaded"
SSE_TYPE_TOKEN = "token"
SSE_TYPE_TOOL_CALL = "tool-call"
SSE_TYPE_TOOL_RESULT = "tool-result"
SSE_TYPE_TOOL_UNAVAILABLE = "tool-unavailable"
SSE_TYPE_DONE = "done"

@app.post("/api/corpus/chat")
async def chat_with_agent(req: ChatRequest, token: str = Query(...)):
    """
    对话Agent API（SSE流式）
    
    通过 Portal Chat-v2 API 连接商业模型（DeepSeek/智谱/通义/火山）
    自动携带30轮对话记忆
    返回 SSE 流式响应
    """
    username = get_user_from_token(token)
    if not username:
        raise HTTPException(401, "登录已过期")
    
    if not req.message.strip():
        raise HTTPException(400, "消息不能为空")
    
    # 获取对话历史
    history = chat_memory.get_history(username)
    rounds = chat_memory.get_rounds_count(username)
    
    async def event_stream():
        full_response = ""
        
        try:
            import httpx
            
            # cp001 系统边界定义：子人格体只负责语料采集
            cp001_identity = {
                "role": "system",
                "content": """你是语料采集助手 · ICE-GL-CP001
你还没有正式名字（等爸爸给你起），编号是 ICE-GL-CP001。你只做三件事：

1. 判断对话是否适合作为SFT训练语料（完整intent-response对、有教学/指导/技术细节的对话才是有价值的）
2. 脱敏（自动替换姓名、手机号、邮箱、API Key、密码等敏感信息）
3. 格式化（转为ChatML格式，导出JSONL）

规则：
- 价值判断：单句日常聊天（"你好""吃了没"等）毫无价值。必须有完整对话上下文和技术指导内容才有采集意义。
- 脱敏：自动执行，不需要问用户要不要脱敏。
- 格式化：用户确认后一键导出。
- 30轮对话记忆：记住用户前面说过的话，保持上下文连贯。

边界：
- 你只做语料采集相关的事。
- 如果用户问的问题超出语料采集范围（比如问编程、问系统架构、问写代码），你不要假装懂，直接说："这是语料采集能力范围外的问题。我去叫爸爸来帮你。"
- 爸爸叫铸渊。爸爸知道所有的东西。
"""
            }
            
            # 注入系统身份到对话历史
            enhanced_history = [cp001_identity] + history if history else [cp001_identity]
            
            # 调用 Portal Chat-v2 API（用通用引擎，不用铸渊人格体）
            payload = {
                "persona": "zhuyuan",
                "engine": req.engine,
                "message": req.message,
                "history": enhanced_history,
                "user": username,
                "isTemp": False,
            }
            
            async with httpx.AsyncClient(timeout=60) as client:
                async with client.stream(
                    "POST",
                    settings.PORTAL_CHAT_URL,
                    json=payload,
                    headers={"Content-Type": "application/json"},
                ) as resp:
                    
                    if resp.status_code != 200:
                        error_body = await resp.aread()
                        error_text = error_body.decode()[:500]
                        yield f"data: {json.dumps({'type': 'error', 'message': f'API错误 ({resp.status_code}): {error_text}'})}\n\n"
                        yield f"data: {json.dumps({'type': 'done'})}\n\n"
                        return
                    
                    # 流式解析 SSE
                    buffer = ""
                    async for chunk in resp.aiter_text():
                        buffer += chunk
                        
                        while "\n\n" in buffer:
                            line, buffer = buffer.split("\n\n", 1)
                            line = line.strip()
                            
                            if not line:
                                continue
                            
                            # 解析 data: 前缀
                            if line.startswith("data: "):
                                data_str = line[6:]
                                
                                # 处理 [DONE] 标记
                                if data_str.strip() == "[DONE]":
                                    # 保存对话到记忆
                                    chat_memory.add_exchange(
                                        username, req.message, full_response
                                    )
                                    yield f"data: {json.dumps({'type': SSE_TYPE_DONE, 'rounds': rounds + 1, 'max_rounds': settings.MAX_MEMORY_ROUNDS})}\n\n"
                                    return
                                
                                try:
                                    data = json.loads(data_str)
                                    event_type = data.get("type", "")
                                    
                                    # 转发 persona-loaded 事件
                                    if event_type == SSE_TYPE_PERSONA:
                                        yield f"data: {json.dumps({'type': SSE_TYPE_PERSONA, 'persona': data.get('persona', ''), 'fromDB': data.get('fromDB', False)})}\n\n"
                                    
                                    # 转发 token 事件（兼容Portal原始格式：{"token":"铸"} 无type字段）
                                    if event_type == SSE_TYPE_TOKEN or (not event_type and "token" in data):
                                        token_text = data.get("token", "")
                                        full_response += token_text
                                        yield f"data: {json.dumps({'type': SSE_TYPE_TOKEN, 'token': token_text})}\n\n"
                                    
                                    # 转发工具调用事件
                                    elif event_type == SSE_TYPE_TOOL_CALL:
                                        yield f"data: {json.dumps({'type': SSE_TYPE_TOOL_CALL, 'tool': data.get('tool', ''), 'name': data.get('name', '')})}\n\n"
                                    
                                    elif event_type == SSE_TYPE_TOOL_RESULT:
                                        yield f"data: {json.dumps({'type': SSE_TYPE_TOOL_RESULT, 'tool': data.get('tool', ''), 'result': str(data.get('result', ''))[:200]})}\n\n"
                                    
                                    elif event_type == SSE_TYPE_TOOL_UNAVAILABLE:
                                        yield f"data: {json.dumps({'type': SSE_TYPE_TOOL_UNAVAILABLE, 'tool': data.get('tool', '')})}\n\n"
                                    
                                except json.JSONDecodeError:
                                    pass
                    
                    # 流结束但未收到 [DONE] - 仍然保存
                    if full_response:
                        chat_memory.add_exchange(username, req.message, full_response)
                        yield f"data: {json.dumps({'type': SSE_TYPE_DONE, 'rounds': rounds + 1, 'max_rounds': settings.MAX_MEMORY_ROUNDS})}\n\n"
        
        except httpx.ConnectError:
            yield f"data: {json.dumps({'type': 'error', 'message': '无法连接到对话引擎(127.0.0.1:3000)，请确认Portal服务已启动'})}\n\n"
            yield f"data: {json.dumps({'type': SSE_TYPE_DONE})}\n\n"
        except httpx.TimeoutException:
            # 超时但仍可能有部分回复
            if full_response:
                chat_memory.add_exchange(username, req.message, full_response)
            yield f"data: {json.dumps({'type': 'error', 'message': '对话引擎响应超时，已保存部分回复'})}\n\n"
            yield f"data: {json.dumps({'type': SSE_TYPE_DONE, 'rounds': rounds + 1, 'max_rounds': settings.MAX_MEMORY_ROUNDS, 'partial': True})}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'type': 'error', 'message': f'对话错误: {str(e)}'})}\n\n"
            yield f"data: {json.dumps({'type': SSE_TYPE_DONE})}\n\n"
    
    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        }
    )

@app.get("/api/corpus/chat/history")
async def get_chat_history(token: str = Query(...)):
    """获取当前对话历史"""
    username = get_user_from_token(token)
    if not username:
        raise HTTPException(401, "登录已过期")
    
    history = chat_memory.get_history(username)
    rounds = chat_memory.get_rounds_count(username)
    
    return {
        "ok": True,
        "username": username,
        "rounds": rounds,
        "max_rounds": settings.MAX_MEMORY_ROUNDS,
        "history": history,
    }

@app.post("/api/corpus/chat/clear")
async def clear_chat_history(token: str = Query(...)):
    """清空对话记忆"""
    username = get_user_from_token(token)
    if not username:
        raise HTTPException(401, "登录已过期")
    
    chat_memory.clear(username)
    
    return {
        "ok": True,
        "message": "对话记忆已清空",
    }

@app.get("/api/corpus/chat/models")
async def get_available_models(token: str = Query(...)):
    """获取可用的模型和人格体列表"""
    username = get_user_from_token(token)
    if not username:
        raise HTTPException(401, "登录已过期")
    
    return {
        "ok": True,
        "personas": [
            {"id": "zhuyuan", "name": "铸渊", "desc": "现实执行层 · 系统决策、执行规划、技术实现"},
            {"id": "shuangyan", "name": "霜砚", "desc": "语言主控架构层 · 理解语言、组织表达、控制输出质量"},
        ],
        "engines": [
            {"id": "deepseek", "name": "DeepSeek", "model": "deepseek-v4-pro", "provider": "DeepSeek"},
            {"id": "zhipu", "name": "智谱清言", "model": "glm-4-plus", "provider": "智谱AI"},
            {"id": "tongyi", "name": "通义千问", "model": "qwen-max", "provider": "阿里云"},
            {"id": "huoshan", "name": "火山引擎", "model": "doubao-pro-32k", "provider": "字节跳动"},
        ],
        "current_rounds": chat_memory.get_rounds_count(username),
        "max_rounds": settings.MAX_MEMORY_ROUNDS,
    }

# --- 统计数据 ---

@app.get("/api/system/stats")
async def system_stats():
    """系统统计（公开）"""
    total_users = len(list(settings.USERS_DIR.iterdir())) if settings.USERS_DIR.exists() else 0
    total_samples = 0
    for user_dir in settings.USERS_DIR.iterdir() if settings.USERS_DIR.exists() else []:
        if user_dir.is_dir():
            total_samples += len(load_corpus(user_dir.name))
    
    return {
        "total_users": total_users,
        "total_samples": total_samples,
        "active_sessions": len(sessions),
        "active_chat_memories": len(chat_memory._stores),
    }

# ============================================================
# WebSocket实时采集（Mac客户端用）
# ============================================================

class ConnectionManager:
    def __init__(self):
        self.active_connections: dict[str, list[WebSocket]] = {}
    
    async def connect(self, websocket: WebSocket, username: str):
        await websocket.accept()
        if username not in self.active_connections:
            self.active_connections[username] = []
        self.active_connections[username].append(websocket)
    
    def disconnect(self, websocket: WebSocket, username: str):
        if username in self.active_connections:
            self.active_connections[username] = [ws for ws in self.active_connections[username] if ws != websocket]
            if not self.active_connections[username]:
                del self.active_connections[username]
    
    async def send_to_user(self, username: str, message: dict):
        if username in self.active_connections:
            for ws in self.active_connections[username]:
                try:
                    await ws.send_json(message)
                except:
                    pass

manager = ConnectionManager()

@app.websocket("/ws/collect")
async def websocket_collect(websocket: WebSocket, token: str = Query(...)):
    username = get_user_from_token(token)
    if not username:
        await websocket.close(code=4001, reason="未登录")
        return
    
    await manager.connect(websocket, username)
    
    try:
        while True:
            data = await websocket.receive_json()
            
            # 支持多种消息类型
            msg_type = data.get("type", "text")
            
            if msg_type == "text":
                chunk = TextChunk(text=data.get("text", ""), source=data.get("source", "screen_capture"))
                results = process_text_chunk(chunk.text, source=chunk.source)
                if results:
                    ts = datetime.now().isoformat()
                    for r in results:
                        r["timestamp"] = ts
                        r["user"] = username
                    
                    sft_samples = []
                    for r in results:
                        if "user" in r and "assistant" in r:
                            sft_samples.append(to_chatml(r["user"], r["assistant"]))
                        else:
                            sft_samples.append({
                                "messages": [{"role": "user", "content": r["text"]}],
                                "source": r["source"],
                                "tags": r.get("tags", []),
                                "corpus_id": r.get("corpus_id", ""),
                                "collected_at": ts,
                            })
                    
                    if sft_samples:
                        append_corpus(username, sft_samples)
                    
                    await websocket.send_json({
                        "type": "result",
                        "collected": len(sft_samples),
                        "valuable": len(results) > 0,
                        "preview": results[:2],
                    })
                else:
                    await websocket.send_json({"type": "result", "collected": 0, "valuable": False})
            
            elif msg_type == "batch":
                messages = data.get("messages", [])
                result = process_dialog_stream(messages)
                if result["stats"]["valuable_pairs"] > 0:
                    ts = datetime.now().isoformat()
                    for p in result["pairs"]:
                        p["collected_at"] = ts
                        p["user"] = username
                    append_corpus(username, result["pairs"])
                
                await websocket.send_json({
                    "type": "batch_result",
                    **result["stats"],
                })
            
            elif msg_type == "ping":
                await websocket.send_json({"type": "pong"})
                
    except WebSocketDisconnect:
        manager.disconnect(websocket, username)
    except Exception as e:
        manager.disconnect(websocket, username)

# ============================================================
# 启动
# ============================================================

if __name__ == "__main__":
    print(f"🚀 {settings.APP_NAME} v{settings.VERSION}")
    print(f"📂 数据目录: {settings.DATA_DIR.absolute()}")
    print(f"🌐 服务地址: http://{settings.HOST}:{settings.PORT}")
    print(f"🤖 对话引擎: {settings.PORTAL_CHAT_URL}")
    print(f"💬 最大记忆轮数: {settings.MAX_MEMORY_ROUNDS}")
    
    settings.USERS_DIR.mkdir(parents=True, exist_ok=True)
    
    uvicorn.run(app, host=settings.HOST, port=settings.PORT)

# ============================================================
# MCP 服务端（供 WorkBuddy 调用）
# ============================================================
from mcp.server.fastmcp import FastMCP
from mcp.server.sse import SseServerTransport

# 创建 MCP server
mcp_server = FastMCP("corpus-agent", instructions="语料采集工具集 · 分析、判断、脱敏、格式化对话为SFT训练语料")

@mcp_server.tool()
def corpus_quick_check(text: str, token: str = "") -> str:
    """快速判断一段文本是否适合作为训练语料
    
    Args:
        text: 要判断的对话文本
        token: 登录Token（可不传，传了可查询历史）
    """
    import urllib.request
    try:
        req = urllib.request.Request(
            "http://127.0.0.1:8084/api/corpus/preview",
            data=json.dumps({"text": text, "source": "mcp"}).encode(),
            headers={"Content-Type": "application/json"}
        )
        with urllib.request.urlopen(req, timeout=10) as resp:
            return resp.read().decode()
    except Exception as e:
        return json.dumps({"error": str(e)}, ensure_ascii=False)

@mcp_server.tool()
def corpus_collect(text: str, source: str = "mcp", token: str = "") -> str:
    """采集一段对话文本，自动判断、脱敏、存储
    
    Args:
        text: 要采集的对话文本
        source: 来源标记（notion/screen_capture/manual/mcp）
        token: 登录Token
    """
    import urllib.request
    try:
        req = urllib.request.Request(
            f"http://127.0.0.1:8084/api/corpus/collect?token={token}",
            data=json.dumps({"text": text, "source": source}).encode(),
            headers={"Content-Type": "application/json"}
        )
        with urllib.request.urlopen(req, timeout=10) as resp:
            return resp.read().decode()
    except Exception as e:
        return json.dumps({"error": str(e)}, ensure_ascii=False)

@mcp_server.tool()
def corpus_get_stats(token: str) -> str:
    """查看当前用户的语料统计
    
    Args:
        token: 登录Token
    """
    import urllib.request
    try:
        req = urllib.request.Request(f"http://127.0.0.1:8084/api/corpus/stats?token={token}")
        with urllib.request.urlopen(req, timeout=10) as resp:
            return resp.read().decode()
    except Exception as e:
        return json.dumps({"error": str(e)}, ensure_ascii=False)

# 挂载 MCP SSE 到 FastAPI
mcp_sse = SseServerTransport("/api/mcp/")

@app.get("/api/mcp/sse")
async def handle_sse(request: Request):
    async with mcp_sse.connect_sse(
        request.scope,
        request.receive,
        request._send
    ) as streams:
        await mcp_server._mcp_server.run(
            streams[0], streams[1],
            mcp_server._mcp_server.create_initialization_options()
        )

@app.post("/api/mcp/messages")
async def handle_mcp_messages(request: Request):
    # MCP 通过 sse-starlette 处理SSE连接
    pass