from __future__ import annotations

import argparse
import base64
import json
import os
import shutil
import subprocess
import time
from pathlib import Path

import yaml


NPX_EXECUTABLE = shutil.which("npx.cmd") or shutil.which("npx") or "npx"
PLAYWRIGHT_CMD = [NPX_EXECUTABLE, "--yes", "--package", "@playwright/cli", "playwright-cli"]
SESSION_NAME = "vip-report-monthly-sales"
DOWNLOAD_CENTER_URL = "https://compass.vip.com/frontend/index.html#/dataCapturing/downloadCenter"


def run_cmd(args: list[str], *, cwd: Path, timeout: int = 300) -> subprocess.CompletedProcess[str]:
    env = os.environ.copy()
    npm_cache = cwd / ".npm-cache"
    npm_cache.mkdir(parents=True, exist_ok=True)
    env["NPM_CONFIG_CACHE"] = str(npm_cache)
    env["npm_config_cache"] = str(npm_cache)
    return subprocess.run(
        args,
        cwd=str(cwd),
        text=True,
        encoding="utf-8",
        errors="replace",
        capture_output=True,
        timeout=timeout,
        check=True,
        env=env,
    )


def build_js(task_id: str) -> str:
    payload = json.dumps({"task_id": task_id, "url": DOWNLOAD_CENTER_URL}, ensure_ascii=False)
    return f"""async function(page) {{
  const spec = {payload};
  await page.goto(spec.url, {{ waitUntil: 'domcontentloaded', timeout: 120000 }});
  await page.waitForTimeout(1000);
  const result = await page.evaluate(async (taskId) => {{
    const body = new URLSearchParams();
    body.set('reqData', JSON.stringify({{ taskId }}));
    const res = await fetch('/report/asyncdownload/getOutputFile', {{
      method: 'POST',
      headers: {{ 'Content-Type': 'application/x-www-form-urlencoded' }},
      body: body.toString(),
      credentials: 'include',
    }});
    const ab = await res.arrayBuffer();
    const bytes = new Uint8Array(ab);
    let binary = '';
    const chunkSize = 0x8000;
    for (let offset = 0; offset < bytes.length; offset += chunkSize) {{
      binary += String.fromCharCode(...bytes.subarray(offset, offset + chunkSize));
    }}
    return {{
      ok: res.ok,
      status: res.status,
      contentType: res.headers.get('content-type') || '',
      disposition: res.headers.get('content-disposition') || '',
      base64: btoa(binary),
    }};
  }}, String(spec.task_id));
  return result;
}}
"""


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Download a ready Compass async task file by task id.")
    p.add_argument("--config", default=r"D:\projects\skills\vip-report\config.yaml")
    p.add_argument("--session", default=SESSION_NAME)
    p.add_argument("--task-id", required=True)
    p.add_argument("--output", required=True)
    return p.parse_args()


def main() -> None:
    args = parse_args()
    cfg = yaml.safe_load(Path(args.config).read_text(encoding="utf-8"))
    workdir = Path(cfg["paths"]["workdir"]).resolve()
    js_path = workdir / f"tmp-download-task-{args.task_id}-{time.time_ns()}.js"
    js_path.write_text(build_js(args.task_id), encoding="utf-8")
    try:
        out = run_cmd(
            PLAYWRIGHT_CMD + ["--session", args.session, "run-code", "--filename", str(js_path)],
            cwd=workdir,
            timeout=420,
        ).stdout
    finally:
        js_path.unlink(missing_ok=True)

    if "### Error" in out:
        raise RuntimeError(out.strip())
    json_text = ""
    marker = "### Result"
    marker_idx = out.find(marker)
    if marker_idx >= 0:
        start = out.find("\n", marker_idx)
        if start >= 0:
            end = out.find("\n###", start + 1)
            json_text = out[start + 1 : end if end >= 0 else None].strip()
    if not json_text:
        for line in out.splitlines():
            s = line.strip()
            if s.startswith("{") and s.endswith("}"):
                json_text = s
                break
    if not json_text:
        raise RuntimeError(f"No JSON result from run-code: {out}")
    result = json.loads(json_text)
    if not result.get("ok", False):
        raise RuntimeError(json.dumps(result, ensure_ascii=False))

    output_path = Path(args.output).resolve()
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_bytes(base64.b64decode(result["base64"]))
    result.pop("base64", None)
    result["savedAs"] = str(output_path)
    print(json.dumps(result, ensure_ascii=True))


if __name__ == "__main__":
    main()
