cm.py

#!/usr/bin/env python3
"""cm - Code Memory CLI. 코드 심볼 단위 경험 기록/조회."""

import argparse
import json
import os
import re
import sqlite3
import sys
from datetime import datetime, timezone
from pathlib import Path

DB_PATH = Path(os.environ.get("CODE_MEMORY_DB", Path(__file__).parent / "code_memory.db"))

SECRET_RE = re.compile(
    r"(sk-[a-zA-Z0-9]{20,}|ghp_[a-zA-Z0-9]{36}|Bearer\s+\S{10,}|AKIA[A-Z0-9]{16})"
)

VALID_STATUSES = {"active", "resolved", "inactive"}

SCHEMA = """
CREATE TABLE IF NOT EXISTS symbols (
    id INTEGER PRIMARY KEY,
    project TEXT NOT NULL,
    name TEXT NOT NULL,
    file TEXT NOT NULL DEFAULT '',
    fail_count INTEGER DEFAULT 0,
    success_count INTEGER DEFAULT 0,
    last_fail TEXT,
    last_success TEXT,
    risk_score REAL DEFAULT 0,
    UNIQUE(project, file, name)
);

CREATE TABLE IF NOT EXISTS experiences (
    id INTEGER PRIMARY KEY,
    symbol_id INTEGER REFERENCES symbols(id),
    type TEXT NOT NULL,
    lesson TEXT NOT NULL,
    detail TEXT,
    session_id TEXT,
    status TEXT DEFAULT 'active',
    created_at TEXT DEFAULT (datetime('now')),
    CHECK(status IN ('active', 'resolved', 'inactive'))
);

CREATE INDEX IF NOT EXISTS idx_exp_symbol ON experiences(symbol_id);
CREATE INDEX IF NOT EXISTS idx_exp_status ON experiences(status);
CREATE INDEX IF NOT EXISTS idx_sym_project ON symbols(project);
CREATE INDEX IF NOT EXISTS idx_sym_lookup ON symbols(project, file, name);
"""


def get_db() -> sqlite3.Connection:
    db = sqlite3.connect(str(DB_PATH))
    db.row_factory = sqlite3.Row
    db.execute("PRAGMA journal_mode=WAL")
    db.execute("PRAGMA busy_timeout=5000")
    version = db.execute("PRAGMA user_version").fetchone()[0]
    if version in (0, 1):
        migrate_v1_to_v2(db)
    return db


def apply_schema(db: sqlite3.Connection) -> None:
    for statement in SCHEMA.split(";"):
        statement = statement.strip()
        if statement:
            db.execute(statement)


def apply_schema_tables(db: sqlite3.Connection) -> None:
    statements = [statement.strip() for statement in SCHEMA.split(";") if statement.strip()]
    for statement in statements[:2]:
        db.execute(statement)


def migrate_v1_to_v2(db: sqlite3.Connection) -> None:
    try:
        db.execute("BEGIN IMMEDIATE")
        if db.execute("PRAGMA user_version").fetchone()[0] == 2:
            db.execute("COMMIT")
            return
        db.execute("DROP TABLE IF EXISTS experiences_old")
        db.execute("DROP TABLE IF EXISTS symbols_old")
        apply_schema_tables(db)

        if needs_schema_rebuild(db):
            rebuild_schema(db)

        apply_schema(db)
        for row in db.execute("SELECT id FROM symbols").fetchall():
            recalculate_risk(db, row["id"])
        db.execute("PRAGMA user_version = 2")
        db.execute("COMMIT")
    except Exception as exc:
        try:
            db.execute("ROLLBACK")
        except sqlite3.Error:
            pass
        print(f"migration failed: {exc}", file=sys.stderr)
        raise SystemExit(1)


def needs_schema_rebuild(db: sqlite3.Connection) -> bool:
    tables = {
        row["name"]
        for row in db.execute("SELECT name FROM sqlite_master WHERE type = 'table'").fetchall()
    }
    if "symbols" not in tables or "experiences" not in tables:
        return False

    sym_cols = table_columns(db, "symbols")
    exp_cols = table_columns(db, "experiences")
    return (
        "risk_score" not in sym_cols
        or sym_cols.get("file", {}).get("notnull") != 1
        or "status" not in exp_cols
        or not has_unique_index(db, "symbols", ["project", "file", "name"])
    )


def rebuild_schema(db: sqlite3.Connection) -> None:
    db.execute("ALTER TABLE symbols RENAME TO symbols_old")
    db.execute("ALTER TABLE experiences RENAME TO experiences_old")
    apply_schema(db)

    old_sym_cols = table_columns(db, "symbols_old")
    file_expr = "COALESCE(file, '')" if "file" in old_sym_cols else "''"
    risk_expr = "COALESCE(risk_score, 0)" if "risk_score" in old_sym_cols else "0"
    db.execute(
        f"""
        INSERT INTO symbols (id, project, name, file, fail_count, success_count, last_fail, last_success, risk_score)
        SELECT id, project, name, {file_expr}, COALESCE(fail_count, 0), COALESCE(success_count, 0),
               last_fail, last_success, {risk_expr}
        FROM symbols_old
        """
    )

    old_exp_cols = table_columns(db, "experiences_old")
    status_expr = (
        "CASE WHEN status IN ('active', 'resolved', 'inactive') THEN status ELSE 'active' END"
        if "status" in old_exp_cols
        else "'active'"
    )
    created_at_expr = "created_at" if "created_at" in old_exp_cols else "datetime('now')"
    db.execute(
        f"""
        INSERT INTO experiences (id, symbol_id, type, lesson, detail, session_id, status, created_at)
        SELECT id, symbol_id, type, lesson, detail, session_id, {status_expr}, {created_at_expr}
        FROM experiences_old
        """
    )
    db.execute("DROP TABLE experiences_old")
    db.execute("DROP TABLE symbols_old")


def table_columns(db: sqlite3.Connection, table: str) -> dict:
    return {row["name"]: dict(row) for row in db.execute(f"PRAGMA table_info({table})")}


def has_unique_index(db: sqlite3.Connection, table: str, columns: list) -> bool:
    for idx in db.execute(f"PRAGMA index_list({table})").fetchall():
        if not idx["unique"]:
            continue
        info = db.execute(f"PRAGMA index_info({idx['name']})").fetchall()
        if [row["name"] for row in info] == columns:
            return True
    return False


def now_iso() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def normalize_file(file: str = None) -> str:
    if file is None or file == "":
        return ""
    normalized = str(file).replace("\\", "/")
    normalized = os.path.normpath(normalized)
    normalized = normalized.replace("\\", "/")
    if os.path.isabs(normalized):
        normalized = os.path.relpath(normalized, os.getcwd())
        normalized = normalized.replace("\\", "/")
    if normalized.startswith("./"):
        normalized = normalized[2:]
    return normalized


def mask_secret(value: str = None) -> str:
    if value is None:
        return None
    return SECRET_RE.sub("***REDACTED***", value)


def risk_level(score: float) -> str:
    if score >= 15:
        return "높음"
    if score >= 10:
        return "주의"
    if score >= 5:
        return "보통"
    return "낮음"


def parse_iso(value: str):
    if not value:
        return None
    try:
        dt = datetime.strptime(value, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc)
    except ValueError:
        try:
            dt = datetime.fromisoformat(value.replace("Z", "+00:00"))
        except ValueError:
            return None
    if dt.tzinfo is None:
        return dt.replace(tzinfo=timezone.utc)
    return dt


def recency_weight(last_fail: str) -> float:
    dt = parse_iso(last_fail)
    if not dt:
        return 0.0
    days = (datetime.now(timezone.utc) - dt).days
    if days <= 7:
        return 1.0
    if days <= 30:
        return 0.5
    return 0.2


def consecutive_fails(db: sqlite3.Connection, sym_id: int, last_success: str) -> int:
    if last_success:
        row = db.execute(
            """
            SELECT COUNT(*) AS n FROM experiences
            WHERE symbol_id = ? AND type = 'fail' AND status = 'active' AND created_at > ?
            """,
            (sym_id, last_success),
        ).fetchone()
    else:
        row = db.execute(
            """
            SELECT COUNT(*) AS n FROM experiences
            WHERE symbol_id = ? AND type = 'fail' AND status = 'active'
            """,
            (sym_id,),
        ).fetchone()
    return row["n"] or 0


def recalculate_risk(db: sqlite3.Connection, sym_id: int) -> float:
    sym = db.execute("SELECT * FROM symbols WHERE id = ?", (sym_id,)).fetchone()
    if not sym:
        return 0.0
    stats = db.execute(
        """
        SELECT
            SUM(CASE WHEN type = 'fail' THEN 1 ELSE 0 END) AS fail_count,
            SUM(CASE WHEN type = 'success' THEN 1 ELSE 0 END) AS success_count,
            MAX(CASE WHEN type = 'fail' THEN created_at END) AS last_fail,
            MAX(CASE WHEN type = 'success' THEN created_at END) AS last_success
        FROM experiences
        WHERE symbol_id = ? AND status = 'active'
        """,
        (sym_id,),
    ).fetchone()
    fail_count = stats["fail_count"] or 0
    success_count = stats["success_count"] or 0
    last_fail = stats["last_fail"]
    last_success = stats["last_success"]
    score = (
        fail_count * 1.0
        + recency_weight(last_fail) * 1.5
        + consecutive_fails(db, sym_id, last_success) * 1.5
        - success_count * 0.3
    )
    score = max(0.0, round(score, 2))
    db.execute(
        """
        UPDATE symbols
        SET fail_count = ?, success_count = ?, last_fail = ?, last_success = ?, risk_score = ?
        WHERE id = ?
        """,
        (fail_count, success_count, last_fail, last_success, score, sym_id),
    )
    return score


def ensure_symbol(db: sqlite3.Connection, project: str, name: str, file: str = None) -> int:
    file = normalize_file(file)
    db.execute(
        "INSERT OR IGNORE INTO symbols (project, name, file) VALUES (?, ?, ?)",
        (project, name, file),
    )
    row = db.execute(
        "SELECT id FROM symbols WHERE project = ? AND file = ? AND name = ?",
        (project, file, name),
    ).fetchone()
    return row["id"]


def as_symbol_dict(row: sqlite3.Row) -> dict:
    data = dict(row)
    data["file"] = data.get("file") or ""
    data["risk_score"] = float(data.get("risk_score") or 0)
    data["risk_level"] = risk_level(data["risk_score"])
    return data


def as_experience_dict(row: sqlite3.Row) -> dict:
    return dict(row)


def print_json(data) -> None:
    print(json.dumps(data, indent=2, ensure_ascii=False))


def query_symbols(db: sqlite3.Connection, args):
    params = [args.project]
    where = ["project = ?"]
    if args.symbol:
        where.append("name LIKE ?")
        params.append(f"%{args.symbol}%")
    if getattr(args, "file", None) is not None:
        where.append("file = ?")
        params.append(normalize_file(args.file))
    order = "risk_score DESC, fail_count DESC, name ASC" if not args.symbol else "name ASC, file ASC"
    limit = " LIMIT ?" if not args.symbol else ""
    if not args.symbol:
        where.append("fail_count > 0")
        params.append(args.top)
    sql = f"SELECT * FROM symbols WHERE {' AND '.join(where)} ORDER BY {order}{limit}"
    return db.execute(sql, params).fetchall()


def get_experiences(db: sqlite3.Connection, symbol_id: int, include_all: bool = False):
    status_clause = "" if include_all else " AND status = 'active'"
    return db.execute(
        f"""
        SELECT * FROM experiences
        WHERE symbol_id = ?{status_clause}
        ORDER BY created_at DESC, id DESC
        LIMIT 10
        """,
        (symbol_id,),
    ).fetchall()


def injection_text(symbol: dict, experiences: list) -> str:
    if symbol["risk_score"] < 10:
        return ""
    lessons = [e["lesson"] for e in experiences if e["status"] == "active"]
    if not lessons:
        return ""
    lines = [
        f"[PAST EXPERIENCE — {symbol['name']} "
        f"(risk: {symbol['risk_score']:.2f}, {symbol['risk_level']})]"
    ]
    lines.extend(f"- {lesson}" for lesson in lessons)
    lines.append("Apply these lessons. Do not repeat these mistakes.")
    return "\n".join(lines)


def cmd_query(args):
    db = get_db()
    rows = query_symbols(db, args)

    if args.inject:
        blocks = []
        for row in rows:
            sym = as_symbol_dict(row)
            exps = [as_experience_dict(e) for e in get_experiences(db, sym["id"], False)]
            block = injection_text(sym, exps)
            if block:
                blocks.append(block)
        output = "\n\n".join(blocks)
        if args.json:
            print_json({"project": args.project, "injection": output})
        elif output:
            print(output)
        else:
            print("No injection needed.")
        return

    if args.json:
        result = []
        for row in rows:
            sym = as_symbol_dict(row)
            if args.symbol:
                sym["experiences"] = [
                    as_experience_dict(e) for e in get_experiences(db, sym["id"], args.all)
                ]
            result.append(sym)
        print_json({"project": args.project, "symbols": result})
        return

    if args.symbol:
        if not rows:
            print(f"[{args.project}] '{args.symbol}' 경험 없음.")
            return
        for row in rows:
            sym = as_symbol_dict(row)
            print(f"\n{'=' * 50}")
            print(f"  {sym['name']}  ({sym['file'] or '파일 미상'})")
            print(
                f"  실패: {sym['fail_count']}  성공: {sym['success_count']}  "
                f"위험도: {sym['risk_score']:.2f} ({sym['risk_level']})"
            )
            if sym["last_fail"]:
                print(f"  최근 실패: {sym['last_fail']}")
            if sym["risk_score"] >= 10:
                print(f"  주의 - 위험도 {sym['risk_score']:.2f}. 과거 교훈 주입 권장.")
            exps = get_experiences(db, sym["id"], args.all)
            if exps:
                print(f"\n  경험 ({len(exps)}건):")
                for e in exps:
                    tag = "x" if e["type"] == "fail" else "ok"
                    status = "" if e["status"] == "active" else f" [{e['status']}]"
                    print(f"    {tag}{status} [{e['created_at'][:10]}] #{e['id']} {e['lesson']}")
                    if e["detail"]:
                        print(f"      - {e['detail']}")
    else:
        if not rows:
            print(f"[{args.project}] 실패 기록 없음.")
            return
        print(f"[{args.project}] 위험 상위 {len(rows)}건:\n")
        for r in rows:
            sym = as_symbol_dict(r)
            file_part = f" ({sym['file']})" if sym["file"] else ""
            print(
                f"  위험도 {sym['risk_score']:>5.2f} ({sym['risk_level']})  "
                f"실패 {sym['fail_count']:>3}회  {sym['name']}{file_part}"
            )


def record_experience(args, kind: str) -> dict:
    db = get_db()
    sym_id = ensure_symbol(db, args.project, args.symbol, args.file)
    created_at = now_iso()
    lesson = mask_secret(args.lesson)
    detail = mask_secret(args.detail)
    db.execute(
        """
        INSERT INTO experiences (symbol_id, type, lesson, detail, session_id, status, created_at)
        VALUES (?, ?, ?, ?, ?, 'active', ?)
        """,
        (sym_id, kind, lesson, detail, args.session, created_at),
    )
    if kind == "fail":
        db.execute(
            "UPDATE symbols SET fail_count = fail_count + 1, last_fail = ? WHERE id = ?",
            (created_at, sym_id),
        )
    else:
        db.execute(
            "UPDATE symbols SET success_count = success_count + 1, last_success = ? WHERE id = ?",
            (created_at, sym_id),
        )
    score = recalculate_risk(db, sym_id)
    db.commit()
    sym = as_symbol_dict(db.execute("SELECT * FROM symbols WHERE id = ?", (sym_id,)).fetchone())
    sym["risk_score"] = score
    return {"symbol": sym, "lesson": lesson, "detail": detail}


def cmd_fail(args):
    data = record_experience(args, "fail")
    if args.json:
        print_json({"ok": True, "type": "fail", **data})
        return
    sym = data["symbol"]
    print(f"x 기록: {sym['name']} - \"{data['lesson']}\" (누적 실패 {sym['fail_count']}회)")


def cmd_ok(args):
    data = record_experience(args, "success")
    if args.json:
        print_json({"ok": True, "type": "success", **data})
        return
    print(f"ok 기록: {data['symbol']['name']} - \"{data['lesson']}\"")


def update_experience_status(args, status: str):
    db = get_db()
    row = db.execute("SELECT * FROM experiences WHERE id = ?", (args.experience_id,)).fetchone()
    if not row:
        if args.json:
            print_json({"ok": False, "error": f"experience {args.experience_id} not found"})
        else:
            print(f"experience {args.experience_id} 없음.", file=sys.stderr)
        return 1
    db.execute("UPDATE experiences SET status = ? WHERE id = ?", (status, args.experience_id))
    score = recalculate_risk(db, row["symbol_id"])
    db.commit()
    updated = as_experience_dict(
        db.execute("SELECT * FROM experiences WHERE id = ?", (args.experience_id,)).fetchone()
    )
    symbol = as_symbol_dict(
        db.execute("SELECT * FROM symbols WHERE id = ?", (row["symbol_id"],)).fetchone()
    )
    if args.json:
        print_json({"ok": True, "experience": updated, "symbol": symbol})
    else:
        print(f"experience #{args.experience_id} -> {status} (risk {score:.2f})")
    return 0


def cmd_resolve(args):
    return update_experience_status(args, "resolved")


def cmd_deactivate(args):
    return update_experience_status(args, "inactive")


def cmd_stats(args):
    db = get_db()
    if args.symbol:
        rows = query_symbols(db, args)
        if not rows:
            if args.json:
                print_json({"project": args.project, "symbols": []})
            else:
                print(f"'{args.symbol}' 없음.")
            return
        data = [as_symbol_dict(row) for row in rows]
        if args.json:
            print_json({"project": args.project, "symbols": data})
        else:
            for sym in data:
                print_json(sym)
        return

    rows = db.execute(
        """
        SELECT project, COUNT(*) AS syms, SUM(fail_count) AS fails, SUM(success_count) AS oks,
               ROUND(AVG(risk_score), 2) AS avg_risk_score,
               ROUND(MAX(risk_score), 2) AS max_risk_score
        FROM symbols
        GROUP BY project
        ORDER BY project
        """
    ).fetchall()
    data = [dict(r) for r in rows]
    if args.json:
        print_json({"projects": data})
        return
    if not rows:
        print("데이터 없음.")
        return
    print(f"{'프로젝트':<20} {'심볼':>6} {'실패':>6} {'성공':>6} {'평균위험도':>10} {'최대위험도':>10}")
    print("-" * 70)
    for r in data:
        print(
            f"{r['project']:<20} {r['syms']:>6} {r['fails'] or 0:>6} {r['oks'] or 0:>6} "
            f"{r['avg_risk_score'] or 0:>10.2f} {r['max_risk_score'] or 0:>10.2f}"
        )


def add_common_flags(parser):
    parser.add_argument("--json", action="store_true", default=argparse.SUPPRESS, help="JSON 출력")


def main():
    parser = argparse.ArgumentParser(prog="cm", description="Code Memory - 코드 심볼 경험 기록")
    parser.add_argument("--project", "-p", default="default", help="프로젝트명")
    parser.add_argument("--json", action="store_true", help="JSON 출력")
    sub = parser.add_subparsers(dest="cmd")

    q = sub.add_parser("query", help="경험 조회")
    q.add_argument("symbol", nargs="?", help="심볼명 (부분 매칭)")
    q.add_argument("--top", type=int, default=10, help="위험 상위 N건")
    q.add_argument("--file", help="파일 경로")
    q.add_argument("--all", action="store_true", help="resolved/inactive 경험도 표시")
    q.add_argument("--inject", action="store_true", help="프롬프트 주입 블록 출력")
    add_common_flags(q)

    f = sub.add_parser("fail", help="실패 기록")
    f.add_argument("symbol", help="심볼명")
    f.add_argument("lesson", help="한 줄 교훈")
    f.add_argument("--detail", "-d", help="추가 맥락")
    f.add_argument("--file", help="파일 경로")
    f.add_argument("--session", help="세션 ID")
    add_common_flags(f)

    o = sub.add_parser("ok", help="성공 기록")
    o.add_argument("symbol", help="심볼명")
    o.add_argument("lesson", nargs="?", default="정상 통과", help="한 줄")
    o.add_argument("--detail", "-d", help="추가 맥락")
    o.add_argument("--file", help="파일 경로")
    o.add_argument("--session", help="세션 ID")
    add_common_flags(o)

    s = sub.add_parser("stats", help="통계")
    s.add_argument("symbol", nargs="?", help="심볼명")
    s.add_argument("--file", help="파일 경로")
    add_common_flags(s)

    r = sub.add_parser("resolve", help="경험을 resolved 상태로 변경")
    r.add_argument("experience_id", type=int)
    add_common_flags(r)

    d = sub.add_parser("deactivate", help="경험을 inactive 상태로 변경")
    d.add_argument("experience_id", type=int)
    add_common_flags(d)

    args = parser.parse_args()
    if hasattr(args, "file") and args.file is not None:
        args.file = normalize_file(args.file)
    if not args.cmd:
        parser.print_help()
        return 0

    result = {
        "query": cmd_query,
        "fail": cmd_fail,
        "ok": cmd_ok,
        "stats": cmd_stats,
        "resolve": cmd_resolve,
        "deactivate": cmd_deactivate,
    }[args.cmd](args)
    return result or 0


if __name__ == "__main__":
    raise SystemExit(main())