"""Read-side context cost estimator for prepared okstra task bundles.

The estimator answers a narrow operational question: how much file/context
surface does the current task bundle ask the lead, analysis workers, and
report-writer to absorb? It does not mutate task artifacts.
"""
from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path
from typing import Iterable

from okstra_ctl.paths import RunRef, runs_dir_of, task_manifest_file
from okstra_ctl.task_target import resolve_task_root, project_rel
from okstra_ctl.json_boundary import load_owned_object


INPUT_FILES = (
    "task-brief.md",
    "analysis-profile.md",
    "analysis-material.md",
    "reference-expectations.md",
    "clarification-response.md",
)
# Per-run hot-path instruction assets the lead loads OUTSIDE the task bundle:
# the lifecycle lead contracts (Phase 1 / 2-5 / 5.5 / 6-7) plus the host-native
# functional agent specs. These are the prompt-diet (perf plan v2 P2) targets; the
# bundle-local metrics below cannot see them, so they get their own metric.
# The contracts ship as runtime resources under ~/.okstra/prompts/lead/, not
# as agent skills.
HOT_PATH_LEAD_RESOURCES = (
    "context-loader",
    "team-contract",
    "convergence",
    "report-writer",
)
WORKER_AGENT_FILES = (
    "claude-worker.md",
    "report-writer-worker.md",
    "translator-worker.md",
)
TIMESTAMPED_ARTIFACT_RE = re.compile(r"\d{4}-\d{2}-\d{2}[_T]\d{2}-\d{2}-\d{2}")


def _file_size(path: Path) -> int:
    return path.stat().st_size if path.is_file() else 0


def _estimate_tokens(paths: Iterable[Path]) -> int:
    """Static token proxy: ~4 ASCII chars per token, non-ASCII (KR/CJK)
    ~1 token per char. Heuristic — real billable cost is the token-usage
    collector's job; this only ranks instruction surfaces against each other.
    """
    total = 0.0
    for path in paths:
        if not path.is_file():
            continue
        text = path.read_text(encoding="utf-8", errors="replace")
        ascii_chars = sum(1 for ch in text if ord(ch) < 128)
        total += ascii_chars / 4 + (len(text) - ascii_chars)
    return round(total)


def _count_files(paths: Iterable[Path]) -> tuple[int, int]:
    file_count = 0
    byte_count = 0
    for path in paths:
        if path.is_file():
            file_count += 1
            byte_count += path.stat().st_size
    return file_count, byte_count


def _all_files(root: Path) -> list[Path]:
    if not root.exists():
        return []
    return [path for path in root.rglob("*") if path.is_file()]


def _load_json(path: Path) -> dict:
    if not path.is_file():
        return {}
    try:
        data = load_owned_object(path, artifact="context cost input")
    except Exception:
        return {}
    return data if isinstance(data, dict) else {}


def _find_current_run_dir(
    task_root: Path, manifest: dict, project_root: Path
) -> Path | None:
    artifacts = manifest.get("artifacts", {})
    prompt_dir = artifacts.get("workerPromptsDirectoryPath")
    if isinstance(prompt_dir, str) and prompt_dir:
        path = project_root / prompt_dir
        if path.is_dir():
            return path.parent

    task_type = manifest.get("workflow", {}).get("currentPhase") or manifest.get("taskType")
    if not isinstance(task_type, str) or not task_type:
        return None
    candidate = RunRef.from_task_root(task_root, task_type).run_dir
    return candidate if candidate.is_dir() else None


def _latest_matching_file(directory: Path, pattern: str) -> Path | None:
    matches = sorted(directory.glob(pattern))
    if not matches:
        return None
    return max(matches, key=lambda path: path.stat().st_mtime)


def _duty_contract_files(
    run_dir: Path | None,
    project_root: Path,
    audience: str,
) -> list[Path]:
    if run_dir is None:
        return []
    manifest_path = _latest_matching_file(
        run_dir / "manifests",
        "run-manifest-*.json",
    )
    manifest = _load_json(manifest_path) if manifest_path else {}
    contract = manifest.get("agentContract")
    duty_root_value = (
        contract.get("dutyRootPath")
        if isinstance(contract, dict)
        else None
    )
    if not isinstance(duty_root_value, str) or not duty_root_value:
        return []
    duty_root = project_root / duty_root_value
    files = [duty_root / "common.md", duty_root / f"{audience}.md"]
    return [path for path in files if path.is_file()]


def _is_timestamped_legacy_artifact(path: Path) -> bool:
    return bool(TIMESTAMPED_ARTIFACT_RE.search(path.name))


def _installed_or_dev(installed: Path, dev_relative: str) -> Path:
    """Prefer the user-machine install (what production runs read); fall back
    to the repo dev tree when running from an uninstalled checkout."""
    if installed.is_file():
        return installed
    return Path(__file__).resolve().parents[2] / dev_relative


def _runtime_template(filename: str) -> Path:
    return _installed_or_dev(
        Path.home() / ".okstra" / "templates" / filename,
        f"templates/{filename}",
    )


def _header_path(prompt_path: Path | None, header: str) -> Path | None:
    if prompt_path is None or not prompt_path.is_file():
        return None
    prefix = f"**{header}:**"
    for line in prompt_path.read_text(encoding="utf-8", errors="replace").splitlines():
        stripped = line.strip()
        if stripped.startswith(prefix):
            value = stripped[len(prefix):].strip().strip("`")
            return Path(value) if value else None
    return None


def _prompt_contract_paths(
    prompt_path: Path | None,
    default_preamble: str,
) -> tuple[Path, Path]:
    selected_preamble = _header_path(prompt_path, "Worker Preamble Path")
    selected_error = _header_path(prompt_path, "Worker Error Contract Path")
    preamble = (
        selected_preamble
        if selected_preamble and selected_preamble.is_file()
        else _runtime_template(
            selected_preamble.name if selected_preamble else default_preamble
        )
    )
    error_contract = (
        selected_error
        if selected_error and selected_error.is_file()
        else _runtime_template("worker-error-contract.md")
    )
    return preamble, error_contract


def _initial_prompt(run_dir: Path | None, *, report_writer: bool) -> Path | None:
    if run_dir is None:
        return None
    candidates = sorted((run_dir / "prompts").glob("*.md"))
    for path in candidates:
        is_report_writer = "report-writer" in path.name
        if is_report_writer == report_writer and "-reverify-r" not in path.name:
            return path
    return None


def _instruction_set_metric(task_root: Path, project_root: Path) -> dict:
    instruction_set = task_root / "instruction-set"
    files = _all_files(instruction_set)
    file_count, byte_count = _count_files(files)
    packet = instruction_set / "analysis-packet.md"
    legacy_packet = instruction_set / "task-packet.md"
    return {
        "path": project_rel(instruction_set, project_root),
        "fileCount": file_count,
        "bytes": byte_count,
        "estimatedTokens": _estimate_tokens(files),
        "analysisPacketBytes": _file_size(packet),
        "legacyTaskPacketBytes": _file_size(legacy_packet),
    }


def _skill_assets_metric() -> dict:
    """Per-run hot-path instruction assets outside the task bundle: lifecycle
    skill bodies + worker agent specs. These dominate the fixed per-run
    instruction footprint and are the prompt-diet ranking input."""
    entries = []
    okstra_home = Path.home() / ".okstra"
    claude_home = Path.home() / ".claude"
    for name in HOT_PATH_LEAD_RESOURCES:
        path = _installed_or_dev(
            okstra_home / "prompts" / "lead" / f"{name}.md",
            f"prompts/lead/{name}.md",
        )
        entries.append((f"resource:{name}", path))
    for fname in WORKER_AGENT_FILES:
        path = _installed_or_dev(
            claude_home / "agents" / fname,
            f"runtime/agents/workers/{fname}",
        )
        entries.append((f"agent:{fname}", path))

    files = []
    for label, path in entries:
        if not path.is_file():
            continue
        files.append({
            "name": label,
            "path": str(path),
            "bytes": _file_size(path),
            "estimatedTokens": _estimate_tokens([path]),
        })
    files.sort(key=lambda row: row["bytes"], reverse=True)
    return {
        "fileCount": len(files),
        "bytes": sum(row["bytes"] for row in files),
        "estimatedTokens": sum(row["estimatedTokens"] for row in files),
        "files": files,
    }


def _lead_phase1_metric(
    task_root: Path, run_dir: Path | None, manifest: dict, project_root: Path
) -> dict:
    active_path = None
    artifacts = manifest.get("artifacts", {})
    active_rel = artifacts.get("activeRunContextPath")
    if isinstance(active_rel, str) and active_rel:
        active_path = project_root / active_rel

    if active_path and active_path.is_file():
        instruction_set = task_root / "instruction-set"
        packet = instruction_set / "analysis-packet.md"
        files = [
            task_manifest_file(task_root),
            active_path,
            instruction_set / "analysis-profile.md",
            packet if packet.is_file() else instruction_set / "task-brief.md",
        ]
        mode = "active-run-context"
    else:
        current_phase = manifest.get("workflow", {}).get("currentPhase") or manifest.get("taskType", "")
        run_manifest = None
        team_state = None
        if run_dir is not None:
            manifests_dir = run_dir / "manifests"
            state_dir = run_dir / "state"
            run_manifest = _latest_matching_file(
                manifests_dir, f"run-manifest-{current_phase}-*.json"
            )
            team_state = _latest_matching_file(
                state_dir, f"team-state-{current_phase}-*.json"
            )
        files = [
            task_manifest_file(task_root),
            task_root / "instruction-set" / "task-brief.md",
            task_root / "instruction-set" / "analysis-profile.md",
        ]
        if run_manifest:
            files.append(run_manifest)
        if team_state:
            files.append(team_state)
        mode = "legacy-five-file"

    duty_files = _duty_contract_files(run_dir, project_root, "lead")
    files.extend(duty_files)
    file_count, byte_count = _count_files(files)
    return {
        "mode": mode,
        "fileCount": file_count,
        "bytes": byte_count,
        "estimatedTokens": _estimate_tokens(files),
        "dutyContractBytes": sum(_file_size(path) for path in duty_files),
        "files": [project_rel(path, project_root) for path in files if path.is_file()],
    }


def _analysis_worker_metric(
    task_root: Path,
    project_root: Path,
    run_dir: Path | None,
    task_type: str,
) -> dict:
    instruction_set = task_root / "instruction-set"
    full_contract_files = [instruction_set / name for name in INPUT_FILES]
    duty_files = _duty_contract_files(
        run_dir,
        project_root,
        "analysis-worker",
    )
    full_contract_files.extend(duty_files)
    default_preamble = (
        "implementation-worker-preamble.md"
        if task_type == "implementation"
        else "worker-prompt-preamble.md"
    )
    preamble, error_contract = _prompt_contract_paths(
        _initial_prompt(run_dir, report_writer=False),
        default_preamble,
    )
    full_contract_files.extend((preamble, error_contract))
    full_file_count, full_byte_count = _count_files(full_contract_files)

    packet = instruction_set / "analysis-packet.md"
    legacy_packet = instruction_set / "task-packet.md"
    primary_packet = packet if packet.is_file() else legacy_packet
    has_packet = primary_packet.is_file()
    packet_files = (
        [primary_packet, preamble, error_contract, *duty_files]
        if has_packet
        else []
    )
    packet_file_count, packet_byte_count = _count_files(packet_files)
    mode = "analysis-packet-primary" if packet.is_file() else "full-input-contract"
    byte_count = packet_byte_count if packet.is_file() else full_byte_count
    file_count = packet_file_count if packet.is_file() else full_file_count
    current_files = packet_files if packet.is_file() else full_contract_files
    reduction_percent = 0
    if full_byte_count and has_packet:
        reduction_percent = round((1 - packet_byte_count / full_byte_count) * 100)

    return {
        "mode": mode,
        "fileCount": file_count,
        "bytesPerWorker": byte_count,
        "estimatedTokensPerWorker": _estimate_tokens(current_files),
        "dutyContractBytes": sum(_file_size(path) for path in duty_files),
        "legacyFullContractBytesPerWorker": full_byte_count,
        "legacyFullContractFileCount": full_file_count,
        "estimatedPacketModeBytesPerWorker": packet_byte_count,
        "estimatedReductionPercent": reduction_percent,
        "promptPreamblePath": str(preamble),
        "workerErrorContractPath": str(error_contract),
        "files": [project_rel(path, project_root) for path in current_files if path.is_file()],
        "legacyFullContractFiles": [
            project_rel(path, project_root)
            for path in full_contract_files
            if path.is_file()
        ],
    }


_SEQ_RESULT_RE = re.compile(r"-(\d{3,})\.md$")


def _current_seq_worker_results(run_dir: Path) -> list[Path]:
    """현재 run seq 의 분석 worker 결과만 추린다.

    같은 worker-results 디렉토리에 seq-less 레거시(`codex-worker.md`)와 이전
    seq 결과가 공존한다(관측: dev-9186 에 5월 레거시 72KB 혼입 → reportWriter
    표면 과대 측정). 실 dispatch 계약은 현재 seq 파일만 열거하므로 max seq 로
    스코핑한다."""
    candidates = [
        path for path in (run_dir / "worker-results").glob("*.md")
        if "-audit-" not in path.name and "report-writer" not in path.name
    ]
    by_seq: dict[str, list[Path]] = {}
    for path in candidates:
        match = _SEQ_RESULT_RE.search(path.name)
        if match:
            by_seq.setdefault(match.group(1), []).append(path)
    if not by_seq:
        return []
    return sorted(by_seq[max(by_seq)])


def _report_writer_metric(run_dir: Path | None, task_root: Path, project_root: Path) -> dict:
    instruction_set = task_root / "instruction-set"
    files = [instruction_set / name for name in INPUT_FILES]
    # 실 dispatch 계약(report-writer prompt 의 required reading)은 packet 이
    # 있으면 analysis-material 대신 analysis-packet 을 읽는다 — worker metric
    # 의 packet-primary 규칙과 동일.
    packet = instruction_set / "analysis-packet.md"
    if packet.is_file():
        files = [path for path in files if path.name != "analysis-material.md"]
        files.append(packet)
    files.extend([
        instruction_set / "final-report-template.md",
        instruction_set / "final-report-schema.json",
    ])
    preamble, error_contract = _prompt_contract_paths(
        _initial_prompt(run_dir, report_writer=True),
        "report-writer-prompt-preamble.md",
    )
    files.extend((preamble, error_contract))
    duty_files = _duty_contract_files(run_dir, project_root, "report-writer")
    files.extend(duty_files)
    if run_dir is not None:
        files.extend(_current_seq_worker_results(run_dir))
        convergence = _latest_matching_file(
            run_dir / "state", f"convergence-{run_dir.name}-*.json"
        )
        if convergence:
            files.append(convergence)
    file_count, byte_count = _count_files(files)
    return {
        "mode": "raw-synthesis-inputs",
        "fileCount": file_count,
        "bytes": byte_count,
        "estimatedTokens": _estimate_tokens(files),
        "dutyContractBytes": sum(_file_size(path) for path in duty_files),
        "promptPreamblePath": str(preamble),
        "workerErrorContractPath": str(error_contract),
        "files": [project_rel(path, project_root) for path in files if path.is_file()],
    }


def analyze_task_bundle(task_root: Path, project_root: Path) -> dict:
    manifest = _load_json(task_manifest_file(task_root))
    run_dir = _find_current_run_dir(task_root, manifest, project_root)
    all_task_files = _all_files(task_root)
    task_file_count, task_bytes = _count_files(all_task_files)
    current_run_file_count, current_run_bytes = _count_files(_all_files(run_dir)) if run_dir else (0, 0)
    legacy_timestamp_files = [
        path for path in runs_dir_of(task_root).rglob("*")
        if path.is_file() and _is_timestamped_legacy_artifact(path)
    ]

    return {
        "ok": True,
        "taskRoot": project_rel(task_root, project_root),
        "projectRoot": str(project_root),
        "taskKey": manifest.get("taskKey", ""),
        "taskType": manifest.get("taskType", ""),
        "currentRunPath": project_rel(run_dir, project_root) if run_dir else "",
        "totals": {
            "taskFileCount": task_file_count,
            "taskBytes": task_bytes,
            "currentRunFileCount": current_run_file_count,
            "currentRunBytes": current_run_bytes,
            "legacyTimestampFileCount": len(legacy_timestamp_files),
        },
        "instructionSet": _instruction_set_metric(task_root, project_root),
        "leadPhase1": _lead_phase1_metric(task_root, run_dir, manifest, project_root),
        "analysisWorker": _analysis_worker_metric(
            task_root,
            project_root,
            run_dir,
            str(manifest.get("taskType") or ""),
        ),
        "reportWriter": _report_writer_metric(run_dir, task_root, project_root),
        "skillAssets": _skill_assets_metric(),
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Estimate context/file-read cost for an okstra task bundle."
    )
    parser.add_argument("target", help="task root path or task-key")
    parser.add_argument("--project-root", default="", help="project root for task-key lookup")
    parser.add_argument("--cwd", default=".", help="cwd for project root resolution")
    args = parser.parse_args(argv)

    task_root, project_root = resolve_task_root(args.target, args.project_root, args.cwd)
    result = analyze_task_bundle(task_root, project_root)
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))