"""The teamlead state file: usage snapshots and the assignment ledger.

State is a hint, not authority. A snapshot records what an agent's budget
looked like when it was measured; it never substitutes for reading the agent's
live status before writing to it. `plan` may run off a stale snapshot on
purpose (planning has no side effects); `apply` always re-checks live status.

Schema (schema_version 2)::

    {
      "schema_version": 2,
      "snapshots":  [ <measure output>, ... ],   # newest last, capped at 20
      "assignments":[ {"schema_version": 2, "at": <ISO-8601>,
                       "role": <str>, "agent": <str>,
                       "status": "applied" | "sent_but_not_started"
                                 | "unknown"}, ... ]
    }

`status` records whether the hand-off was confirmed: `applied` counts toward
an agent's role history, `sent_but_not_started` does not (see
UNCOUNTED_STATUSES), and `unknown` marks a version-1 row migrated without the
information. Version 1 documents and rows carry no `status`; the 1 -> 2
migration below stamps them `unknown`.

Every RECORD carries its own `schema_version`, not just the document: a ledger
row outlives the document it arrived in, and a version on the row is what makes
a later migration auditable row by row.

teamlead owns this file and is the only writer. Writes are atomic: a temp file
in the same directory followed by `os.replace`, so an interrupted run leaves
the previous state intact rather than a truncated file.

Reading follows one rule per direction:

* **Older** document or record: the owner migrates it through MIGRATIONS and
  rewrites the upgraded file. Nobody else migrates.
* **Newer** document or record: this build is the lagging reader, not the
  migrator. It takes the no-usable-prior-state path -- an empty document in
  memory, the file left exactly as found, a warning on stderr.
* **Corrupt** file: no usable prior state, same treatment. Discarding a
  snapshot ring costs one re-measure; overwriting an unread file costs the
  ledger.
"""

import json
import os
import tempfile
from pathlib import Path

from .diagnostics import stderr_warn as _warn
from .errors import StateError

#: The version this build writes, for the document and for every record in it.
#: They move together: one owner, one file, one release train.
STATE_SCHEMA_VERSION = 2

#: An assignment row records what teamlead did, including what did not work.
#: `applied` -- the hand-off was CONFIRMED: the brief appeared in the agent's
#:     transcript, or the agent left idle. Either is enough, and neither
#:     proves the work finished -- see `landed` and `started` on the apply
#:     record for which one it was.
#: `sent_but_not_started` -- the paste went out and NEITHER was observed (see
#:     teamlead/composer.py send_message).
#: `unknown` -- written before rows carried a status.
STATUS_APPLIED = "applied"
STATUS_NOT_STARTED = "sent_but_not_started"
STATUS_UNKNOWN = "unknown"

#: Statuses that do NOT count toward "held this role N times". A hand-off
#: nobody started is not experience, and letting it count would push the next
#: round's tie-break away from an agent that never did the work.
#: Deny-list, not an allow-list: rows migrated from before the field are
#: `unknown`, and those were real hand-offs whose history should not vanish.
UNCOUNTED_STATUSES = frozenset({STATUS_NOT_STARTED})

#: The version a document or record carries when it has no `schema_version` at
#: all -- the pre-versioning shape. Reading an absent key as 0 is what lets the
#: migration chain start below the first stamped version.
UNVERSIONED = 0

MAX_SNAPSHOTS = 20


def default_state_path():
    """`$XDG_STATE_HOME/teamlead/state.json`, falling back to `~/.local/state`."""
    base = os.environ.get("XDG_STATE_HOME")
    root = Path(base) if base else Path.home() / ".local" / "state"
    return root / "teamlead" / "state.json"


def empty_state():
    """A fresh, valid state document."""
    return {"schema_version": STATE_SCHEMA_VERSION, "snapshots": [], "assignments": []}


class _NoUsableState(Exception):
    """Internal signal: this file cannot be read, and must not be written."""


def _migrate_record_0_to_1(record):
    """Pre-versioning assignment row -> version 1: stamp it."""
    record["schema_version"] = 1
    return record


def _migrate_record_1_to_2(record):
    """Assignment row version 1 -> 2: stamp a status.

    Version 1 rows were appended whether or not the agent started, so their
    real outcome is not recoverable -- `unknown` says so rather than claiming
    `applied`. Unknown still counts toward role history: these were genuine
    hand-offs, and the failure the status was added for is the rare case.
    """
    record["schema_version"] = 2
    record.setdefault("status", STATUS_UNKNOWN)
    return record


def _migrate_document_1_to_2(payload):
    """Document version 1 -> 2: carry every assignment row up with it."""
    payload["schema_version"] = 2
    payload["assignments"] = [
        _migrate_record_1_to_2(record) if isinstance(record, dict) else record
        for record in payload.get("assignments", [])
    ]
    return payload


def _migrate_document_0_to_1(payload):
    """Pre-versioning document -> version 1.

    The pre-versioning shape is this one without the stamps, so the upgrade is
    the stamps: the document gains `schema_version`, and every assignment row
    gains its own. The entry earns its place as much for its shape as for its
    work -- adding a real 1->2 step later is one more row in MIGRATIONS, not a
    rewrite of the reader.
    """
    payload["schema_version"] = 1
    payload["assignments"] = [
        _migrate_record_0_to_1(record) if isinstance(record, dict) else record
        for record in payload.get("assignments", [])
    ]
    return payload


#: Document migrations, keyed by the version being upgraded FROM. Each value is
#: (version_produced, upgrade_callable). `_apply_migrations` walks the chain
#: until it reaches STATE_SCHEMA_VERSION, so a future 1->2 is one entry.
MIGRATIONS = {
    UNVERSIONED: (1, _migrate_document_0_to_1),
    1: (2, _migrate_document_1_to_2),
}

#: The same table for one assignment record, walked the same way.
RECORD_MIGRATIONS = {
    UNVERSIONED: (1, _migrate_record_0_to_1),
    1: (2, _migrate_record_1_to_2),
}


def _version_of(payload):
    """The declared version, with an absent key reading as UNVERSIONED."""
    version = payload.get("schema_version", UNVERSIONED)
    if isinstance(version, bool) or not isinstance(version, int):
        raise _NoUsableState("schema_version {!r} is not an integer".format(version))
    return version


def _apply_migrations(payload, table, label):
    """Walk `payload` up the migration chain. Returns (payload, migrated?)."""
    version = _version_of(payload)
    migrated = False
    seen = set()
    while version != STATE_SCHEMA_VERSION:
        if version > STATE_SCHEMA_VERSION:
            raise _NoUsableState(
                "{} is at schema_version {}; this build owns {}".format(
                    label, version, STATE_SCHEMA_VERSION
                )
            )
        if version in seen:
            raise _NoUsableState(
                "{} migration chain loops at version {}".format(label, version)
            )
        seen.add(version)
        step = table.get(version)
        if step is None:
            raise _NoUsableState(
                "{} is at schema_version {}, which no migration upgrades".format(
                    label, version
                )
            )
        produced, upgrade = step
        payload = upgrade(payload)
        version = produced
        migrated = True
    return payload, migrated


def _validate(payload, path):
    """Return (state, migrated?) or raise _NoUsableState.

    Every rejection here is a no-usable-prior-state signal, never an
    instruction to the operator to delete their ledger.
    """
    if not isinstance(payload, dict):
        raise _NoUsableState("the document is not a JSON object")

    payload, migrated = _apply_migrations(payload, MIGRATIONS, "state file")

    for key in ("snapshots", "assignments"):
        if not isinstance(payload.get(key), list):
            raise _NoUsableState("{!r} is not an array".format(key))

    rows = []
    for record in payload["assignments"]:
        if not isinstance(record, dict):
            raise _NoUsableState("an assignment row is not a JSON object")
        record, row_migrated = _apply_migrations(record, RECORD_MIGRATIONS, "an assignment row")
        migrated = migrated or row_migrated
        rows.append(record)
    payload["assignments"] = rows

    # A snapshot is a whole `measure` document and arrives already stamped. One
    # that is not an object, or is stamped ahead of this build, is the same
    # lagging-reader case as the document itself.
    for snapshot in payload["snapshots"]:
        if not isinstance(snapshot, dict):
            raise _NoUsableState("a snapshot entry is not a JSON object")
        found = _version_of(snapshot)
        if found > STATE_SCHEMA_VERSION:
            raise _NoUsableState(
                "a snapshot is at schema_version {}; this build owns {}".format(
                    found, STATE_SCHEMA_VERSION
                )
            )
    return payload, migrated


def load_state(path, warn=None):
    """Read the state file. See `load_state_checked` for the full contract."""
    state, _usable = load_state_checked(path, warn=warn)
    return state


def load_state_checked(path, warn=None):
    """Read the state file, migrating an older one and rewriting it.

    Returns `(state, usable)`. `usable` is False when a file EXISTS and could
    not be read as prior state -- corrupt, or written by a newer build. A
    caller that intends to WRITE must refuse on False: this function promised
    to leave that file exactly as found, and saving over it would destroy the
    ledger it just preserved. A missing file is usable: there is nothing to
    lose, and an empty document is the honest starting point.

    A missing file and a never-written file are the same thing: no prior state.
    So is a corrupt one, and so is one written by a newer build -- in both of
    those the file is left exactly as found, and the caller gets an empty
    document plus a warning on stderr. A tool failure (unreadable permissions,
    a directory in the way) still raises: that is not "no state", it is "this
    environment is wrong", and it carries a fix.
    """
    warn = warn or _warn
    path = Path(path)
    try:
        text = path.read_text(encoding="utf-8")
    except FileNotFoundError:
        return empty_state(), True
    except PermissionError as exc:
        raise StateError(
            "State file {} is not readable: {} - fix its permissions with "
            "`chmod u+rw {}`.".format(path, exc.strerror, path),
            {"path": str(path)},
        ) from None
    except IsADirectoryError:
        raise StateError(
            "State path {} is a directory - point --state at the state.json "
            "file itself.".format(path),
            {"path": str(path)},
        ) from None

    try:
        payload = json.loads(text)
    except json.JSONDecodeError as exc:
        warn(
            "state file {} is not valid JSON ({} at line {} column {}) - starting "
            "from an empty ledger; the file is left untouched.".format(
                path, exc.msg, exc.lineno, exc.colno
            )
        )
        return empty_state(), False

    try:
        state, migrated = _validate(payload, path)
    except _NoUsableState as exc:
        warn(
            "state file {}: {} - starting from an empty ledger; the file is left "
            "untouched.".format(path, exc)
        )
        return empty_state(), False

    if migrated:
        try:
            save_state(path, state)
        except StateError as exc:
            # The upgrade holds in memory even when the rewrite cannot land.
            warn(
                "state file {} migrated to schema_version {} in memory, but the "
                "rewrite failed ({}) - the next run migrates it again.".format(
                    path, STATE_SCHEMA_VERSION, exc
                )
            )
    return state, True


def save_state(path, state):
    """Write `state` atomically, creating the parent directory when needed."""
    path = Path(path)
    try:
        path.parent.mkdir(parents=True, exist_ok=True)
    except (PermissionError, FileExistsError, NotADirectoryError) as exc:
        raise StateError(
            "Cannot create the state directory {}: {} - pass --state to point "
            "at a writable location.".format(path.parent, exc),
            {"path": str(path)},
        ) from None

    handle = None
    temp_name = None
    try:
        fd, temp_name = tempfile.mkstemp(
            prefix=path.name + ".", suffix=".tmp", dir=str(path.parent)
        )
        handle = os.fdopen(fd, "w", encoding="utf-8")
        json.dump(state, handle, indent=2, sort_keys=False)
        handle.write("\n")
        handle.flush()
        os.fsync(handle.fileno())
        handle.close()
        handle = None
        os.replace(temp_name, str(path))
        temp_name = None
    except OSError as exc:
        raise StateError(
            "Cannot write the state file {}: {} - pass --state to point at a "
            "writable location.".format(path, exc),
            {"path": str(path)},
        ) from None
    finally:
        if handle is not None:
            handle.close()
        if temp_name is not None and os.path.exists(temp_name):
            os.unlink(temp_name)
    return path


def add_snapshot(state, snapshot):
    """Append `snapshot` and keep only the newest MAX_SNAPSHOTS entries."""
    snapshots = list(state.get("snapshots", []))
    snapshots.append(snapshot)
    state["snapshots"] = snapshots[-MAX_SNAPSHOTS:]
    return state


def add_assignment(state, at, role, agent, status=STATUS_APPLIED):
    """Append one role-to-agent assignment to the ledger.

    Every hand-off is recorded, including one that never started -- the ledger
    is what teamlead did, and a round that went out and died is exactly the
    thing worth being able to look up afterwards. `status` is what keeps that
    honesty from corrupting the role history: see UNCOUNTED_STATUSES.

    The row carries its own `schema_version`, so a later migration can walk the
    ledger row by row rather than inferring a row's shape from the document.
    """
    state.setdefault("assignments", []).append(
        {
            "schema_version": STATE_SCHEMA_VERSION,
            "at": at,
            "role": role,
            "agent": agent,
            "status": status,
        }
    )
    return state


def latest_snapshot(state):
    """The most recent snapshot, or None when nothing has been measured yet."""
    snapshots = state.get("snapshots", [])
    return snapshots[-1] if snapshots else None


def role_counts(state):
    """How many times each agent has previously held each role.

    Returns ``{role: {agent: count}}``. The planner uses it to break headroom
    ties toward the agent that has held the role least often, which spreads
    roles around instead of pinning one agent to `developer` forever.

    Rows whose status is in UNCOUNTED_STATUSES are skipped: an assignment that
    was sent but never started is not experience of the role, and counting it
    would steer the next round away from the agent that never did the work.
    """
    counts = {}
    for record in state.get("assignments", []):
        if not isinstance(record, dict):
            continue
        if record.get("status") in UNCOUNTED_STATUSES:
            continue
        role = record.get("role")
        agent = record.get("agent")
        if role is None or agent is None:
            continue
        per_role = counts.setdefault(role, {})
        per_role[agent] = per_role.get(agent, 0) + 1
    return counts