#!/usr/bin/env python3
"""Measure a skill's SKILL.md / description / references footprint.

Token counts are a CJK-aware heuristic, not a model tokenizer.
Budgets are defined in references/token-structure.md; this script only measures.
"""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path

SOFT_SKILL_LINES = 200
HARD_SKILL_LINES = 500
SOFT_SKILL_TOKENS = 2500
HARD_SKILL_TOKENS = 5000
SOFT_DESC_TOKENS = 120
HARD_DESC_CHARS = 1024

_CJK_RE = re.compile(r"[\u4e00-\u9fff]")
_LATIN_RE = re.compile(r"[A-Za-z0-9_]+")
_OTHER_RE = re.compile(r"[^\s\w\u4e00-\u9fff]")
_DESC_BLOCK_RE = re.compile(
    r"^description:\s*(?:>-|>\||>|-)?\s*\n((?:[ \t].+\n?)+)",
    re.MULTILINE,
)
_DESC_INLINE_RE = re.compile(r"^description:\s*(.+)$", re.MULTILINE)
_HEADING_RE = re.compile(r"^## .+$", re.MULTILINE)
_LINK_RE = re.compile(r"\[[^\]]*\]\(([^)]+)\)")
_VAGUE_REFS_RE = re.compile(
    r"见\s*`?references/?`?|详见\s*`?references/?`?|see\s+references/?",
    re.IGNORECASE,
)


def approx_tokens(text: str) -> int:
    cjk = len(_CJK_RE.findall(text))
    latin = len(_LATIN_RE.findall(text))
    other = len(_OTHER_RE.findall(text))
    return int(cjk + latin * 1.3 + other * 0.5)


def parse_frontmatter(text: str) -> tuple[str, str, str]:
    if not text.startswith("---"):
        return "", "", text
    end = text.find("\n---", 3)
    if end == -1:
        return "", "", text
    frontmatter = text[4:end]
    body = text[end + 4 :]
    name = ""
    name_match = re.search(r"^name:\s*(.+)$", frontmatter, re.MULTILINE)
    if name_match:
        name = name_match.group(1).strip().strip("\"'")
    desc = ""
    block = _DESC_BLOCK_RE.search(frontmatter)
    if block:
        desc = " ".join(
            line.strip() for line in block.group(1).splitlines() if line.strip()
        )
    else:
        inline = _DESC_INLINE_RE.search(frontmatter)
        if inline:
            desc = inline.group(1).strip().strip("\"'")
    return name, desc, body


def resolve_skill_dir(spec: str) -> Path:
    path = Path(spec).expanduser()
    if (path / "SKILL.md").is_file():
        return path.resolve()
    if path.is_file() and path.name == "SKILL.md":
        return path.parent.resolve()

    home = Path.home() / ".pouch"
    for candidate in (home / "skills" / spec, home / ".drafts" / spec):
        if (candidate / "SKILL.md").is_file():
            return candidate.resolve()
    raise FileNotFoundError(
        f"找不到 skill: {spec}（需要目录内有 SKILL.md，或 ~/.pouch/skills/<name>）"
    )


def iter_reference_files(skill_dir: Path) -> list[Path]:
    files: list[Path] = []
    refs_dir = skill_dir / "references"
    if refs_dir.is_dir():
        files.extend(sorted(p for p in refs_dir.glob("*.md") if p.is_file()))
    loose = skill_dir / "reference.md"
    if loose.is_file():
        files.append(loose)
    return files


def relative_links(text: str) -> list[str]:
    found: list[str] = []
    for raw in _LINK_RE.findall(text):
        target = raw.strip().split("#", 1)[0].split("?", 1)[0]
        if not target or "://" in target or target.startswith(("mailto:", "/")):
            continue
        found.append(target)
    return found


def section_sizes(body: str) -> list[tuple[str, int, int]]:
    matches = list(_HEADING_RE.finditer(body))
    rows: list[tuple[str, int, int]] = []
    if not matches:
        title = body.strip().splitlines()[0] if body.strip() else "(body)"
        rows.append((title[:60], body.count("\n") + 1, approx_tokens(body)))
        return rows
    preamble = body[: matches[0].start()]
    if preamble.strip():
        rows.append(
            (
                "(preamble)",
                preamble.count("\n") + 1,
                approx_tokens(preamble),
            )
        )
    for index, match in enumerate(matches):
        start = match.start()
        end = matches[index + 1].start() if index + 1 < len(matches) else len(body)
        chunk = body[start:end]
        rows.append(
            (
                match.group(0)[:60],
                chunk.count("\n") + 1,
                approx_tokens(chunk),
            )
        )
    return rows


def bulk_load_sections(body: str, ref_names: set[str]) -> list[str]:
    flagged: list[str] = []
    matches = list(_HEADING_RE.finditer(body))
    spans: list[tuple[str, str]] = []
    if matches:
        for index, match in enumerate(matches):
            start = match.start()
            end = matches[index + 1].start() if index + 1 < len(matches) else len(body)
            spans.append((match.group(0), body[start:end]))
    else:
        spans.append(("(body)", body))
    for heading, chunk in spans:
        hits = {name for name in ref_names if name in chunk}
        if len(hits) >= 3:
            flagged.append(f"{heading} → {', '.join(sorted(hits))}")
    return flagged


def audit(skill_dir: Path) -> tuple[str, int]:
    skill_md = skill_dir / "SKILL.md"
    text = skill_md.read_text(encoding="utf-8")
    name, desc, body = parse_frontmatter(text)
    lines = text.count("\n") + 1
    skill_tokens = approx_tokens(text)
    desc_tokens = approx_tokens(desc)
    desc_chars = len(desc)
    sections = section_sizes(body)
    ref_files = iter_reference_files(skill_dir)
    links = relative_links(text)
    ref_names = {path.name for path in ref_files}
    linked_refs = {
        Path(target).name
        for target in links
        if Path(target).name in ref_names or target.startswith("references/")
    }
    orphans = sorted(ref_names - linked_refs)
    vague = bool(_VAGUE_REFS_RE.search(text))
    bulk = bulk_load_sections(body, ref_names)

    warnings: list[str] = []
    notes: list[str] = []
    if lines > HARD_SKILL_LINES or skill_tokens > HARD_SKILL_TOKENS:
        warnings.append(
            f"HARD SKILL.md {lines} 行 / {skill_tokens} token "
            f"(硬顶 {HARD_SKILL_LINES} 行 / {HARD_SKILL_TOKENS} token)"
        )
    elif lines > SOFT_SKILL_LINES or skill_tokens > SOFT_SKILL_TOKENS:
        warnings.append(
            f"SOFT SKILL.md {lines} 行 / {skill_tokens} token "
            f"(目标 {SOFT_SKILL_LINES} 行 / {SOFT_SKILL_TOKENS} token)"
        )
    if desc_chars > HARD_DESC_CHARS:
        warnings.append(
            f"HARD description {desc_chars} 字符 (硬顶 {HARD_DESC_CHARS})"
        )
    elif desc_tokens > SOFT_DESC_TOKENS:
        warnings.append(
            f"SOFT description ~{desc_tokens} token (目标 ≤{SOFT_DESC_TOKENS})"
        )
    if vague:
        warnings.append("SKILL.md 含空泛「见 references/」，应改成「若 X 则读 Y.md」")
    for item in bulk:
        warnings.append(f"无条件批量加载风险: {item}")
    for orphan in orphans:
        notes.append(f"reference 未被 SKILL.md 链接: {orphan}")

    if any(item.startswith("HARD ") for item in warnings):
        status = "over-hard-budget"
    elif warnings:
        status = "over-soft-budget"
    else:
        status = "within-budget"

    out: list[str] = [
        f"skill: {name or skill_dir.name}",
        f"path: {skill_dir}",
        f"status: {status}",
        f"SKILL.md: {lines} lines, ~{skill_tokens} tokens",
        f"description: {desc_chars} chars, ~{desc_tokens} tokens",
        "sections:",
    ]
    for heading, sec_lines, sec_tokens in sections:
        out.append(f"  {sec_tokens:5d} tok  {sec_lines:4d} lines  {heading}")
    out.append("references:")
    if not ref_files:
        out.append("  (none)")
    for path in ref_files:
        ref_text = path.read_text(encoding="utf-8")
        rel = path.relative_to(skill_dir)
        out.append(
            f"  {approx_tokens(ref_text):5d} tok  "
            f"{ref_text.count(chr(10)) + 1:4d} lines  {rel}"
        )
    out.append("warnings:")
    if not warnings:
        out.append("  (none)")
    else:
        out.extend(f"  - {item}" for item in warnings)
    out.append("notes:")
    if not notes:
        out.append("  (none)")
    else:
        out.extend(f"  - {item}" for item in notes)
    return "\n".join(out) + "\n", 0


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Audit a pouch skill's token/structure footprint."
    )
    parser.add_argument("skill", help="skill 名，或含 SKILL.md 的目录")
    args = parser.parse_args(argv)
    try:
        skill_dir = resolve_skill_dir(args.skill)
    except FileNotFoundError as exc:
        print(exc, file=sys.stderr)
        return 2
    report, code = audit(skill_dir)
    sys.stdout.write(report)
    return code


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