#!/usr/bin/env python3
"""
register_token.py — 将用户的 hc_session token 注册到 AI Token 表，
                     供 AI 直接调用后端 API。

用户只需提供一次，后端存储后 AI 后续就能用这个 token 访问。

用法：
    python3 register_token.py
    python3 register_token.py <token值>   # 静默模式（用于自动化）
"""

import hashlib
import json
import os
import secrets
import sys
import webbrowser
from datetime import datetime, timedelta
from pathlib import Path
from urllib.parse import urlencode

# ---------- 项目路径 ----------
SCRIPT_DIR = Path(__file__).resolve().parent
BACKEND_DIR = SCRIPT_DIR.parent
sys.path.insert(0, str(BACKEND_DIR))

from config import Config
from database import connection, now_iso


def _sha256(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


def _normalize_user_from_token(token: str) -> dict | None:
    """用 token 查找用户，不存在则返回 None"""
    session_hash = _sha256(token)
    current = now_iso()
    with connection() as conn:
        row = conn.execute(
            """SELECT s.id AS session_id,u.* FROM auth_sessions s
               JOIN auth_users u ON u.id=s.user_id
               WHERE s.session_hash=? AND s.revoked_at IS NULL AND s.expires_at>?""",
            (session_hash, current),
        ).fetchone()
        if not row:
            return None
        conn.execute(
            "UPDATE auth_sessions SET last_seen_at=? WHERE id=?",
            (current, row["session_id"]),
        )
        return {
            "id": str(row["ding_user_id"]),
            "name": row["name"],
            "ding_user_id": row["ding_user_id"],
        }


def ensure_ai_tokens_table() -> None:
    """建 AI Token 表（如果不存在）"""
    with connection() as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS ai_access_tokens (
                id             INTEGER PRIMARY KEY AUTOINCREMENT,
                token_hash     TEXT NOT NULL UNIQUE,
                user_id        TEXT NOT NULL,
                user_name      TEXT NOT NULL,
                description    TEXT DEFAULT '',
                created_at     TEXT NOT NULL,
                expires_at     TEXT NOT NULL,
                revoked_at     TEXT,
                last_used_at   TEXT,
                use_count      INTEGER NOT NULL DEFAULT 0
            )
        """)
        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_ai_token_hash ON ai_access_tokens(token_hash)"
        )


def register_ai_token(token: str, user_info: dict, description: str = "") -> dict:
    """注册 / 更新 AI 访问 token"""
    token_hash = _sha256(token)
    now = now_iso()
    expires = (datetime.now().astimezone() + timedelta(days=365)).isoformat(
        timespec="seconds"
    )
    with connection() as conn:
        existing = conn.execute(
            "SELECT id FROM ai_access_tokens WHERE token_hash=?",
            (token_hash,),
        ).fetchone()
        if existing:
            conn.execute(
                """UPDATE ai_access_tokens
                   SET last_used_at=?, use_count=use_count+1
                   WHERE token_hash=?""",
                (now, token_hash),
            )
            action = "updated"
        else:
            conn.execute(
                """INSERT INTO ai_access_tokens
                   (token_hash,user_id,user_name,description,created_at,expires_at,last_used_at,use_count)
                   VALUES (?,?,?,?,?,?,?,1)""",
                (
                    token_hash, user_info["ding_user_id"],
                    user_info["name"], description, now, expires, now,
                ),
            )
            action = "registered"
    return {"action": action, "user": user_info["name"], "expires_at": expires}


def revoke_ai_token(token: str) -> bool:
    token_hash = _sha256(token)
    with connection() as conn:
        conn.execute(
            "UPDATE ai_access_tokens SET revoked_at=? WHERE token_hash=? AND revoked_at IS NULL",
            (now_iso(), token_hash),
        )
        affected = conn.total_changes
    return affected > 0


def list_ai_tokens() -> list[dict]:
    with connection() as conn:
        rows = conn.execute(
            """SELECT user_name,description,created_at,expires_at,last_used_at,use_count,revoked_at
               FROM ai_access_tokens WHERE revoked_at IS NULL
               ORDER BY last_used_at DESC"""
        ).fetchall()
    return [dict(r) for r in rows]


# ---------- 主入口：交互式注册流程 ----------
def main() -> None:
    ensure_ai_tokens_table()

    backend = os.environ.get("HC_BACKEND_URL", f"http://localhost:{Config.PORT}").rstrip("/")

    print()
    print("═" * 55)
    print("  🔐  Holy Crab · AI Token 注册")
    print("═" * 55)
    print()
    print("  本工具用于将你的 hc_session token 注册给 AI 访问权限。")
    print()
    print("  Step 1：请从浏览器获取你的 hc_session token")
    print()
    print("  操作步骤：")
    print("    1. 在浏览器中打开 Holy Crab 前端并完成钉钉登录")
    print("    2. 打开浏览器 DevTools (F12) → Application → Cookies")
    print("    3. 找到你的后端地址，点击展开 Cookie 列表")
    print("    4. 复制 'hc_session' 的 Value（很长的一串字符）")
    print()

    # 检查是否传入参数
    if len(sys.argv) > 1:
        token = sys.argv[1].strip()
        if token in ("--help", "-h"):
            print("用法: python3 register_token.py [token]")
            sys.exit(0)
        print(f"  检测到命令行传入 token: {token[:8]}...{token[-4:]}")
    else:
        print("  直接回车打开浏览器说明网页（推荐）：")
        print()
        print("  或者手动输入 token（粘贴后回车）：")
        token = input("  > ").strip()

    if not token:
        # 打开说明网页
        help_url = (
            f"{backend}/api/v1/auth/dingtalk/login"
            f"?redirect={urlencode({'r': '/__help__'})}"
        )
        print()
        print(f"  正在打开说明网页：{help_url}")
        webbrowser.open(help_url)
        print()
        print("  请按上述步骤获取 token 后再次运行：")
        print(f"    python3 {sys.argv[0]} <your_token>")
        sys.exit(0)

    print()
    print("  🔍 正在验证 token...")
    user = _normalize_user_from_token(token)
    if not user:
        print()
        print("  ❌ token 无效或已过期")
        print("  请确认：")
        print("    1. token 是否正确（注意不要有前后空格）")
        print("    2. 是否已通过钉钉登录")
        print("    3. token 是否已过期（默认14天）")
        sys.exit(1)

    print(f"  ✅ 验证成功！用户：{user['name']}")
    print()
    result = register_ai_token(
        token, user,
        description=f"AI助手访问 (注册IP/来源: register_script)"
    )
    print(f"  ✅ AI Token 注册成功")
    print(f"     用户：{result['user']}")
    print(f"     有效期至：{result['expires_at']}")
    print()
    print("  📋 AI 后续使用方式：")
    print()
    print("  方式 A — 设置环境变量（推荐，每次会话有效）：")
    print(f"    export HC_SESSION='{token}'")
    print()
    print("  方式 B — 写入文件（持久化）：")
    print(f"    echo '{token}' > ~/.holy_crab_cookie")
    print()
    print("  方式 C — 写入配置（自动读取）：")
    print(f"    echo 'HC_SESSION={token}' >> ~/.holy_crab_env")
    print("    # scripts 会自动从 ~/.holy_crab_env 读取环境变量")
    print()
    print("═" * 55)
    print()


if __name__ == "__main__":
    main()