"""Upload client-selected data to an AEStudio job without putting bytes in chat.

Requires only Python 3's standard library. The output is a compact JSON status.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path


def main() -> None:
    parser = argparse.ArgumentParser(description="Upload local files into an existing AEStudio MCP task")
    parser.add_argument("--server", default="https://aestudioai.com")
    parser.add_argument("--job-id", required=True)
    parser.add_argument("--token", help="Task token; alternatively read it from --job-json.")
    parser.add_argument("--job-json", type=Path, help="JSON file containing job_id/token, or raw MCP structuredContent.")
    parser.add_argument("--file", type=Path, action="append", default=[], help="Files in the same order as the preparation tool.")
    parser.add_argument("--download-result", type=Path, help="Download the completed private result ZIP instead of uploading files.")
    parser.add_argument("--start", action="store_true", help="Start processing after upload.")
    parser.add_argument("--api-key-env", default="AESTUDIO_MCP_KEY", help="Environment variable holding the assigned MCP API key.")
    args = parser.parse_args()
    if not re.fullmatch(r"\d{8}-\d{6}-[a-f0-9]{8,32}", args.job_id):
        parser.error("Invalid task ID.")
    if not args.file and not args.download_result:
        parser.error("Provide --file for upload or --download-result for downloading.")
    if args.file and args.download_result:
        parser.error("Choose either upload or download in one invocation.")
    server = args.server.rstrip("/")
    parsed_url = urllib.parse.urlsplit(server)
    if parsed_url.scheme != "https" and not (
        parsed_url.scheme == "http" and parsed_url.hostname in {"127.0.0.1", "localhost"}
    ):
        parser.error("Use HTTPS, or localhost HTTP for testing.")
    token = args.token
    if args.job_json:
        credentials = json.loads(args.job_json.read_text(encoding="utf-8-sig"))
        credentials = credentials.get("structuredContent", credentials)
        if credentials.get("job_id") != args.job_id:
            parser.error("The credentials file belongs to a different task.")
        token = credentials.get("token")
    if not token:
        parser.error("Provide --token or --job-json.")
    if args.download_result:
        if args.download_result.exists():
            parser.error("The result path already exists. Choose a new destination.")
        headers = {"X-AEStudio-Upload-Token": token, "User-Agent": "AEStudio-MCP/downloader"}
        api_key = os.environ.get(args.api_key_env)
        if api_key:
            headers["Authorization"] = "Bearer " + api_key
        req = urllib.request.Request(server + "/downloads/process/" + args.job_id + ".zip", headers=headers)
        with urllib.request.urlopen(req, timeout=60) as response:
            if not response.headers.get("Content-Type", "").startswith("application/zip"):
                raise ValueError("The result service did not return a ZIP archive.")
            args.download_result.parent.mkdir(parents=True, exist_ok=True)
            with args.download_result.open("xb") as output:
                while chunk := response.read(1024 * 1024):
                    output.write(chunk)
        print(json.dumps({"job_id": args.job_id, "downloaded": str(args.download_result), "bytes": args.download_result.stat().st_size}, ensure_ascii=False))
        return
    for file in args.file:
        if not file.is_file():
            parser.error(f"File does not exist: {file}")
    if sum(file.stat().st_size for file in args.file) > 1024 * 1024 * 1024:
        parser.error("Each task is limited to 1 GB. Split larger datasets into separate tasks.")

    def request(method: str, path: str, body: bytes | None = None, content_type: str = "application/json") -> dict:
        headers = {"X-AEStudio-Upload-Token": token, "Content-Type": content_type, "User-Agent": "AEStudio-MCP/uploader"}
        api_key = os.environ.get(args.api_key_env)
        if api_key:
            headers["Authorization"] = "Bearer " + api_key
        req = urllib.request.Request(server + path, data=body, headers=headers, method=method)
        with urllib.request.urlopen(req, timeout=60) as response:
            return json.load(response)

    for index, file in enumerate(args.file):
        offset = 0
        with file.open("rb") as source:
            while chunk := source.read(8 * 1024 * 1024):
                query = urllib.parse.urlencode({"job_id": args.job_id, "file_index": index, "offset": offset})
                for attempt in range(3):
                    try:
                        result = request("POST", "/api/process/chunk?" + query, chunk, "application/octet-stream")
                        break
                    except urllib.error.HTTPError as exc:
                        if exc.code < 500 or attempt == 2:
                            raise
                        time.sleep(attempt + 1)
                    except (OSError, TimeoutError):
                        if attempt == 2:
                            raise
                        time.sleep(attempt + 1)
                offset = int(result["received"])
                print(f"file {index + 1}/{len(args.file)}: {offset}/{file.stat().st_size} bytes", file=sys.stderr)
    if args.start:
        status = request("POST", "/api/process/start", json.dumps({"job_id": args.job_id}).encode("utf-8"))
    else:
        status = request("GET", "/api/process/status?" + urllib.parse.urlencode({"job_id": args.job_id}))
    print(json.dumps({"job_id": args.job_id, "status": status.get("status"), "uploaded_bytes": status.get("uploaded_bytes")}, ensure_ascii=False))


if __name__ == "__main__":
    try:
        main()
    except urllib.error.HTTPError as exc:
        message = exc.read(4000).decode("utf-8", errors="replace")
        print(f"AEStudio upload failed (HTTP {exc.code}): {message}", file=sys.stderr)
        raise SystemExit(1)
    except (OSError, ValueError) as exc:
        print(f"AEStudio upload failed: {exc}", file=sys.stderr)
        raise SystemExit(1)