"""Command-line access to the file API and HTTP service."""

import argparse
import json
import os
from pathlib import Path

from .files import prune_manifest


def main(argv=None):
    parser = argparse.ArgumentParser(prog="tensor-prune", description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    prune = commands.add_parser("prune", help="prune a memory-mapped manifest")
    prune.add_argument("--input-root", required=True, type=Path)
    prune.add_argument("--manifest", required=True)
    prune.add_argument("--output", required=True, type=Path)
    prune.add_argument("--scratch-dir", type=Path)
    prune.add_argument("--threads", type=int, default=1)
    prune.add_argument("--seed", type=int, default=0)

    serve = commands.add_parser("serve", help="run the bounded HTTP job service")
    serve.add_argument("--input-root", required=True, type=Path)
    serve.add_argument("--job-root", required=True, type=Path)
    serve.add_argument("--host", default="127.0.0.1")
    serve.add_argument("--port", type=int, default=8080)
    serve.add_argument("--workers", type=int, default=1, help="simultaneous jobs in this process")
    serve.add_argument("--queue-size", type=int, default=2)
    serve.add_argument("--threads-per-job", type=int, default=1)
    serve.add_argument("--max-weights", type=int, default=100_000_000)

    args = parser.parse_args(argv)
    if args.command == "prune":
        result = prune_manifest(args.input_root, args.manifest, args.output,
                                scratch_dir=args.scratch_dir, threads=args.threads, seed=args.seed)
        print(json.dumps(result))
        return

    try:
        import uvicorn
        from .service import create_app
    except ImportError:
        parser.error("install the HTTP extras with pip install '.[http]'")
    token = os.environ.get("TENSOR_PRUNING_API_TOKEN")
    if args.host not in {"127.0.0.1", "::1", "localhost"} and not token:
        parser.error("set TENSOR_PRUNING_API_TOKEN before binding to a non-loopback address")
    app = create_app(args.input_root, args.job_root, workers=args.workers,
                     queue_size=args.queue_size, threads_per_job=args.threads_per_job,
                     max_weights=args.max_weights, token=token)
    uvicorn.run(app, host=args.host, port=args.port)


if __name__ == "__main__":
    main()