"""aa_env — AlphaAgent environment checkpoint helper (customer-downloadable).

This module lets a long-running analysis SURVIVE the environment Lambda's wall.
A single Lambda invocation is bounded (≤15 min, and the handler caps it lower to
leave room for the workspace sync). A genuinely heavy job — a backtest, a
parameter sweep, a large simulation — expresses its work as PREDECLARED,
well-defined UNITS. Each unit writes durable output to the workspace and is
idempotent ("done" ⟺ its output exists). When the invocation nears its deadline
the script CHECKPOINTS and exits with a continue-sentinel; the coder harness
re-invokes a FRESH Lambda that resumes from the manifest. Nothing in interpreter
memory is ever needed across invocations — the S3-backed workspace is the truth.

Design invariants (correctness spine):
  * PURE STDLIB. No boto3, no third-party imports — so this file drops into any
    environment image (including a customer's custom image) with a single COPY.
    All S3 sync is done by the handler (workdir↔S3 on every invoke).
  * The manifest (``_manifest.json``) is written ATOMICALLY (temp + os.replace),
    so a crash never leaves it half-written.
  * commit() is the ONLY durability point and its ordering is: write the shard →
    advance the cursor → write the manifest. A hard-kill mid-chunk leaves the
    shard un-referenced and the cursor unmoved, so that chunk is idempotently
    RECOMPUTED (and its shard overwritten) on resume — never double-counted.
  * If ``AA_DEADLINE_EPOCH`` is unset (e.g. an older handler that predates this
    contract), should_yield() is always False: the script runs single-shot and
    degrades gracefully to the handler's hard timeout. Safe on every handler.

Typical usage (fine-grained cursor over a long accumulation)::

    import aa_env as aa
    aa.declare(["fetch", "backtest", "aggregate"])
    if not aa.is_done("fetch"):
        fetch_and_save("master.csv"); aa.done("fetch")
    aa.maybe_yield()
    if not aa.is_done("backtest"):
        for t in aa.resume_range("backtest", total=6387):
            aa.append("backtest", compute_day(t))   # buffered
            aa.tick("backtest", t + 1)               # advance (in-memory)
            if aa.should_yield():
                aa.commit("backtest"); aa.yield_now()
        aa.commit("backtest"); aa.done("backtest")
    if not aa.is_done("aggregate"):
        rows = aa.read_rows("backtest"); write_result(rows); aa.done("aggregate")
    aa.finalize()   # status=complete → return normally (exit 0)
"""

from __future__ import annotations

import json
import os
import sys
import time
from typing import Any, Dict, Iterable, Iterator, List, Optional

# Exit code the script uses to tell the coder harness "checkpointed, re-invoke
# me". MUST match code-interpreter/services/coder_worker.py::EXIT_CONTINUE.
EXIT_CONTINUE = 75

MANIFEST_NAME = "_manifest.json"
_SHARD_DIR = "_aa_shards"

# Cached, lazily-loaded manifest for this invocation.
_state: Optional[Dict[str, Any]] = None
# Per-unit in-memory row buffer for the current (uncommitted) chunk.
_buffers: Dict[str, List[Any]] = {}
# Per-unit pending cursor (advanced by tick(), persisted by commit()).
_pending_cursor: Dict[str, int] = {}


# --------------------------------------------------------------------------- #
# manifest load / atomic save
# --------------------------------------------------------------------------- #
def _fresh() -> Dict[str, Any]:
    return {"units": [], "done": [], "cursor": {}, "shards": {},
            "status": "in_progress", "version": 1}


def _load() -> Dict[str, Any]:
    global _state
    if _state is not None:
        return _state
    try:
        with open(MANIFEST_NAME, "r", encoding="utf-8") as f:
            _state = json.load(f)
        # Defensive: ensure required keys exist even if an older manifest is read.
        for k, v in _fresh().items():
            _state.setdefault(k, v)
    except (FileNotFoundError, ValueError):
        _state = _fresh()
    return _state


def _save() -> None:
    """Atomically persist the manifest to the CWD (the handler syncs it to S3)."""
    state = _load()
    tmp = MANIFEST_NAME + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(state, f, default=str)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, MANIFEST_NAME)  # atomic on POSIX


# --------------------------------------------------------------------------- #
# deadline / yield
# --------------------------------------------------------------------------- #
def deadline_epoch() -> Optional[float]:
    """The soft deadline (unix epoch) injected by the handler, or None."""
    raw = os.environ.get("AA_DEADLINE_EPOCH", "").strip()
    if not raw:
        return None
    try:
        return float(raw)
    except ValueError:
        return None


def should_yield() -> bool:
    """True once we are within the safety margin of the invocation deadline.

    Always False when no deadline is set (older handler) → the script simply
    runs to completion in a single invocation.
    """
    dl = deadline_epoch()
    return dl is not None and time.time() >= dl


def yield_now() -> "None":  # noqa: D401 - never returns
    """Exit the process with the continue-sentinel so the harness re-invokes.

    Call this ONLY after commit()/_save() has persisted progress. The manifest
    is flushed here too as a backstop.
    """
    _save()
    sys.stdout.flush()
    sys.exit(EXIT_CONTINUE)


def maybe_yield() -> None:
    """Coarse gate: if the deadline is near, persist and yield. Use between
    whole units when you are not inside a resume_range loop."""
    if should_yield():
        yield_now()


# --------------------------------------------------------------------------- #
# units (coarse) + cursor (fine)
# --------------------------------------------------------------------------- #
def declare(units: Iterable[str]) -> None:
    """Declare the ordered units of work. Idempotent: on resume this loads the
    existing manifest and preserves done/cursor/shards; on a fresh run it seeds
    a new manifest. Call this once at the top of the script."""
    state = _load()
    state["units"] = list(units)
    # Prune progress for units that no longer exist (defensive on re-declare).
    state["done"] = [u for u in state["done"] if u in state["units"]]
    _save()


def is_done(unit: str) -> bool:
    return unit in _load()["done"]


def done(unit: str) -> None:
    state = _load()
    if unit not in state["done"]:
        state["done"].append(unit)
    _save()


def cursor(unit: str) -> int:
    return int(_load()["cursor"].get(unit, 0))


def resume_range(unit: str, total: int) -> range:
    """A range that starts at the persisted cursor for ``unit`` and ends at
    ``total`` — so a resumed invocation continues exactly where it left off."""
    return range(cursor(unit), int(total))


def tick(unit: str, next_index: int) -> None:
    """Record that everything strictly BEFORE ``next_index`` has been produced
    in this chunk. Persisted only by commit() — advance-then-commit ordering."""
    _pending_cursor[unit] = int(next_index)


def append(unit: str, row: Any) -> None:
    """Buffer one output row for ``unit``. Flushed to a durable shard by commit()."""
    _buffers.setdefault(unit, []).append(row)


def commit(unit: str) -> None:
    """Durably checkpoint ``unit``: write the buffered rows to an immutable shard,
    then advance the cursor, then persist the manifest — in that order.

    A crash before this returns leaves the shard un-referenced and the cursor
    unmoved, so the chunk is recomputed (and its shard overwritten) on resume.
    """
    state = _load()
    buf = _buffers.get(unit, [])
    shards = state["shards"].setdefault(unit, [])
    if buf:
        os.makedirs(_SHARD_DIR, exist_ok=True)
        # Name by the CURRENT shard count so a mid-commit retry overwrites the
        # same orphaned file rather than creating a duplicate.
        shard_name = os.path.join(_SHARD_DIR, f"{unit}.part_{len(shards):04d}.jsonl")
        tmp = shard_name + ".tmp"
        with open(tmp, "w", encoding="utf-8") as f:
            for row in buf:
                f.write(json.dumps(row, default=str) + "\n")
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp, shard_name)
        shards.append(shard_name)          # (1) shard referenced
    if unit in _pending_cursor:
        state["cursor"][unit] = _pending_cursor[unit]   # (2) cursor advanced
    _save()                                              # (3) manifest persisted
    _buffers[unit] = []


def read_rows(unit: str) -> Iterator[Any]:
    """Yield every committed row for ``unit`` across all shards, in order. Use in
    an aggregate/finalize unit to assemble the full result."""
    for shard in _load()["shards"].get(unit, []):
        try:
            with open(shard, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        yield json.loads(line)
        except FileNotFoundError:
            # Shard referenced but not on this fresh workdir — the handler's
            # resume sync-down should have fetched it; skip defensively.
            continue


# --------------------------------------------------------------------------- #
# completion
# --------------------------------------------------------------------------- #
def is_complete() -> bool:
    return _load().get("status") == "complete"


def finalize() -> None:
    """Mark the run complete and persist. After this the script returns normally
    (exit 0), which the harness reads as done."""
    state = _load()
    state["status"] = "complete"
    _save()


def progress() -> Dict[str, Any]:
    """A small progress snapshot for diagnostics/telemetry."""
    state = _load()
    return {"units": state["units"], "done": state["done"],
            "cursor": state["cursor"], "status": state["status"]}