#!/usr/bin/env python3
"""Orthogonal ImageFind WebDAV upload stress test with isolated cleanup.

The default profile is CI-friendly. Pass ``--full`` to execute the requested
10 MiB / 1 GiB / 5 GiB singles, ten-file batch and fifty 10 MiB-file batch.
The API token is used for both WebDAV Basic authentication and cleanup.
"""

from __future__ import annotations

import argparse
import base64
import concurrent.futures
import hashlib
import os
import sys
import threading
import time
import uuid
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from urllib.parse import quote

import httpx

MIB = 1024**2
GIB = 1024**3
CHUNK = MIB


@dataclass(frozen=True)
class Case:
    virtual_path: str
    size: int
    resume: bool = False


@dataclass
class TimedBody:
    name: str
    size: int
    prefix: bytes = b""
    start: int = 0
    end: int | None = None
    finished_at: float | None = None

    def __iter__(self) -> Iterator[bytes]:
        yield from blocks(
            self.name,
            self.size,
            prefix=self.prefix,
            start=self.start,
            end=self.end,
        )
        # httpx resumes the iterator after it has handed the final request
        # bytes to the transport. The remaining interval is server response
        # latency, including any accidental post-PUT whole-file work.
        self.finished_at = time.monotonic()


def blocks(
    name: str,
    size: int,
    *,
    prefix: bytes = b"",
    start: int = 0,
    end: int | None = None,
) -> Iterator[bytes]:
    end = size if end is None else end
    seed = hashlib.sha256(name.encode()).digest()
    block = (seed * (CHUNK // len(seed) + 1))[:CHUNK]
    position = start
    while position < end:
        if position < len(prefix):
            length = min(len(prefix) - position, end - position)
            yield prefix[position : position + length]
            position += length
            continue
        offset = position % CHUNK
        length = min(CHUNK - offset, end - position)
        yield block[offset : offset + length]
        position += length


def encoded_path(path: str) -> str:
    return "/".join(quote(part, safe="") for part in path.split("/"))


class StressRun:
    def __init__(
        self,
        base_url: str,
        token: str,
        full: bool,
        large: bool,
        workers: int,
        timeout: float,
        probe_timeout: float,
        media_prefix: bytes = b"",
    ):
        self.base_url = base_url.rstrip("/")
        self.token = token
        self.full = full
        self.large = large
        self.workers = workers
        self.timeout = timeout
        self.probe_timeout = probe_timeout
        self.media_prefix = media_prefix
        self.run_id = f"ifstress-{uuid.uuid4().hex[:10]}"
        basic = base64.b64encode(f"imagefind:{token}".encode()).decode()
        self.dav_headers = {"Authorization": f"Basic {basic}"}
        self.api_headers = {"Authorization": f"Bearer {token}"}
        self.upload_ids: set[str] = set()
        self.request_timeout = httpx.Timeout(connect=15, write=60, read=120, pool=30)
        self.probe_latencies: list[float] = []
        self.probe_errors: list[str] = []

    def probe(self, stopping: threading.Event) -> None:
        timeout = httpx.Timeout(connect=3, write=3, read=self.probe_timeout, pool=3)
        while not stopping.is_set():
            started = time.monotonic()
            try:
                response = httpx.get(f"{self.base_url}/api/v1/status", timeout=timeout)
                response.raise_for_status()
                self.probe_latencies.append(time.monotonic() - started)
            except Exception as exc:
                if len(self.probe_errors) < 10:
                    self.probe_errors.append(f"{type(exc).__name__}: {exc}")
            stopping.wait(0.25)

    def cases(self) -> list[Case]:
        if self.large:
            return [
                Case(f"{self.run_id}/large-{index}.mp4", 256 * MIB)
                for index in range(4)
            ]
        single_sizes = [10 * MIB, GIB, 5 * GIB] if self.full else [10 * MIB]
        cases = [
            Case(f"{self.run_id}/single-{size}.mp4", size, resume=index == 0)
            for index, size in enumerate(single_sizes)
        ]
        cases.append(Case(f"root-{self.run_id}.mp4", 10 * MIB if self.full else MIB))
        for index in range(10):
            depth = (1, 5, 10)[index % 3]
            groups = "/".join(f"level-{level}" for level in range(1, depth + 1))
            size = 10 * MIB if self.full else MIB
            cases.append(Case(f"{self.run_id}/{groups}/{self.run_id}-ten-{index:02d}.mp4", size))
        for index in range(50):
            size = 10 * MIB if self.full else 256 * 1024
            cases.append(Case(f"{self.run_id}/fifty/{self.run_id}-fifty-{index:02d}.mp4", size))
        return cases

    def put(self, case: Case) -> tuple[str, int, str | None, float]:
        url = f"{self.base_url}/webdav/{encoded_path(case.virtual_path)}"
        headers = dict(self.dav_headers)
        upload_id = None
        response_body: TimedBody
        with httpx.Client(timeout=self.request_timeout) as client:
            if case.resume:
                split = case.size // 2
                first_body = TimedBody(
                    case.virtual_path, case.size, self.media_prefix, end=split
                )
                first = client.put(
                    url,
                    headers={
                        **headers,
                        "Content-Length": str(split),
                        "Content-Range": f"bytes 0-{split - 1}/{case.size}",
                    },
                    content=first_body,
                )
                first.raise_for_status()
                if first.headers.get("upload-offset") != str(split):
                    raise RuntimeError(f"{case.virtual_path}: first offset was not persisted")
                head = client.head(url, headers=headers)
                head.raise_for_status()
                if head.headers.get("upload-offset") != str(split):
                    raise RuntimeError(f"{case.virtual_path}: HEAD did not report the partial offset")
                response_body = TimedBody(
                    case.virtual_path, case.size, self.media_prefix, start=split
                )
                response = client.put(
                    url,
                    headers={
                        **headers,
                        "Content-Length": str(case.size - split),
                        "Content-Range": f"bytes {split}-{case.size - 1}/{case.size}",
                    },
                    content=response_body,
                )
            else:
                response_body = TimedBody(case.virtual_path, case.size, self.media_prefix)
                response = client.put(
                    url,
                    headers={**headers, "Content-Length": str(case.size)},
                    content=response_body,
                )
            response.raise_for_status()
            upload_id = response.headers.get("x-imagefind-upload-id")
            response_latency = time.monotonic() - (response_body.finished_at or time.monotonic())
        return case.virtual_path, case.size, upload_id, response_latency

    def wait(self, expected: int) -> None:
        deadline = time.monotonic() + self.timeout
        while time.monotonic() < deadline:
            response = httpx.get(
                f"{self.base_url}/api/v1/uploads?limit=500",
                headers=self.api_headers,
                timeout=30,
            )
            response.raise_for_status()
            rows = [row for row in response.json() if row.get("id") in self.upload_ids]
            failures = [row for row in rows if row.get("status") == "failed"]
            if failures:
                raise RuntimeError(f"upload failures: {failures}")
            if len(rows) == expected and all(row.get("status") == "completed" for row in rows):
                return
            time.sleep(2)
        raise TimeoutError("uploads did not complete before the stress timeout")

    def upload_rows(self) -> list[dict]:
        response = httpx.get(
            f"{self.base_url}/api/v1/uploads?limit=500",
            headers=self.api_headers,
            timeout=self.request_timeout,
        )
        response.raise_for_status()
        return response.json()

    def verify_dedup(self, case: Case) -> None:
        url = f"{self.base_url}/webdav/{encoded_path(case.virtual_path)}"
        before = {row["id"] for row in self.upload_rows()}
        body = TimedBody(case.virtual_path, case.size, self.media_prefix)
        response = httpx.put(
            url,
            headers={
                **self.dav_headers,
                "Content-Length": str(case.size),
            },
            content=body,
            timeout=self.request_timeout,
        )
        response.raise_for_status()
        if response.headers.get("x-imagefind-deduplicated") != "true":
            raise RuntimeError("full retry after a simulated lost response was not deduplicated")
        after = {row["id"] for row in self.upload_rows()}
        if after != before:
            raise RuntimeError("full retry created another upload task")
        latency = time.monotonic() - (body.finished_at or time.monotonic())
        print(
            f"deduplicated lost-response retry response_latency={latency:.3f}s",
            flush=True,
        )

    def cleanup(self) -> None:
        try:
            videos = httpx.get(
                f"{self.base_url}/api/v1/videos?limit=500", headers=self.api_headers, timeout=30
            ).json()
            for video in videos:
                if self.run_id in str(video.get("display_name")) or self.run_id in str(video.get("source_key")):
                    response = httpx.delete(
                        f"{self.base_url}/api/v1/videos/{video['id']}?delete_source=true",
                        headers=self.api_headers,
                        timeout=60,
                    )
                    if response.status_code != 404:
                        response.raise_for_status()
            collections = httpx.get(
                f"{self.base_url}/api/v1/collections", headers=self.api_headers, timeout=30
            ).json()
            for collection in collections:
                if collection.get("name") == self.run_id:
                    httpx.delete(
                        f"{self.base_url}/api/v1/collections/{collection['id']}",
                        headers=self.api_headers,
                        timeout=30,
                    ).raise_for_status()
            trash = httpx.get(
                f"{self.base_url}/api/v1/trash", headers=self.api_headers, timeout=30
            ).json()
            for item in trash:
                if self.run_id in str(item.get("display_name")) or self.run_id in str(item.get("original_key")):
                    response = httpx.delete(
                        f"{self.base_url}/api/v1/trash/{item['id']}",
                        headers=self.api_headers,
                        timeout=60,
                    )
                    if response.status_code != 404:
                        response.raise_for_status()
        except Exception as exc:
            print(f"cleanup warning: {exc}", file=sys.stderr)

    def execute(self, keep: bool) -> None:
        cases = self.cases()
        started = time.monotonic()
        probe_stopping = threading.Event()
        probe_thread = threading.Thread(
            target=self.probe,
            args=(probe_stopping,),
            name="imagefind-stress-probe",
            daemon=True,
        )
        probe_thread.start()
        try:
            with concurrent.futures.ThreadPoolExecutor(max_workers=self.workers) as executor:
                futures = [executor.submit(self.put, case) for case in cases]
                for future in concurrent.futures.as_completed(futures):
                    path, size, upload_id, response_latency = future.result()
                    if upload_id:
                        self.upload_ids.add(upload_id)
                    print(
                        f"received {path} ({size} bytes) "
                        f"response_latency={response_latency:.3f}s",
                        flush=True,
                    )
            self.wait(len(self.upload_ids))
            self.verify_dedup(cases[0])
            if self.probe_errors:
                raise RuntimeError(
                    "ImageFind API became unreachable during WebDAV transfer: "
                    + "; ".join(self.probe_errors)
                )
            maximum_probe = max(self.probe_latencies, default=0)
            print(
                f"PASS run={self.run_id} files={len(cases)} bytes={sum(case.size for case in cases)} "
                f"seconds={time.monotonic() - started:.1f} max_probe_latency={maximum_probe:.3f}s",
                flush=True,
            )
        finally:
            probe_stopping.set()
            probe_thread.join(timeout=self.probe_timeout + 5)
            if not keep:
                self.cleanup()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", required=True, help="ImageFind native URL, e.g. http://nas:8765")
    parser.add_argument("--token", default=os.getenv("IMAGEFIND_WEBDAV_TOKEN"))
    profile = parser.add_mutually_exclusive_group()
    profile.add_argument("--full", action="store_true", help="include 1 GiB, 5 GiB and full-size batches")
    profile.add_argument("--large", action="store_true", help="run four concurrent 256 MiB uploads")
    parser.add_argument("--workers", type=int, default=8)
    parser.add_argument("--timeout", type=float, default=7200)
    parser.add_argument(
        "--probe-timeout",
        type=float,
        default=10,
        help="fail when the API is unreachable for this many seconds during transfers",
    )
    parser.add_argument("--keep", action="store_true", help="keep isolated stress data for inspection")
    parser.add_argument(
        "--media-prefix",
        help="optional valid media file prepended before deterministic padding",
    )
    args = parser.parse_args()
    if not args.token:
        parser.error("--token or IMAGEFIND_WEBDAV_TOKEN is required")
    media_prefix = Path(args.media_prefix).read_bytes() if args.media_prefix else b""
    StressRun(
        args.base_url,
        args.token,
        args.full,
        args.large,
        min(10, max(1, args.workers)),
        args.timeout,
        max(1, args.probe_timeout),
        media_prefix,
    ).execute(args.keep)


if __name__ == "__main__":
    main()