#!/usr/bin/env python3

"""Submit, follow, and download a user-defined Vanellus simulation."""

import argparse
from collections.abc import Iterator
import io
import json
import os
from pathlib import Path
import threading
import time
from typing import Any
import zipfile

# These imports support the shared runtime inserted at the INCLUDE marker below.
# API responses contain dynamically shaped JSON values that are validated by the server.
JsonObject = dict[str, Any]
DOWNLOADABLE_STATES = {
    "canceled",
    "dry_run",
    "max_iterations_reached",
    "monitor_converged",
    "residual_converged",
}


# ANCHOR: build_request
def build_request() -> JsonObject:
    """Build and return the complete simulation request supplied by the user."""
    # Replace this exception with the Python code that constructs your request.
    raise NotImplementedError("Add your simulation request to build_request()")


# ANCHOR_END: build_request


def requests_client() -> Any:
    """Load Requests only when the downloaded client starts network activity."""
    # Documentation builds import build_request(), which does not need the HTTP dependency.
    try:
        import requests
    except ModuleNotFoundError as error:
        raise RuntimeError("Install the client requirements before running a simulation") from error
    return requests


def api_url(path: str) -> str:
    """Resolve a path against the public Vanellus API."""
    return f"https://api.vanellus.tech/{path.lstrip('/')}"


def request_headers(api_key: str, accept: str) -> dict[str, str]:
    """Return authentication and response-format headers for one request."""
    # Keep the API key in a header rather than placing it in a URL or request body.
    return {"X-API-Key": api_key, "Accept": accept}


def require_success(response: Any, action: str) -> None:
    """Raise a concise error containing the API response body when a request fails."""
    # The response body normally contains more useful validation detail than a traceback.
    if response.status_code != 200:
        raise RuntimeError(f"{action} failed with HTTP {response.status_code}: {response.text}")


# ANCHOR: submit_request
def submit_request(request: JsonObject, api_key: str) -> int:
    """Submit one simulation request and return its server-generated ID."""
    requests = requests_client()
    response = requests.post(
        api_url("/simulations"),
        json=request,
        headers=request_headers(api_key, "application/json"),
        timeout=120,
    )
    require_success(response, "Simulation submission")

    submission = response.json()
    simulation_id = int(submission["id"])
    print(f"Submitted simulation {simulation_id}")
    # Accepted requests can still contain setup warnings.
    for warning in submission.get("warnings", []):
        print(f"API warning: {warning}")
    return simulation_id


# ANCHOR_END: submit_request


# ANCHOR: progress_updates
def iter_sse_data(response: Any) -> Iterator[str]:
    """Yield complete data payloads from a server-sent-events response."""
    data_lines: list[str] = []
    # One SSE event ends at a blank line and may contain more than one data line.
    for raw_line in response.iter_lines(decode_unicode=True):
        if raw_line is None:
            continue
        if isinstance(raw_line, bytes):
            raw_line = raw_line.decode(response.encoding or "utf-8", errors="replace")
        line = raw_line.rstrip("\r\n")
        if not line:
            if data_lines:
                yield "\n".join(data_lines)
                data_lines.clear()
        elif line.startswith("data:"):
            data_lines.append(line.removeprefix("data:").removeprefix(" "))
    if data_lines:
        yield "\n".join(data_lines)


def format_progress_value(value: Any) -> str:
    """Format one live diagnostic, preserving the API's null marker for non-finite output."""
    return "NaN" if value is None else f"{float(value):.3e}"


def print_progress_updates(
    api_key: str,
    simulation_id: int,
    stop_event: threading.Event,
) -> None:
    """Print live residual updates until the run or the client thread stops."""
    requests = requests_client()
    try:
        with requests.get(
            api_url(f"/simulations/{simulation_id}/progress_stream"),
            headers=request_headers(api_key, "text/event-stream"),
            stream=True,
            timeout=(10, None),
        ) as response:
            require_success(response, "Progress stream")
            for payload in iter_sse_data(response):
                if stop_event.is_set():
                    return
                update = json.loads(payload)
                residuals = update.get("residuals", {})
                residual_text = ", ".join(
                    f"{name}={format_progress_value(value)}" for name, value in sorted(residuals.items())
                )
                print(f"Iteration {update['iteration']}: {residual_text}", flush=True)
    except (requests.RequestException, RuntimeError, json.JSONDecodeError) as error:
        # Status polling remains authoritative if the optional stream disconnects.
        if not stop_event.is_set():
            print(f"Progress stream stopped: {error}")


# ANCHOR_END: progress_updates


def get_status(api_key: str, simulation_id: int) -> JsonObject:
    """Fetch the current lifecycle and solver status for one simulation."""
    requests = requests_client()
    response = requests.get(
        api_url(f"/simulations/{simulation_id}/status"),
        headers=request_headers(api_key, "application/json"),
        timeout=120,
    )
    require_success(response, "Status request")
    return response.json()


def cancel_simulation(api_key: str, simulation_id: int) -> None:
    """Request a recoverable stop at the next solver iteration boundary."""
    requests = requests_client()
    # Cancellation preserves partial results, unlike an immediate kill request.
    response = requests.post(
        api_url(f"/simulations/{simulation_id}/cancel"),
        headers=request_headers(api_key, "application/json"),
        timeout=120,
    )
    require_success(response, "Cancellation")
    print(f"Cancellation requested for simulation {simulation_id}")


# ANCHOR: wait_for_completion
def wait_for_completion(
    api_key: str,
    simulation_id: int,
    dry_run: bool = False,
) -> JsonObject:
    """Stream progress while polling the authoritative simulation status."""
    stop_event = threading.Event()
    reported_cells = False
    progress_thread: threading.Thread | None = None
    # Dry runs have no iterative residual history, so they only need status polling.
    if not dry_run:
        progress_thread = threading.Thread(
            target=print_progress_updates,
            kwargs={
                "api_key": api_key,
                "simulation_id": simulation_id,
                "stop_event": stop_event,
            },
            daemon=True,
        )
        progress_thread.start()
    try:
        while True:
            status = get_status(api_key=api_key, simulation_id=simulation_id)
            if status.get("num_cells") and not reported_cells:
                print(f"Mesh cells: {int(status['num_cells']):,}", flush=True)
                reported_cells = True
            if status.get("status") in {"diverged", "error", "killed"} or status.get("error"):
                raise RuntimeError(f"Simulation stopped without results: {status}")
            if status.get("completed"):
                return status
            time.sleep(2)
    except KeyboardInterrupt:
        # Ctrl-C asks the server to preserve a partial result at an iteration boundary.
        cancel_simulation(api_key=api_key, simulation_id=simulation_id)
        raise SystemExit(130) from None
    finally:
        stop_event.set()
        if progress_thread is not None:
            progress_thread.join(timeout=5)


# ANCHOR_END: wait_for_completion


def extract_zip_safely(archive_bytes: bytes, output_directory: Path) -> None:
    """Extract a result archive after rejecting paths outside the destination."""
    output_directory.mkdir(parents=True, exist_ok=True)
    output_root = output_directory.resolve()
    # Validate every member before extracting any file from the downloaded archive.
    with zipfile.ZipFile(io.BytesIO(archive_bytes)) as archive:
        for member in archive.infolist():
            destination = (output_directory / member.filename).resolve()
            if destination != output_root and not destination.is_relative_to(output_root):
                raise RuntimeError(f"Unsafe path in result archive: {member.filename!r}")
        archive.extractall(output_directory)


# ANCHOR: download_results
def download_results(
    api_key: str,
    simulation_id: int,
    output_directory: Path,
) -> None:
    """Download and safely extract every available result artifact."""
    requests = requests_client()
    response = requests.get(
        api_url(f"/simulations/{simulation_id}/download"),
        headers=request_headers(api_key, "application/zip"),
        timeout=None,
    )
    require_success(response, "Results download")
    extract_zip_safely(response.content, output_directory)
    print(f"Extracted results to {output_directory}")


# ANCHOR_END: download_results
# ANCHOR: run_simulation
def run_simulation(
    request: JsonObject,
    api_key: str,
    output_root: Path,
    output_name: str | None = None,
) -> Path:
    """Submit, follow, download, and return one simulation's result directory."""
    # Reject an occupied custom folder before submitting a simulation that may use credits.
    if output_name is not None:
        planned_output = output_root / output_name
        if planned_output.exists() and (not planned_output.is_dir() or any(planned_output.iterdir())):
            raise FileExistsError(f"Output directory is not empty: {planned_output}")

    simulation_id = submit_request(request, api_key=api_key)
    status = wait_for_completion(api_key=api_key, simulation_id=simulation_id)
    output_directory = output_root / (output_name or f"simulation-{simulation_id}")
    download_results(
        api_key=api_key,
        simulation_id=simulation_id,
        output_directory=output_directory,
    )

    # Preserve the exact client-side request alongside the server-normalized request.json.
    submitted_path = output_directory / "submitted-request.json"
    submitted_path.write_text(json.dumps(request, indent=2) + "\n", encoding="utf-8")

    final_state = str(status["status"])
    print(f"Simulation finished with status {final_state}")
    if final_state == "max_iterations_reached":
        print("Warning: inspect convergence before using this result")
    elif final_state == "canceled":
        print("Warning: the downloaded result is from a canceled simulation")
    elif final_state not in DOWNLOADABLE_STATES:
        raise RuntimeError(f"Simulation ended as {final_state}; diagnostic results were retained")
    return output_directory


# ANCHOR_END: run_simulation


def parse_arguments() -> argparse.Namespace:
    """Parse the result location."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=Path("results"), help="parent directory for downloaded results")
    parser.add_argument("--name", help="result folder name (default: simulation ID)")
    arguments = parser.parse_args()
    # Keep custom names inside the selected output directory.
    if arguments.name and (Path(arguments.name).name != arguments.name or arguments.name in {".", ".."}):
        parser.error("--name must be a folder name; use --output to choose its parent directory")
    return arguments


def main() -> None:
    """Read credentials, build the request, and run the selected workflow."""
    arguments = parse_arguments()
    # Read the secret only from the environment so it is not stored in source or shell history.
    api_key = os.environ.get("VANELLUS_API_KEY")
    if not api_key:
        raise SystemExit("Set VANELLUS_API_KEY before running this script")

    # Model construction stays separate from the reusable API workflow below it.
    request = build_request()
    run_simulation(
        request,
        api_key=api_key,
        output_root=arguments.output,
        output_name=arguments.name,
    )


if __name__ == "__main__":
    main()