"""Create an experimental per-channel W8 variant of one validated typed package.

This is weight-only compression. The caller must separately verify native
decisions, Core ML placement, and full-request latency before publishing it.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from typed_coreml import package_sha256


def package_bytes(package: Path) -> int:
    return sum(path.stat().st_size for path in package.rglob("*") if path.is_file())


def quantize(source: Path, destination: Path, scope: str = "all", linear_mode: str = "linear_symmetric") -> dict:
    import coremltools as ct
    from coremltools.optimize.coreml import (
        OpLinearQuantizerConfig,
        OptimizationConfig,
        get_weights_metadata,
        linear_quantize_weights,
    )

    source = source.resolve()
    destination = destination.resolve()
    if source == destination or destination.exists():
        raise ValueError("Destination must be a distinct, nonexistent package")
    if source.suffix != ".mlpackage" or destination.suffix != ".mlpackage":
        raise ValueError("Source and destination must be .mlpackage directories")
    if not source.is_dir():
        raise FileNotFoundError(source)
    source_report = source.with_suffix(".json")
    if not source_report.is_file():
        raise FileNotFoundError(source_report)
    metadata = json.loads(source_report.read_text())
    if metadata.get("coreml_choice_agreement") is not True:
        raise ValueError(
            "The source package must have successful conversion validation"
        )
    original = ct.models.MLModel(str(source), skip_model_load=True)
    quantizer = OpLinearQuantizerConfig(
        mode=linear_mode, dtype="int8", granularity="per_channel", weight_threshold=2048,
    )
    selected_weights = None
    if scope == "embedding":
        weights = get_weights_metadata(original, weight_threshold=2048)
        selected_weights = [
            name for name, weight in weights.items()
            if len(weight.val.shape) == 2
            and any(op.op_type == "gather" for op in weight.child_ops)
            and "embeddings_tok_embeddings" in name
        ]
        if len(selected_weights) != 1:
            raise ValueError(f"expected one token embedding, found {selected_weights}")
        config = OptimizationConfig(op_name_configs={selected_weights[0]: quantizer})
    elif scope == "all":
        config = OptimizationConfig(global_config=quantizer)
    else:
        raise ValueError(f"unsupported compression scope: {scope}")
    compressed = linear_quantize_weights(original, config=config)
    destination.parent.mkdir(parents=True, exist_ok=True)
    compressed.save(str(destination))
    result = {
        "kind": metadata["kind"],
        "shape": metadata["shape"],
        "source_repo": metadata["source_repo"],
        "source_revision": metadata["source_revision"],
        "package_sha256": package_sha256(destination),
        "package_bytes": package_bytes(destination),
        "compressed_from": {
            "package": source.name,
            "sha256": package_sha256(source),
            "bytes": package_bytes(source),
        },
        "compression": {
            "type": "weight-only linear quantization",
            "mode": linear_mode,
            "dtype": "int8",
            "granularity": "per_channel",
            "weight_threshold": 2048,
            "scope": scope,
            "selected_weights": selected_weights,
        },
        "validated": False,
    }
    destination.with_suffix(".json").write_text(json.dumps(result, indent=2) + "\n")
    return result


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--scope", choices=("all", "embedding"), default="all")
    parser.add_argument("--linear-mode", choices=("linear_symmetric", "linear"), default="linear_symmetric")
    args = parser.parse_args()
    print(json.dumps(quantize(args.source, args.output, args.scope, args.linear_mode), indent=2))


if __name__ == "__main__":
    main()