"""Immutable carry-in for the original JAIF backend budget, never old-state repair."""
from __future__ import annotations

import hashlib
import json
import math
from pathlib import Path

CAP_SECONDS = 28800
ORIGINAL_DEADLINE = 1792677918.181115
FORMAT = "hessian-budget-carry-v1"
RELEASE_FORMAT = "hessian-budget-allocation-release-v1"


def _require(value, message):
    if not value:
        raise ValueError(message)


def _number(value):
    _require(type(value) in (int, float) and math.isfinite(value) and value >= 0,
             "Invalid budget accounting number")
    return value


def _read(path):
    path = Path(path)
    _require(path.is_absolute() and path.resolve() == path and path.is_file()
             and not path.is_symlink(), "Budget evidence must be a regular absolute path")
    from .training import _unique_pairs
    raw = path.read_bytes()
    value = json.loads(raw, object_pairs_hook=_unique_pairs,
                       parse_constant=lambda value: (_ for _ in ()).throw(ValueError(value)))
    _require(isinstance(value, dict), "Budget evidence must be an object")
    return value, hashlib.sha256(raw).hexdigest()


def reference(path):
    """Describe existing evidence only; no creation, mutation or scheduler access."""
    path = Path(path).absolute()
    _, digest = _read(path)
    return {"path": str(path), "sha256": digest}


def _bound(ref):
    _require(isinstance(ref, dict) and set(ref) == {"path", "sha256"}, "Invalid budget evidence reference")
    value, digest = _read(ref["path"])
    _require(digest == ref["sha256"], "Budget evidence hash changed")
    return value


def validate(carry, *, backend):
    """Verify a closed ownership projection and derive a conservative old charge.

    The release receipt is a separately reviewed, hash-bound projection of actual
    Master task/allocation reads. This offline validator does not attest a server.
    It cannot accept an active resource, a missing end time, or a lower charge.
    """
    _require(set(carry) == {"format", "backend", "cap_seconds", "absolute_deadline",
        "carry_in_seconds", "prior_plan", "prior_journal", "release_receipt",
        "controller_binding", "controller_progress", "owner_controller_binding",
        "owner_controller_progress"}, "Invalid budget lineage schema")
    _require(backend in {"nemo", "openenv"} and carry["format"] == FORMAT and carry["backend"] == backend
             and carry["cap_seconds"] == CAP_SECONDS
             and carry["absolute_deadline"] == ORIGINAL_DEADLINE, "Budget cap/deadline/backend changed")
    prior, journal = _bound(carry["prior_plan"]), _bound(carry["prior_journal"])
    release = _bound(carry["release_receipt"])
    binding, progress = _bound(carry["controller_binding"]), _bound(carry["controller_progress"])
    canonical = hashlib.sha256(json.dumps(binding, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
    _require(progress.get("status") == "failed" and progress.get("binding_sha256") == canonical
             and progress.get("deadline") == ORIGINAL_DEADLINE
             and progress.get("started_at") + binding.get("max_wall_seconds") == ORIGINAL_DEADLINE,
             "Original controller deadline/binding changed")
    entry = binding["backends"][backend]
    _require(entry["plan"] == carry["prior_plan"]["path"]
             and entry["plan_sha256"] == carry["prior_plan"]["sha256"]
             and entry["root"] == prior["root"], "Carry-in is not the original backend plan")
    _require(prior.get("max_gpu_wall_seconds") == CAP_SECONDS
             and prior.get("release_eligible") is False
             and (prior.get("backend") == backend if backend == "openenv" else prior.get("backend") is None)
             and carry["prior_plan"]["path"] == str(Path(prior["root"]) / "experiment-plan.json")
             and carry["prior_journal"]["path"] == str(Path(prior["root"]) / "experiment-status.json")
             and journal.get("plan_sha256") == carry["prior_plan"]["sha256"]
             and journal.get("release_eligible") is False
             and journal.get("status") in {"running", "failed", "qualification_completed"},
             "Prior plan/journal is not an unfinished original backend")
    _require(release.get("format") == RELEASE_FORMAT and release.get("backend") == backend
             and release.get("project_id") == 363 and release.get("workspace_id") == 123
             and release.get("experiment_plan_sha256") == carry["prior_plan"]["sha256"]
             and release.get("journal_sha256") == carry["prior_journal"]["sha256"]
             and release.get("released") is True and release.get("active_task_ids") == [],
             "Actual prior allocation release evidence is incomplete")
    experiment = release.get("experiment_id")
    _require(type(experiment) is int and experiment > 0
             and type(release.get("trial_id")) is int and release["trial_id"] > 0,
             "Prior task ownership is missing")
    owner = _bound(carry["owner_controller_binding"])
    owner_progress = _bound(carry["owner_controller_progress"])
    owner_canonical = hashlib.sha256(json.dumps(owner, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
    _require(owner_progress.get("status") in {"failed", "canceled"}
             and owner_progress.get("binding_sha256") == owner_canonical,
             "Prior owning controller is not closed/failed with its original binding")
    owner_entry = owner["backends"][backend]
    name = owner_entry["configuration"]["name"] + "-training"
    child = owner_progress.get("children", {}).get(name, {})
    _require(owner_entry["plan"] == carry["prior_plan"]["path"]
             and owner_entry["plan_sha256"] == carry["prior_plan"]["sha256"]
             and child.get("id") == experiment and child.get("creation_started") is True
             and child.get("stage") == {"backend": backend, "operation": "training", "kind": "experiment", "name": name},
             "Released experiment is not the exact recorded owning-controller child")
    tasks = release.get("tasks")
    _require(isinstance(tasks, list) and tasks, "Prior task release inventory is missing")
    sources = release.get("source_receipts", [])
    _require(isinstance(sources, list) and sources, "Missing underlying release evidence inventory")
    for source in sources:
        _bound(source)
    ends, task_ids, allocation_ids = [], set(), set()
    for task in tasks:
        task_id = task["task_id"]
        end = _number(task["end_time_epoch"])
        _require(isinstance(task_id, str) and task_id.startswith(str(experiment) + ".")
                 and task_id not in task_ids and end > 0, "Invalid released task ownership/end time")
        task_ids.add(task_id)
        allocations = task.get("allocations")
        _require(isinstance(allocations, list) and allocations, "Prior allocation inventory is missing")
        for allocation in allocations:
            identifier = allocation["allocation_id"]
            start, finish = _number(allocation["start_time_epoch"]), _number(allocation["end_time_epoch"])
            _require(isinstance(identifier, str) and identifier.startswith(task_id + ".")
                     and identifier not in allocation_ids and allocation["state"] == "STATE_TERMINATED"
                     and 0 < start <= finish <= end, "An owned allocation is not terminal/released")
            allocation_ids.add(identifier)
            ends.append(finish)
    used = _number(journal["gpu_wall_seconds"])
    active, attempts = journal.get("active_attempt"), journal.get("attempts")
    _require(isinstance(attempts, list) and attempts, "Prior attempt accounting is missing")
    unknown = 0.0
    if active is not None:
        _require(type(active) is int and active == len(attempts) - 1,
                 "Prior unclosed attempt identity is invalid")
        attempt = attempts[active]
        accounted, started = _number(attempt["accounted_through_epoch"]), _number(attempt["started_at_epoch"])
        _require(started <= accounted <= max(ends), "Invalid prior unclosed accounting interval")
        unknown = max(ends) - accounted
    else:
        _require(all("finished_at_epoch" in item for item in attempts), "Prior closed attempts lack end times")
    derived = used + unknown
    _require(_number(carry["carry_in_seconds"]) == derived, "Carry-in differs from conservative released charge")
    return {"carry_in_seconds": derived, "unclosed_seconds": unknown,
            "absolute_deadline": ORIGINAL_DEADLINE, "cap_seconds": CAP_SECONDS}


def load(plan):
    ref = plan.get("budget_carry")
    if ref is None:
        return None
    carry = _bound(ref)
    result = validate(carry, backend=plan.get("backend", "nemo"))
    _require(plan["max_gpu_wall_seconds"] == CAP_SECONDS, "Original backend cap changed")
    return result | {"reference_sha256": ref["sha256"]}


def remaining(plan, state, now):
    available = plan["max_gpu_wall_seconds"] - _number(state["gpu_wall_seconds"])
    lineage = load(plan)
    if lineage is not None:
        _require(state.get("budget_carry") == lineage, "Journal lost its immutable budget carry-in")
        active = state.get("active_attempt")
        _require(active is None or (type(active) is int and active == len(state["attempts"]) - 1),
                 "Carried journal active attempt identity changed")
        charges = [_number(item["charged_gpu_wall_seconds"]) for item in state["attempts"]]
        for item, charged in zip(state["attempts"], charges, strict=True):
            elapsed = _number(item["accounted_through_epoch"]) - _number(item["started_at_epoch"])
            _require(elapsed >= 0 and charged + 1 >= elapsed,
                     "Carried journal attempt accounting decreased/reset")
        _require(state["gpu_wall_seconds"] == lineage["carry_in_seconds"] + sum(charges),
                 "Carried journal cumulative accounting changed/reset")
        available = min(available, lineage["absolute_deadline"] - now)
    if not math.isfinite(available) or available <= 0:
        raise TimeoutError("The original cumulative eight-hour GPU budget or absolute deadline is exhausted")
    return available