#!/usr/bin/env python3
"""Collect machine-readable facts about the Piko-9b checkpoint.

Everything reported here is read from files.  Nothing is inferred, and nothing
is filled in from the model card.

Usage
-----
python scripts/audit_repository.py --model <path> --output reports/repository_audit.json
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
from pathlib import Path
from typing import Any

from safetensors import safe_open

# Patterns that must never appear in a published repository.
SECRET_PATTERNS = {
    "hf_token": re.compile(r"\bhf_[A-Za-z0-9]{34,}\b"),
    "openai_key": re.compile(r"\bsk-[A-Za-z0-9]{20,}\b"),
    "aws_access_key": re.compile(r"\bAKIA[0-9A-Z]{16}\b"),
    "private_key_block": re.compile(r"-----BEGIN (?:RSA |OPENSSH |EC )?PRIVATE KEY-----"),
    "generic_bearer": re.compile(r"\bBearer\s+[A-Za-z0-9\-._~+/]{24,}"),
}

# Absolute paths that leak the author's machine layout.
PATH_PATTERNS = {
    "wsl_mount": re.compile(r"/mnt/[a-z]/"),
    "windows_drive": re.compile(r"\b[A-Z]:\\\\?"),
    "unix_home": re.compile(r"/home/[A-Za-z0-9_.-]+/"),
    "windows_users": re.compile(r"[Cc]:[\\/]+Users[\\/]+"),
}

TEXT_SUFFIXES = {".json", ".md", ".txt", ".jinja", ".yaml", ".yml", ".py", ".cfg", ".toml"}


def sha256_file(path: Path, limit: int | None = None) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        read = 0
        while chunk := handle.read(1 << 20):
            digest.update(chunk)
            read += len(chunk)
            if limit is not None and read >= limit:
                break
    return digest.hexdigest()


def safetensors_header(path: Path) -> dict[str, Any]:
    """Read the JSON header of a safetensors file without loading tensors."""
    with path.open("rb") as handle:
        length = int.from_bytes(handle.read(8), "little")
        header = json.loads(handle.read(length))
    metadata = header.pop("__metadata__", {})
    dtypes: dict[str, int] = {}
    parameters = 0
    for entry in header.values():
        dtypes[entry["dtype"]] = dtypes.get(entry["dtype"], 0) + 1
        count = 1
        for dimension in entry["shape"]:
            count *= dimension
        parameters += count
    return {
        "file": path.name,
        "bytes": path.stat().st_size,
        "tensor_count": len(header),
        "dtypes": dtypes,
        "parameters": parameters,
        "metadata": metadata,
    }


def scan_text_files(root: Path) -> dict[str, Any]:
    secrets: list[dict[str, Any]] = []
    paths: list[dict[str, Any]] = []
    for file in sorted(root.rglob("*")):
        if not file.is_file() or file.suffix not in TEXT_SUFFIXES:
            continue
        if file.stat().st_size > 8 * 1024 * 1024:
            continue  # tokenizer.json and friends: no secrets, huge cost
        try:
            text = file.read_text(encoding="utf-8", errors="replace")
        except OSError:
            continue
        relative = str(file.relative_to(root))
        for name, pattern in SECRET_PATTERNS.items():
            if pattern.search(text):
                secrets.append({"file": relative, "pattern": name})
        for name, pattern in PATH_PATTERNS.items():
            for match in set(pattern.findall(text)):
                # Record the surrounding line so a human can judge it.
                for line in text.splitlines():
                    if match in line:
                        paths.append(
                            {"file": relative, "pattern": name, "line": line.strip()[:200]}
                        )
                        break
    return {"secrets": secrets, "absolute_paths": paths}


def audit(model_root: Path) -> dict[str, Any]:
    report: dict[str, Any] = {"model_root": str(model_root)}

    config = json.loads((model_root / "config.json").read_text(encoding="utf-8"))
    text_config = config.get("text_config", {})
    vision_config = config.get("vision_config", {})
    rope = text_config.get("rope_parameters", {})

    layer_types = text_config.get("layer_types", [])
    report["architecture"] = {
        "architectures": config.get("architectures"),
        "model_type": config.get("model_type"),
        "declared_transformers_version": config.get("transformers_version"),
        "tie_word_embeddings": config.get("tie_word_embeddings"),
        "text": {
            "hidden_size": text_config.get("hidden_size"),
            "intermediate_size": text_config.get("intermediate_size"),
            "num_hidden_layers": text_config.get("num_hidden_layers"),
            "num_attention_heads": text_config.get("num_attention_heads"),
            "num_key_value_heads": text_config.get("num_key_value_heads"),
            "head_dim": text_config.get("head_dim"),
            "vocab_size": text_config.get("vocab_size"),
            "max_position_embeddings": text_config.get("max_position_embeddings"),
            "layer_type_counts": {
                kind: layer_types.count(kind) for kind in sorted(set(layer_types))
            },
            "full_attention_interval": text_config.get("full_attention_interval"),
            "linear_attention": {
                "conv_kernel_dim": text_config.get("linear_conv_kernel_dim"),
                "key_head_dim": text_config.get("linear_key_head_dim"),
                "num_key_heads": text_config.get("linear_num_key_heads"),
                "num_value_heads": text_config.get("linear_num_value_heads"),
                "value_head_dim": text_config.get("linear_value_head_dim"),
            },
            "dtype": text_config.get("dtype"),
        },
        "vision": {
            "depth": vision_config.get("depth"),
            "hidden_size": vision_config.get("hidden_size"),
            "num_heads": vision_config.get("num_heads"),
            "patch_size": vision_config.get("patch_size"),
            "spatial_merge_size": vision_config.get("spatial_merge_size"),
            "temporal_patch_size": vision_config.get("temporal_patch_size"),
            "out_hidden_size": vision_config.get("out_hidden_size"),
            "num_position_embeddings": vision_config.get("num_position_embeddings"),
            "deepstack_visual_indexes": vision_config.get("deepstack_visual_indexes"),
        },
        "rope": rope,
        "special_token_ids": {
            "image_token_id": config.get("image_token_id"),
            "video_token_id": config.get("video_token_id"),
            "vision_start_token_id": config.get("vision_start_token_id"),
            "vision_end_token_id": config.get("vision_end_token_id"),
            "eos_token_id": text_config.get("eos_token_id"),
        },
    }

    # Non-standard keys are exactly where provenance leaks tend to hide.
    known_top_level = {
        "architectures",
        "image_token_id",
        "model_type",
        "text_config",
        "tie_word_embeddings",
        "transformers_version",
        "video_token_id",
        "vision_config",
        "vision_end_token_id",
        "vision_start_token_id",
        "dtype",
        "eos_token_id",
        "pad_token_id",
        "hidden_size",
        "use_cache",
    }
    report["nonstandard_config_keys"] = {
        key: config[key] for key in config if key not in known_top_level
    }

    shards = sorted(model_root.glob("model-*.safetensors"))
    headers = [safetensors_header(shard) for shard in shards]
    total_parameters = sum(h["parameters"] for h in headers)
    dtypes: dict[str, int] = {}
    for header in headers:
        for dtype, count in header["dtypes"].items():
            dtypes[dtype] = dtypes.get(dtype, 0) + count

    index_path = model_root / "model.safetensors.index.json"
    weight_map = (
        json.loads(index_path.read_text(encoding="utf-8"))["weight_map"]
        if index_path.is_file()
        else {}
    )
    vision_keys = [k for k in weight_map if ".visual." in k]
    language_keys = [k for k in weight_map if ".visual." not in k]

    vision_parameters = 0
    language_parameters = 0
    for shard in shards:
        with safe_open(shard, framework="pt", device="cpu") as handle:
            for key in handle.keys():
                view = handle.get_slice(key)
                count = 1
                for dimension in view.get_shape():
                    count *= int(dimension)
                if ".visual." in key:
                    vision_parameters += count
                else:
                    language_parameters += count

    report["weights"] = {
        "shard_count": len(shards),
        "total_bytes": sum(h["bytes"] for h in headers),
        "tensor_count": sum(h["tensor_count"] for h in headers),
        "total_parameters": total_parameters,
        "language_parameters": language_parameters,
        "vision_parameters": vision_parameters,
        "language_tensor_count": len(language_keys),
        "vision_tensor_count": len(vision_keys),
        "dtype_histogram": dtypes,
        "safetensors_metadata": {h["file"]: h["metadata"] for h in headers},
        "shards": headers,
    }

    generation_path = model_root / "generation_config.json"
    report["generation_config"] = (
        json.loads(generation_path.read_text(encoding="utf-8"))
        if generation_path.is_file()
        else None
    )

    tokenizer_config = json.loads(
        (model_root / "tokenizer_config.json").read_text(encoding="utf-8")
    )
    added = tokenizer_config.get("added_tokens_decoder", {})
    report["tokenizer"] = {
        "tokenizer_class": tokenizer_config.get("tokenizer_class"),
        "processor_class": tokenizer_config.get("processor_class"),
        "model_max_length": tokenizer_config.get("model_max_length"),
        "eos_token": tokenizer_config.get("eos_token"),
        "pad_token": tokenizer_config.get("pad_token"),
        "bos_token": tokenizer_config.get("bos_token"),
        "padding_side": tokenizer_config.get("padding_side"),
        "added_token_count": len(added),
        "special_tokens": tokenizer_config.get("model_specific_special_tokens"),
        "has_chat_template_file": (model_root / "chat_template.jinja").is_file(),
    }

    for name in ("preprocessor_config.json", "video_preprocessor_config.json"):
        path = model_root / name
        report[name.replace(".json", "")] = (
            json.loads(path.read_text(encoding="utf-8")) if path.is_file() else None
        )

    report["files"] = (
        sorted(
            {
                "name": str(p.relative_to(model_root)),
                "bytes": p.stat().st_size,
            }
            for p in model_root.rglob("*")
            if p.is_file() and ".cache" not in p.parts
        )
        if False
        else [
            {"name": str(p.relative_to(model_root)), "bytes": p.stat().st_size}
            for p in sorted(model_root.rglob("*"))
            if p.is_file() and ".cache" not in p.parts
        ]
    )

    report["remote_code"] = {
        "auto_map_present": "auto_map" in config,
        "python_files_in_repo": [f["name"] for f in report["files"] if f["name"].endswith(".py")],
        "trust_remote_code_required": "auto_map" in config,
    }

    report["hygiene"] = scan_text_files(model_root)

    small_hashes = {}
    for entry in report["files"]:
        path = model_root / entry["name"]
        if entry["bytes"] <= 4 * 1024 * 1024:
            small_hashes[entry["name"]] = sha256_file(path)
    report["file_sha256_small"] = small_hashes

    return report


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True, type=Path)
    parser.add_argument("--output", type=Path, default=Path("reports/repository_audit.json"))
    args = parser.parse_args()

    report = audit(args.model)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")

    weights = report["weights"]
    print(f"Wrote {args.output}")
    print(f"  parameters       : {weights['total_parameters']:,}")
    print(f"    language       : {weights['language_parameters']:,}")
    print(f"    vision         : {weights['vision_parameters']:,}")
    print(f"  shards           : {weights['shard_count']}")
    print(f"  dtypes           : {weights['dtype_histogram']}")
    print(f"  secrets found    : {len(report['hygiene']['secrets'])}")
    print(f"  absolute paths   : {len(report['hygiene']['absolute_paths'])}")
    for entry in report["hygiene"]["absolute_paths"]:
        print(f"      {entry['file']}: {entry['line'][:120]}")


if __name__ == "__main__":
    main()