AirPlan-V2/lib/air_runtime/modes/ctx_mode.py

"""
AirContext mode — V2 上下文管理器。
V2 改进：压缩质量校验、自适应 Token 估算、陈旧锁检测、三级降级压缩。
"""

from __future__ import annotations

import os
import re
from pathlib import Path

from air_runtime.io import atomic_json_write, safe_json_load
from air_runtime.paths import airplan_root, event_log_path
from air_runtime.events import EventLog, CONTEXT_COMPACTED
from air_runtime.utils import now_iso


class CompressionLevel:
    """三级降级"""
    RETRY = "retry"                # 1. 重试一次
    FALLBACK_MODEL = "fallback_model"  # 2. 换模型
    TRUNCATE = "truncate"          # 3. 激进截断


DEFAULT_TRUNCATION_KEEP = 10  # 保留最近 10 轮

CHARS_PER_TOKEN = {
    "chinese": 1.5,
    "english": 4.0,
    "code": 3.0,
    "markup": 5.0,
}

MUST_PRESERVE_PATTERNS = [
    r"[A-Za-z0-9_\-/]+\.(py|ts|js|cpp|h|md|json|yaml)",
    r"ADR-\d{4}",
    r"TODO|FIXME|HACK",
    r"INV-\d+",
]


def _ctx_paths(project_root: Path) -> dict[str, Path]:
    root = airplan_root(project_root) / "state" / "aircontext"
    return {"root": root, "state": root / "state.json", "lock": root / "compactor.lock"}


def estimate_tokens(text: str) -> int:
    chinese = len(re.findall(r"[一-鿿]", text))
    code = len(re.findall(r"[{}()\[\];=<>]", text))
    markup = len(re.findall(r"[#*\-`|]", text))
    english = max(0, len(text) - chinese - code - markup)
    tokens = (
        chinese / CHARS_PER_TOKEN["chinese"]
        + code / CHARS_PER_TOKEN["code"]
        + markup / CHARS_PER_TOKEN["markup"]
        + english / CHARS_PER_TOKEN["english"]
    )
    return int(tokens)


def validate_compression(original: str, summary: str) -> dict:
    missing = []
    for pattern in MUST_PRESERVE_PATTERNS:
        orig_matches = set(re.findall(pattern, original))
        summary_matches = set(re.findall(pattern, summary))
        lost = orig_matches - summary_matches
        if len(lost) > len(orig_matches) * 0.3 and len(orig_matches) > 3:
            missing.append({"pattern": pattern, "lost": list(lost)[:10]})
    return {"ok": len(missing) == 0, "missing": missing, "originalTokens": estimate_tokens(original),
            "summaryTokens": estimate_tokens(summary)}


def _compress_basic(text: str, max_tokens: int) -> str:
    """基础压缩：token 估算 + 截断"""
    estimated_tokens = len(text) // 3
    if estimated_tokens <= max_tokens:
        return text
    # 按行截断
    lines = text.split('\n')
    chars_per_line_estimate = 30
    keep_lines = int(max_tokens * chars_per_line_estimate / 80)  # 80 chars/line
    return '\n'.join(lines[-keep_lines:])


def _simplify_prompt(text: str) -> str:
    """简化 prompt：移除详细上下文，保留核心"""
    lines = text.split('\n')
    # 只保留前 3 行 + 包含 "def " / "class " / "#" 的行
    kept = lines[:3]
    kept.extend([l for l in lines[3:] if 'def ' in l or 'class ' in l or l.startswith('#')])
    return '\n'.join(kept)


def compress_with_fallback(context: str, max_tokens: int = 4000) -> dict:
    """
    三级降级压缩：
    - 尝试正常压缩
    - 失败则换模型重试
    - 再失败则激进截断
    返回: {"level": "...", "result": "...", "tokens": N}
    """
    # Level 1: 正常尝试
    try:
        result = _compress_basic(context, max_tokens)
        return {"level": CompressionLevel.RETRY, "result": result, "tokens": len(result.split())}
    except Exception:
        pass

    # Level 2: 换模型（更简单的 prompt + 更宽松的 max_tokens）
    try:
        simplified = _simplify_prompt(context)
        result = _compress_basic(simplified, int(max_tokens * 1.5))
        return {"level": CompressionLevel.FALLBACK_MODEL, "result": result, "tokens": len(result.split())}
    except Exception:
        pass

    # Level 3: 激进截断
    lines = context.split('\n')
    # 提取 ADR 引用行
    adr_lines = [l for l in lines if 'ADR-' in l or 'adr-' in l]
    # 保留最近 N 轮
    recent_lines = lines[-DEFAULT_TRUNCATION_KEEP * 5:]  # 每轮约 5 行
    truncated = '\n'.join(recent_lines + adr_lines)
    return {
        "level": CompressionLevel.TRUNCATE,
        "result": truncated,
        "tokens": len(truncated.split()),
        "warning": f"truncated to {DEFAULT_TRUNCATION_KEEP * 5} recent lines + {len(adr_lines)} ADR lines"
    }


def validate_compression_with_fallback(project_root: Path, context_path: Path) -> dict:
    """验证压缩有效性，失败时触发三级降级"""
    content = context_path.read_text()
    original_len = len(content)

    # 先用当前配置尝试
    result = compress_with_fallback(content)

    validation = {
        "original_chars": original_len,
        "result_chars": len(result["result"]),
        "level": result["level"],
        "tokens": result.get("tokens", 0),
    }

    if result["level"] == CompressionLevel.TRUNCATE:
        validation["warning"] = result.get("warning", "")

    return validation


def acquire_compactor_lock(lock_path: Path) -> bool:
    try:
        fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
        os.write(fd, str(os.getpid()).encode())
        os.close(fd)
        return True
    except FileExistsError:
        try:
            pid = int(lock_path.read_text().strip())
            os.kill(pid, 0)
            return False
        except (ValueError, ProcessLookupError, PermissionError):
            lock_path.unlink(missing_ok=True)
            return acquire_compactor_lock(lock_path)


def ctx_enter(project_root: Path) -> dict:
    paths = _ctx_paths(project_root)
    paths["root"].mkdir(parents=True, exist_ok=True)
    atomic_json_write(paths["state"], {"enabled": True, "updatedAt": now_iso(),
                                        "projectRoot": str(project_root)})
    return {"state_path": str(paths["state"])}


def main(args) -> None:
    project_root = Path(args.project).expanduser().resolve()
    sub = args.sub or "status"
    paths = _ctx_paths(project_root)

    if sub == "enter":
        result = ctx_enter(project_root)
        print(f"airplan_mode=ctx\nstate_path={result['state_path']}")
    elif sub == "estimate":
        text = "sample"  # 实际使用时从 stdin 或文件读取
        tokens = estimate_tokens(text)
        print(f"airplan_mode=ctx\ntokens={tokens}")
    elif sub == "validate":
        ctx_path = project_root / "AirPlan" / "context.md"
        if ctx_path.exists():
            validation = validate_compression_with_fallback(project_root, ctx_path)
            ok = validation["level"] != CompressionLevel.TRUNCATE
            print(f"airplan_mode=ctx\nvalidation_ok={ok}\nlevel={validation['level']}\ntokens={validation['tokens']}")
            if "warning" in validation:
                print(f"warning={validation['warning']}")
        else:
            print("airplan_mode=ctx\nvalidation_ok=true")
    elif sub == "compress":
        # 读取 context 文件
        ctx_path = project_root / "AirPlan" / "context.md"
        if not ctx_path.exists():
            print("error: context.md not found")
            return

        # 调用三级降级压缩
        content = ctx_path.read_text()
        result = compress_with_fallback(content)

        log = EventLog(event_log_path(project_root))
        log.emit(CONTEXT_COMPACTED, {
            "compressionLevel": result["level"],
            "originalChars": len(content),
            "resultChars": len(result["result"]),
            "tokens": result.get("tokens", 0),
        })

        print(f"airplan_mode=ctx")
        print(f"compression_level={result['level']}")
        print(f"original_chars={len(content)}")
        print(f"result_chars={len(result['result'])}")
        if 'warning' in result:
            print(f"warning={result['warning']}")

        # 可选：写回压缩结果
        if getattr(args, "write_back", False):
            compressed_path = project_root / "AirPlan" / "context.compressed.md"
            compressed_path.write_text(result['result'])
            print(f"written_to={compressed_path}")
    else:
        state = safe_json_load(paths["state"]) or {}
        print(f"airplan_mode=ctx\nenabled={state.get('enabled', False)}")