# ==========================================
# 🚀 VolcConsume — 火山免费额度自动消耗任务
# ==========================================
# 与 volc_limit_sync.py 合并方案的一部分（见 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_HOST = os.getenv("ARK_HOST", "ark.cn-beijing.volces.com")
# 与 config.yaml 的 os.environ/HUOSHAN_BASE_URL 一致；用于生成与 LiteLLM 路径相同的 usage key
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": "根据以下日志评估系统负载特征，指出瓶颈迹象。"},
]

# 日志源候选（容器内大概率不存在，由合成兜底）
LOG_SOURCES = [
    "/var/log/syslog",
    "/var/log/messages",
    "/var/log/kern.log",
    "/var/log/auth.log",
]

# 敏感信息脱敏规则：(正则, 替换值)
_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=5000):
    """
    生成时间线连续、格式多样的合成日志（语料不足时的兜底）。
    从约 2 小时前开始，每行间隔约 1.5 秒，模拟真实系统日志流。
    """
    hosts = ["web-01", "web-02", "api-01", "db-01", "cache-01", "gateway-01"]
    templates = [
        "kernel: [TCP] Possible SYN flooding on port {port}. Sending cookies.",
        "sshd: Accepted publickey for deploy from x.x.x.x port {port} ssh2",
        "nginx: worker process, pid={pid}, GET /api/v2/status HTTP/1.1 200 0.00{ms}",
        "systemd: Started Daily apt download and clean activities.",
        "python3[{pid}]: heartbeat ok, workers={n}, queue_depth={q}",
        "postgres[{pid}]: LOG: checkpoint complete: wrote {n} buffers",
        "redis-server[{pid}]: Saving DB to disk complete",
        "kernel: EXT4-fs (sda1): mounted filesystem with ordered data mode",
    ]
    start_ts = datetime.now().timestamp() - count * 1.5
    lines = []
    for i in range(count):
        ts = datetime.fromtimestamp(start_ts + i * 1.5).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(1000, 9999),
            n=random.randint(10, 500),
            q=random.randint(0, 200),
            ms=random.randint(1, 9),
        )
        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(5000)
        order.append("synthetic")

    # 去除相邻重复行（日志轮转/重复打印的常见现象）
    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]

    window_size = min(len(source_lines), 5000)
    start_idx = random.randint(0, max(0, len(source_lines) - window_size))
    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）。"""
    url = f"https://{ARK_HOST}/api/v3/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):
    """构造消耗请求：日志语料 + 任务提示词，max_tokens 限制输出。"""
    dynamic_input = generate_dynamic_payload(source_groups, order)
    tpl = random.choice(TASK_PROMPTS)
    return {
        "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,
        "max_tokens": 1000,
    }


# ==========================================
# 📈 用量记账（与 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, dry_run=False):
    """
    对单个监控模型执行"洗过期包"消耗。
    目标 = Σ窗口内包计划（均分×1.1/冲刺），clamp：
      - 采集上限：plan ≤ CONSUME_DAILY_CAP − 当日已消耗
      - 防扣费：plan ≤ 余额快照 − BUFFER
    停止条件（每轮检查）：daily_usage ≥ 2M / 余额快照−已烧 ≤ 0 / 已烧 ≥ plan
    """
    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

    # 速率估算：单次 ~4.3 万 token（输入 3-5.5 万字符 + 输出 ~1000）
    est_per_req = (MIN_INPUT_CHARS + MAX_INPUT_CHARS) / 2 + 1000
    est_requests = max(1, plan / est_per_req)
    interval = int(min(60, max(2, 3600 / est_requests)))
    cur_min = max(1, interval - 5)
    cur_max = interval + 5

    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())}"
    )

    for model in models:
        try:
            run_consume_model(model, source_groups, order, 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()