# ==========================================
# 🚀 VolcConsume — 火山免费额度自动消耗任务
# ==========================================
# 与 volc_limit_sync.py 合并方案的一部分（见 docs/merge-auto-consume.md）：
#   - 每日北京时间 23:00 由 volc_limit_sync 守护线程触发（可手动 --force）
#   - 洗掉 CONSUME_MONITOR_DAYS 天内过期的资源包额度（FIFO 下即队列头部），
#     防止 30 天过期蒸发导致"启动资金"缩水（免费额度规则：消耗 X → 次日返还 X）
#   - 直连 Ark API（Bearer token，绕过 LiteLLM 限额，避免消耗任务自杀）
#   - 消耗量手动 INCR 进 daily_usage（与 LiteLLM 回调同 key）：
#     ① 算准采集上限缺口，② 空窗期限额能正确拦截真实用户（防扣费）
#   - 三条停止条件：daily_usage ≥ 2M / 余额快照 − 已烧 ≤ 0 / 已烧 ≥ plan
#
# 手动运行：
#   python volc_consume.py --dry-run    # 只算计划不消耗
#   python volc_consume.py --force      # 忽略 23:00 时间限制立即执行
import os
import sys
import time
import math
import random
import uuid
import re
import traceback
from datetime import datetime, timezone

import requests

# 复用 volc_limit_sync 的拉包/归一化/Redis/日志/用量 key
import volc_limit_sync as vls

# ==========================================
# 🛠️ Configuration
# ==========================================
# Ark endpoint 单源：URL 直接拼自 HUOSHAN_BASE_URL（与 config.yaml 同一 env）
HUOSHAN_BASE_URL = os.getenv(
    "HUOSHAN_BASE_URL", "https://ark.cn-beijing.volces.com/api/v3"
)
ARK_API_KEY = os.getenv("HUOSHAN_API_KEY", "")

# 采集上限：火山免费额度规则，每模型每天最多返还 2M（2026-08 确认所有模型统一）
CONSUME_DAILY_CAP = int(os.getenv("CONSUME_DAILY_CAP", "2000000"))
# 计划低于该阈值不消耗（避免无谓请求）
CONSUME_MIN_PLAN = int(os.getenv("CONSUME_MIN_PLAN", "50000"))
# 洗包窗口：ExpiryTime - now <= N 天的包视为"快过期"
CONSUME_MONITOR_DAYS = int(os.getenv("CONSUME_MONITOR_DAYS", "9"))
# 计划超额系数（旧脚本逻辑：多烧 10% 防积压）
PLAN_OVERAGE = 1.1

# 请求伪装参数（移植自 auto_consume.py）
MIN_INPUT_CHARS = 30000
MAX_INPUT_CHARS = 55000
MIN_SLEEP_BASE = 15
MAX_SLEEP_BASE = 40
MAX_RUNTIME_SECONDS = 21600  # 单模型任务最长 6h
MAX_CONSEC_FAILS = 5

# 任务提示词池（移植自 auto_consume config.py）
TASK_PROMPTS = [
    {"sys": "你是资深运维专家。", "user": "请分析以下日志数据，找出潜在的风险点。"},
    {"sys": "你是数据结构化助手。", "user": "请忽略乱码，提取以下文本中的关键信息。"},
    {"sys": "你是安全审计员。", "user": "扫描以下文本，查找是否存在敏感操作记录。"},
    {
        "sys": "你是系统可靠性工程师。",
        "user": "请判断以下日志片段中是否出现服务降级或异常重启的迹象。",
    },
    {"sys": "你是日志分析工具。", "user": "汇总以下日志中的错误级别事件，并给出时间分布。"},
    {"sys": "你是性能优化顾问。", "user": "根据以下日志评估系统负载特征，指出瓶颈迹象。"},
]

# 日志源候选：宿主挂载（/host-var-log，docker-compose 只读挂载 /var/log）优先，
# 容器内 /var/log 兜底；均不存在时由合成语料兜底
LOG_SOURCES = [
    "/host-var-log/syslog",
    "/host-var-log/messages",
    "/host-var-log/kern.log",
    "/host-var-log/auth.log",
    "/host-var-log/journal.log",
    "/var/log/syslog",
    "/var/log/messages",
    "/var/log/kern.log",
    "/var/log/auth.log",
]

# 单次运行内的语料游标（来源 -> 下一窗口起点），避免连续请求取到重复片段
_CURSOR = {}

# 敏感信息脱敏规则：(正则, 替换值)
_SENSITIVE_PATTERNS = [
    (re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"), "x.x.x.x"),
    (re.compile(r"([0-9a-fA-F]{2}:){5}[0-9a-fA-F]{2}"), "xx:xx:xx:xx:xx:xx"),
    (re.compile(r"[\w.+-]+@[\w-]+\.[\w.]+"), "user@example.com"),
    (re.compile(r"(/home/|/Users/)[^\s\"':]+"), r"\1user"),
]


# ==========================================
# 📄 语料生成（移植自 auto_consume.py，原样保留）
# ==========================================
def _sanitize_line(line):
    """过滤日志行中的敏感信息（IP、MAC、邮箱、用户目录等）"""
    for pattern, repl in _SENSITIVE_PATTERNS:
        line = pattern.sub(repl, line)
    return line


def _synthetic_lines(count=30000):
    """
    生成时间线连续、格式多样的合成日志（语料不足时的兜底）。
    覆盖最近 24 小时（count 行均匀分布），多主机多风格，模拟真实系统日志流。
    """
    hosts = [
        "web-01",
        "web-02",
        "web-03",
        "api-01",
        "api-02",
        "db-01",
        "cache-01",
        "gateway-01",
        "worker-01",
        "k8s-node-01",
        "mq-01",
        "log-01",
    ]
    templates = [
        # 内核 / 网络
        "kernel: [TCP] Possible SYN flooding on port {port}. Sending cookies.",
        "kernel: EXT4-fs (sda1): mounted filesystem with ordered data mode",
        "kernel: nf_conntrack: table full, dropping packet.",
        "kernel: br0: port 2(veth{hash}) entered blocking state",
        "NetworkManager[{pid}]: <info> device (eth0): state change: activated -> connected",
        # 系统 / systemd / 安全
        "systemd[{pid}]: Started Daily apt download and clean activities.",
        "systemd[{pid}]: Stopped User Manager for UID 1000.",
        "systemd[{pid}]: Reloading nginx.service.",
        'cron[{pid}]: (root) CMD (find /var/log -name "*.gz" -mtime +7 -delete)',
        "sudo: pam_unix(sudo:session): session opened for user deploy by (uid=0)",
        "sshd: Accepted publickey for deploy from x.x.x.x port {port} ssh2",
        "sshd[{pid}]: Failed password for invalid user admin from {ip} port {port} ssh2",
        # 应用
        "nginx: worker process, pid={pid}, GET /api/v2/status HTTP/1.1 200 0.00{ms}",
        'nginx: {ip} - - ["GET /v1/users HTTP/1.1"] 200 512',
        "python3[{pid}]: heartbeat ok, workers={n}, queue_depth={q}",
        "gunicorn[{pid}]: POST /api/chat 200 {ms}ms",
        "celery[{pid}]: Task mail.send[abc{hash}] succeeded in 12.3s",
        "node[{pid}]: memory usage 62%, heap {heap}MB, gc paused {gc}ms",
        "app[{pid}]: slow query detected: SELECT * FROM metrics WHERE ts > now() - interval '1h'",
        # 数据库 / 中间件
        "postgres[{pid}]: LOG: checkpoint complete: wrote {n} buffers ({p}%)",
        "postgres[{pid}]: LOG: autovacuum: processing database postgres",
        "redis-server[{pid}]: Saving DB to disk complete",
        "redis-server[{pid}]: 100 changes in 60 seconds. Saving...",
        "mysql[{pid}]: [Note] InnoDB: Buffer pool(s) load completed",
        "rabbitmq[{pid}]: connection from {ip}:{port} opened",
        # 日志轮转 / 采集
        "logrotate: rotating /var/log/nginx/access.log, log->access.log.1",
        "rsyslogd: action 'action 18' resumed",
        "fluent-bit[{pid}]: [info] tail input plugin started",
    ]
    # 24 小时覆盖，行间隔 ≈ 86400/count 秒
    step = 86400.0 / count
    start_ts = datetime.now().timestamp() - 86400.0
    lines = []
    for i in range(count):
        ts = datetime.fromtimestamp(start_ts + i * step).strftime("%b %d %H:%M:%S")
        host = random.choice(hosts)
        tpl = random.choice(templates)
        msg = tpl.format(
            port=random.randint(1024, 65535),
            pid=random.randint(500, 9999),
            n=random.randint(10, 500),
            q=random.randint(0, 200),
            ms=random.randint(1, 99),
            hash="%04x" % random.randint(0, 0xFFFF),
            a=random.randint(0, 255),
            b=random.randint(0, 255),
            heap=random.randint(200, 900),
            gc=random.randint(1, 60),
            p=random.randint(1, 100),
            ip=f"10.0.{random.randint(0, 255)}.{random.randint(1, 254)}",
        )
        lines.append(f"{ts} {host} {msg}")
    return lines


def load_base_logs():
    """
    加载日志语料（按来源分组，保留行序，去除相邻重复、过滤敏感信息）。
    :return: (来源列表, {来源: 日志行列表})
    """
    source_groups = {}
    order = []

    for log_path in LOG_SOURCES:
        if os.path.exists(log_path) and os.access(log_path, os.R_OK):
            try:
                with open(log_path, "r", errors="ignore") as f:
                    lines = [l.strip() for l in f.readlines()[-8000:] if l.strip()]
                if lines:
                    key = os.path.basename(log_path)
                    source_groups[key] = [_sanitize_line(l) for l in lines]
                    order.append(key)
            except Exception as e:
                vls.log(f"⚠️ 读取 {log_path} 失败: {e}", level="WARNING")
                continue

    # 语料不足 1000 行时，用时间线连续的合成日志补充
    total_lines = sum(len(v) for v in source_groups.values())
    if total_lines < 1000:
        vls.log("📄 本地日志不足，使用合成日志补齐语料")
        source_groups["synthetic"] = _synthetic_lines(30000)
        order.append("synthetic")

    # 每次运行重置语料游标（新语料从头取窗口，跨请求时间轴推进）
    _CURSOR.clear()

    # 去除相邻重复行（日志轮转/重复打印的常见现象）
    for key in source_groups:
        deduped = []
        for line in source_groups[key]:
            if not deduped or line != deduped[-1]:
                deduped.append(line)
        source_groups[key] = deduped

    return order, source_groups


def generate_dynamic_payload(source_groups, order):
    """
    生成随机会话头 + 时间线连贯的日志片段。
    从单个来源中随机起点取连续窗口（保持时间线连续性），混入少量轮转标记。
    """
    target_chars = random.randint(MIN_INPUT_CHARS, MAX_INPUT_CHARS)

    available = [k for k in order if len(source_groups[k]) >= 50]
    if not available:
        available = [k for k in order]
    source_key = random.choice(available)
    source_lines = source_groups[source_key]

    # 游标取窗口：单次运行内连续前进，避免重复片段，且时间轴随请求推进
    #（语料按时间排序；耗尽后环绕，跨请求重复可接受，跨天语料重建）
    total = len(source_lines)
    window_size = min(total, 5000)
    start_idx = _CURSOR.get(source_key, 0)
    if start_idx + window_size > total:
        start_idx = 0
    _CURSOR[source_key] = (start_idx + window_size) % total
    sampled = source_lines[start_idx : start_idx + window_size]

    content_lines = []
    total_chars = 0
    for line in sampled:
        content_lines.append(line)
        total_chars += len(line) + 1
        if total_chars >= target_chars:
            break

    markers = [
        "-- Reboot --",
        "-- 日志轮转，新文件开始 --",
        "-- 会话超时，连接已断开 --",
        "# 服务健康检查快照",
        "",
    ]
    for _ in range(random.randint(1, 3)):
        if len(content_lines) > 10:
            pos = random.randint(0, len(content_lines) - 1)
            content_lines.insert(pos, random.choice(markers))

    content = "\n".join(content_lines)

    session_id = uuid.uuid4()
    ts = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
    header = f"[LOG SESSION {session_id} | 来源: {source_key} | 采集时间: {ts} UTC]"

    return f"{header}\n{content}\n[END OF SESSION {session_id}]"


# ==========================================
# 📊 消耗计划（旧 calculate_plan 核心 + v2 clamp）
# ==========================================
def _pkg_models(pkg):
    """解析资源包归属的监控模型（与 process_packages_breakdown 同口径）。"""
    raw_code = pkg.get("ConfigurationCode", "")
    raw_name = pkg.get("ConfigurationName", "Unknown")
    target_raw = raw_code if raw_code else raw_name
    norm = vls.normalize_volc_name(target_raw)
    targets = vls.MONITORED_MODEL_ALIASES.get(norm, [norm])
    if isinstance(targets, str):
        targets = [targets]
    return targets


def calculate_plan_for_model(packages, model):
    """
    计算某模型的当日消耗计划与余额快照。
    :return: (plan, details, balance)  balance = 该模型全部 Effective token 包剩余
    逻辑（保留旧脚本）：
      - 窗口内包（ExpiryTime - now ≤ CONSUME_MONITOR_DAYS）：
        days_left ≤ 1 → 冲刺全量；否则 remaining/ceil(days)×1.1，多包相加
      - 窗口外包不消耗（留给真实用户，FIFO 下被自然消耗）
    """
    total_plan = 0
    balance = 0
    details = []
    now = datetime.now(timezone.utc)

    for pkg in packages:
        if model not in _pkg_models(pkg):
            continue
        if pkg.get("Status") != "Effective":
            continue
        unit = pkg.get("Unit", "token")
        if str(unit).lower() != "token":
            continue
        try:
            remaining = int(float(pkg.get("AvailableAmount", 0)))
        except Exception:
            continue
        if remaining < 1000:
            continue
        balance += remaining

        expiry_str = pkg.get("ExpiryTime", "")
        if not expiry_str:
            continue
        try:
            if expiry_str.endswith("Z"):
                expiry_str = expiry_str[:-1] + "+00:00"
            expiry_date = datetime.fromisoformat(expiry_str)
        except Exception:
            continue

        days_left = (expiry_date - now).total_seconds() / 86400.0
        if days_left <= 0 or days_left > CONSUME_MONITOR_DAYS:
            continue

        if days_left <= 1.0:
            plan = remaining
            strategy = "🚨 冲刺(最后24h)"
        else:
            div = math.ceil(days_left)
            plan = int((remaining / div) * PLAN_OVERAGE)
            if plan > remaining:
                plan = remaining
            strategy = f"📉 分{div}天(加量10%)"

        if plan > 0:
            total_plan += plan
            details.append(
                f"   - [{pkg.get('ConfigurationName')}] 剩:{remaining:,} | 计划:{plan:,} | {strategy}"
            )

    return total_plan, details, balance


# ==========================================
# 🔌 Ark 直连请求
# ==========================================
def _send_ark(payload, timeout=120):
    """直连 Ark Chat Completions（Bearer token，不经过 LiteLLM）。
    单源配置：endpoint 直接取 HUOSHAN_BASE_URL（与 config.yaml 同一 env），防 region 漂移。"""
    url = f"{HUOSHAN_BASE_URL.rstrip('/')}/chat/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {ARK_API_KEY}",
    }
    return requests.post(url, headers=headers, json=payload, timeout=timeout)


def _build_payload(model, source_groups, order):
    """构造消耗请求：日志语料 + 任务提示词。
    模型参数与旧 MONITOR_TARGETS 实测配置一致：
    - doubao-seed 系列 → max_completion_tokens + reasoning_effort minimal（火山方舟规范）
    - 其他（glm/deepseek）→ max_tokens
    """
    dynamic_input = generate_dynamic_payload(source_groups, order)
    tpl = random.choice(TASK_PROMPTS)
    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": tpl["sys"]},
            {
                "role": "user",
                "content": f"{tpl['user']}\n\n[DATA START]\n{dynamic_input}\n[DATA END]",
            },
        ],
        "stream": False,
    }
    if model.startswith("doubao"):
        payload["reasoning_effort"] = "minimal"
        payload["max_completion_tokens"] = 1000
    else:
        payload["max_tokens"] = 1000
    return payload


# ==========================================
# 📈 用量记账（与 LiteLLM 回调同 key）
# ==========================================
def _usage_key(model):
    return vls.get_usage_key(model, HUOSHAN_BASE_URL)


def _current_daily_usage(model):
    if not vls.redis_client:
        return 0
    val = vls.redis_client.get(_usage_key(model))
    return float(val or 0)


def _record_consumption(model, tokens):
    if not vls.redis_client:
        return
    try:
        key = _usage_key(model)
        vls.redis_client.incrby(key, int(tokens))
        vls.redis_client.expire(key, 172800)
    except Exception as e:
        vls.log(f"❌ 消耗记账失败: {e}", level="ERROR")


# ==========================================
# 🎯 单模型消耗执行
# ==========================================
def run_consume_model(model, source_groups, order, packages=None, dry_run=False):
    """
    对单个监控模型执行"洗过期包"消耗。
    目标 = Σ窗口内包计划（均分×1.1/冲刺），clamp：
      - 采集上限：plan ≤ CONSUME_DAILY_CAP − 当日已消耗
      - 防扣费：plan ≤ 余额快照 − BUFFER
    停止条件（每轮检查）：daily_usage ≥ 2M / 余额快照−已烧 ≤ 0 / 已烧 ≥ plan
    :param packages: 复用 run_consume_all 的一次拉包结果；None 时自行拉取
    """
    if packages is None:
        try:
            packages = vls.fetch_all_packages()
        except Exception as e:
            vls.log(f"❌ [{model}] 拉取资源包失败: {e}", level="ERROR")
            return

    plan, details, balance = calculate_plan_for_model(packages, model)
    for d in details:
        vls.log(d)

    if plan <= 0:
        vls.log(f"💤 [{model}] 无快过期包，跳过")
        return

    usage_now = _current_daily_usage(model)
    # clamp 1: 采集上限（超出部分不返还，白烧）
    plan = min(plan, CONSUME_DAILY_CAP - usage_now)
    # clamp 2: 余额快照 − BUFFER（脚本自身不烧穿余额，防扣费）
    plan = min(plan, max(0, balance - vls.BUFFER))
    plan = int(max(0, plan))

    if plan < CONSUME_MIN_PLAN:
        vls.log(f"💤 [{model}] 计划 {plan:,} < 阈值 {CONSUME_MIN_PLAN:,}，跳过")
        return

    vls.log(
        f"▶️ [{model}] 目标消耗 {plan:,} | 余额 {balance:,} | 当日已用 {usage_now:.0f}"
    )
    if dry_run:
        vls.log(f"🔍 [dry-run] 仅计算，不发送请求")
        return

    _run_requests(model, plan, balance, source_groups, order)


def _run_requests(model, plan, balance_snapshot, source_groups, order):
    """
    循环发送消耗请求。速率：目标 1 小时内完成（23:00 触发，00:00 结算）。
    """
    burned = 0
    fail_count = 0
    deadline = time.time() + MAX_RUNTIME_SECONDS

    # 速率三档（对齐旧 auto_consume run_stealthy_consume）：
    #   任务轻松（单次间隔 > MAX_SLEEP_BASE）→ 正常模式 15-40s 伪装，不赶时间
    #   任务极重（间隔 < 2s）→ 极速模式 1-3s（本场景 plan≤2M 实际不会触发）
    #   其余 → 动态 ±5s
    # 目标：23:00 触发、00:00 结算，1 小时内完成（plan=2M ≈ 47 次 × 27s ≈ 21 分钟）
    est_per_req = (MIN_INPUT_CHARS + MAX_INPUT_CHARS) / 2 + 1000
    est_requests = max(1, plan / est_per_req)
    interval = int(3600 / est_requests)
    if interval > MAX_SLEEP_BASE:
        cur_min, cur_max = MIN_SLEEP_BASE, MAX_SLEEP_BASE
        mode_str = "正常模式 🐢"
    elif interval < 2:
        cur_min, cur_max = 1, 3
        mode_str = "极速模式 🐇"
    else:
        cur_min, cur_max = max(1, interval - 5), interval + 5
        mode_str = f"动态调整 🚗 ({cur_min}~{cur_max}s)"
    vls.log(
        f"⚡ 速率控制: {mode_str} | 预计 {int(est_requests)} 次请求 (单次 ≈{int(est_per_req):,} token)"
    )

    while burned < plan:
        if time.time() >= deadline:
            vls.log(f"⏰ [{model}] 已达运行时长上限，停止本次任务", level="WARNING")
            break
        if fail_count > MAX_CONSEC_FAILS:
            vls.log(f"❌ [{model}] 失败次数过多，跳过剩余任务", level="ERROR")
            break
        # 停止条件 1: 采集上限
        if _current_daily_usage(model) >= CONSUME_DAILY_CAP:
            vls.log(f"🛑 [{model}] 当日采集已达上限 {CONSUME_DAILY_CAP:,}，停止", level="WARNING")
            break
        # 停止条件 2: 余额快照耗尽（防扣费）
        if (balance_snapshot - burned) <= 0:
            vls.log(f"🛑 [{model}] 余额快照已耗尽，停止（防扣费）", level="WARNING")
            break
        # 停止条件 3: 计划完成
        if burned >= plan:
            break

        payload = _build_payload(model, source_groups, order)
        try:
            resp = _send_ark(payload)
            if resp.status_code == 200:
                fail_count = 0
                total = resp.json().get("usage", {}).get("total_tokens", 0)
                if total > 0:
                    _record_consumption(model, total)
                burned += total
                percent = min(100, burned / plan * 100)
                vls.log(
                    f"✨ [{model}] 进度 {percent:.1f}% | 本次 {total:,} | 累计 {burned:,}/{plan:,}"
                )
                if burned < plan:
                    time.sleep(random.randint(cur_min, cur_max))
            else:
                if resp.status_code == 401:
                    vls.log(f"❌ [{model}] 认证失败 (401)，API Key 可能已失效，中止", level="ERROR")
                    break
                fail_count += 1
                backoff = min(60, 2**fail_count) + random.uniform(0, 1)
                vls.log(
                    f"⏳ [{model}] API 错误 {resp.status_code}，{backoff:.1f}s 后重试",
                    level="ERROR",
                )
                time.sleep(backoff)
        except Exception as e:
            fail_count += 1
            vls.log(f"❌ [{model}] 请求异常: {e}", level="ERROR")
            time.sleep(10)

    vls.log(f"✅ [{model}] 消耗结束。实际消耗: {burned:,}/{plan:,}")


# ==========================================
# 🚀 总入口
# ==========================================
def run_consume_all(dry_run=False):
    """遍历所有监控模型执行消耗。模型列表复用 volc_limit_sync 的 config.yaml 发现。"""
    if not ARK_API_KEY:
        vls.log("❌ 未配置 HUOSHAN_API_KEY，无法消耗", level="ERROR")
        return

    if not vls.MONITORED_MODEL_PAIRS:
        vls.discover_monitored_models()
    models = sorted({m for m, _h in vls.MONITORED_MODEL_PAIRS})
    if not models:
        vls.log("❌ 未发现监控模型（config.yaml 无 volcengine 条目）", level="ERROR")
        return

    vls.log(f"📚 加载语料库...")
    order, source_groups = load_base_logs()
    vls.log(
        f"📄 语料来源: {', '.join(f'{k}({len(v)}行)' for k, v in source_groups.items())}"
    )

    # 只拉一次资源包，全部模型复用（每模型独立 fetch 会重复分页 6 次）
    try:
        packages = vls.fetch_all_packages()
    except Exception as e:
        vls.log(f"❌ 拉取资源包失败: {e}", level="ERROR")
        packages = []

    for model in models:
        try:
            run_consume_model(model, source_groups, order, packages=packages, dry_run=dry_run)
        except Exception as e:
            vls.log(f"❌ [{model}] 消耗任务异常: {e}\n{traceback.format_exc()}", level="ERROR")

    vls.log("🎉 消耗任务完成")


def main():
    dry_run = "--dry-run" in sys.argv
    force = "--force" in sys.argv

    if not force and not dry_run:
        beijing_hour = (time.gmtime(time.time()).tm_hour + 8) % 24
        if beijing_hour != 23:
            vls.log(
                f"⏰ 当前北京时间 {beijing_hour} 时，非 23:00 触发窗口。加 --force 强制执行",
                level="WARNING",
            )
            return

    run_consume_all(dry_run=dry_run)


if __name__ == "__main__":
    main()