#!/usr/bin/env python3
"""Upload a folder to an AO project, or download a project's files to a folder.

  python3 ao_files.py upload   <folder> <project>
  python3 ao_files.py download <project> <folder>

<project> is the project's name, or its id if several projects share the name.
Set AO_URL (such as https://ao.example.com) and AO_EMAIL; the script asks for your
password. On MCP-enabled AO, you can set AO_TOKEN to an API token instead.
"""

import getpass
import hashlib
import json
import os
import re
import sys
from pathlib import Path

import requests

API = os.environ["AO_URL"].rstrip("/") + "/api/v1"
session = requests.Session()
session.headers["User-Agent"] = "ao-files-script"
refresh_token = None


def use_tokens(response):
    """Keeps the tokens from a sign-in or renewal's user-context header."""
    global refresh_token
    response.raise_for_status()
    tokens = json.loads(response.headers["user-context"])
    session.headers["Authorization"] = "Bearer " + tokens["auth_token"]
    refresh_token = tokens["refresh_token"]


def sign_in():
    if os.environ.get("AO_TOKEN"):
        session.headers["Authorization"] = "Bearer " + os.environ["AO_TOKEN"]
        return
    email = os.environ["AO_EMAIL"]
    password = getpass.getpass(f"AO password for {email}: ")
    use_tokens(
        session.post(
            f"{API}/users/login",
            headers={"X-Email": email, "X-Password": password},
            timeout=60,
        )
    )


def sign_out():
    """Ends a password sign-in's session, so it doesn't count toward an account's 10."""
    if refresh_token:
        try:
            call("POST", "/users/logout", timeout=60)
        except requests.RequestException:
            pass


def call(method, path, upload=None, **options):
    """Calls AO. After a 401, renews the sign-in once and tries again."""
    options.setdefault("timeout", (30, 3600))
    for attempt in (1, 2):
        if upload:
            with open(upload, "rb") as f:
                response = session.request(
                    method, API + path, files={"file": (upload.name, f)}, **options
                )
        else:
            response = session.request(method, API + path, **options)
        if response.status_code == 401 and refresh_token and attempt == 1:
            use_tokens(
                session.post(
                    f"{API}/users/refresh-token",
                    headers={"X-Refresh-Token": refresh_token},
                    timeout=60,
                )
            )
            continue
        response.raise_for_status()
        return response


def find_project(name_or_id):
    if re.fullmatch(
        r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", name_or_id
    ):
        return name_or_id
    me = call("GET", "/users/me").json()["items"]
    projects = call("GET", f"/surveys/users/{me['userid']}/preview").json()["items"][0][
        "surveys"
    ]
    found = [p for p in projects if p["name"].lower() == name_or_id.lower()]
    if len(found) != 1:
        listed = "".join(
            f"\n  {p['uuid']}  {p['name']} ({p['org_name']})" for p in found
        )
        sys.exit(
            f"{len(found)} projects are named {name_or_id!r}. "
            f"Give the project's id instead.{listed}"
        )
    return found[0]["uuid"]


def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def key(file_path, name):
    """AO matches a file by its folder and name, ignoring case."""
    return ((file_path or "").lower(), name.lower())


def upload(folder, project):
    project_id = find_project(project)
    local = {}
    for path in sorted(Path(folder).rglob("*")):
        relative = path.relative_to(folder)
        if path.is_file() and not any(part.startswith(".") for part in relative.parts):
            file_path = (
                "" if relative.parent == Path(".") else relative.parent.as_posix() + "/"
            )
            local[key(file_path, path.name)] = (file_path, path)
    sent, already, not_sent, failed = 0, 0, [], []
    keys = list(local)
    for start in range(0, len(keys), 200):
        batch = keys[start : start + 200]
        records = [
            {
                "name": local[k][1].name,
                "filePath": local[k][0],
                "hash": sha256(local[k][1]),
                "on_conflict": "SKIP",
            }
            for k in batch
        ]
        answer = call(
            "POST", f"/surveys/{project_id}/files", json={"files": {"files": records}}
        ).json()
        statuses = {
            key(s.get("filePath"), s["name"]): s
            for s in answer["items"][0]["fileCreationStatus"]
        }
        for k in batch:
            file_path, path = local[k]
            status = statuses.get(k)
            if status is None:
                not_sent.append(file_path + path.name)
            elif status["status"] == "dup":
                already += 1
            else:
                try:
                    call("POST", f"/file-management/geo/{status['fid']}", upload=path)
                    sent += 1
                    print("sent", file_path + path.name)
                except requests.RequestException as error:
                    failed.append(f"{file_path}{path.name}: {error}")
    print(
        f"{sent} sent, {already} already in AO, "
        f"{len(not_sent)} not sent, {len(failed)} failed"
    )
    for name in not_sent:
        print("  not sent, because a different file with this name is in AO:", name)
    for name in failed:
        print("  failed (run again to retry):", name)
    return not not_sent and not failed


def download(project, folder):
    project_id = find_project(project)
    files = call("GET", f"/surveys/{project_id}/files").json()["items"][0][
        "fileStatusList"
    ]
    fetched, same, failed = 0, 0, []
    for f in files:
        if f["isVirtualDir"]:
            continue
        target = Path(folder, f["filePath"] or "", f["name"])
        if not target.resolve().is_relative_to(Path(folder).resolve()):
            failed.append(f"{target}: outside {folder}")
            continue
        if target.exists() and sha256(target) == f["hash"]:
            same += 1
            continue
        target.parent.mkdir(parents=True, exist_ok=True)
        partial = target.with_name(target.name + ".part")
        try:
            with (
                call(
                    "GET", f"/file-management/geo/{f['fileUuid']}", stream=True
                ) as response,
                open(partial, "wb") as out,
            ):
                for chunk in response.iter_content(1024 * 1024):
                    out.write(chunk)
        except requests.RequestException as error:
            failed.append(f"{target}: {error}")
            continue
        if sha256(partial) == f["hash"]:
            partial.replace(target)
            fetched += 1
            print("downloaded", target)
        else:
            failed.append(f"{target}: the download doesn't match AO's copy")
        partial.unlink(missing_ok=True)
    print(f"{fetched} downloaded, {same} already here, {len(failed)} failed")
    for name in failed:
        print("  failed (run again to retry):", name)
    return not failed


if __name__ == "__main__":
    if len(sys.argv) != 4 or sys.argv[1] not in ("upload", "download"):
        sys.exit(__doc__)
    sign_in()
    try:
        ok = (
            upload(sys.argv[2], sys.argv[3])
            if sys.argv[1] == "upload"
            else download(sys.argv[2], sys.argv[3])
        )
    finally:
        sign_out()
    sys.exit(0 if ok else 1)