"""Deterministic route admission, account concentration and atomic reservations."""
from __future__ import annotations

import hashlib
import json
import math
import sqlite3
import time
from contextlib import contextmanager
from datetime import datetime
from pathlib import Path

ACTIVE = ("reserved", "dispatched", "uncertain")
TERMINAL = ("completed", "refused_before_start", "failed_native_terminal")
ROUTES = {"anthropic": {"claude-seat"}, "openai": {"conductor-route"},
          "xai": {"grok-conductor", "cursor-grok"}, "stealth": {"managed-stealth"}}


def canonical(value):
    return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False)


def digest(value):
    return hashlib.sha256(canonical(value).encode()).hexdigest()


def number(value, lo=0, hi=float("inf")):
    return type(value) in (float, int) and math.isfinite(value) and lo <= value <= hi


def timestamp(value):
    if number(value):
        return float(value)
    if isinstance(value, str):
        try:
            parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
            return parsed.timestamp() if parsed.tzinfo else None
        except ValueError:
            pass
    return None


def fresh(observed, now, max_age=600):
    value = timestamp(observed)
    return value is not None and 0 <= now - value <= max_age


def permitted_model(option, task):
    model, provider = option.get("model"), option.get("provider")
    if option.get("route") not in ROUTES.get(provider, set()):
        return False
    if task.get("fixed_model") and model != task["fixed_model"]:
        return False
    if task.get("origin") == "original_lead":
        return bool(task.get("fixed_model") and task.get("original_session"))
    if task.get("origin") != "worker":
        return False
    return ((provider == "anthropic" and model == "claude-opus-5")
            or (provider == "xai" and isinstance(model, str) and model.startswith("grok-")
                and option.get("highest_native_model") is True)
            or (provider == "stealth" and model == "stealth-alpha"))


def option_errors(option, task, now):
    errors = []
    if not isinstance(option, dict):
        return ["route_option_must_be_object"]
    if not all(isinstance(option.get(k), str) for k in ("id", "account", "pool", "model", "provider", "route")):
        return ["identity_missing"]
    if not isinstance(option.get("capabilities", []), list) or not all(isinstance(x, str) for x in option.get("capabilities", [])):
        return ["capabilities_invalid"]
    if not permitted_model(option, task):
        errors.append("model_or_provider_prohibited")
    if option.get("held") is not False or task.get("held") is not False:
        errors.append("hold_not_clear")
    for field in ("model_verified_at", "identity_verified_at", "billing_verified_at"):
        if not fresh(option.get(field), now):
            errors.append(field + "_stale_or_unknown")
    if option.get("billing_mode") != "included_only":
        errors.append("included_billing_unproved")
    if not set(task.get("required_capabilities", [])) <= set(option.get("capabilities", [])):
        errors.append("capability_mismatch")
    if not all(isinstance(option.get(k), str) and option[k] for k in ("id", "account", "pool", "model")):
        errors.append("identity_missing")
    machine = option.get("machine", {})
    if not isinstance(machine, dict):
        return errors + ["machine_unproved_or_saturated"]
    if (not isinstance(machine.get("observed_work_keys", []), list)
            or not all(isinstance(x, str) for x in machine.get("observed_work_keys", []))):
        return errors + ["machine_reservations_invalid"]
    if (machine.get("state") != "OK" or machine.get("reachable") is not True
            or not fresh(machine.get("observed_at"), now)
            or not number(machine.get("load_per_core"), 0, 1.499999)
            or not number(machine.get("memory_used_percent"), 0, 91.999999)
            or machine.get("fit") != "FITS_FLOOR"
            or not number(machine.get("new_slots"), 1)
            or not isinstance(machine.get("id"), str)):
        errors.append("machine_unproved_or_saturated")
    if option.get("ownership_clear") is not True:
        errors.append("ownership_unproved")
    windows = option.get("windows")
    if not isinstance(windows, list) or not windows:
        errors.append("quota_unknown")
        return errors
    if not all(isinstance(w, dict) and isinstance(w.get("key"), str) for w in windows):
        return errors + ["quota_stale_or_invalid"]
    if len({w.get("key") for w in windows}) != len(windows):
        errors.append("duplicate_quota_window")
    for w in windows:
        if (not isinstance(w.get("key"), str) or not w["key"]
                or not fresh(w.get("observed_at"), now)
                or not number(w.get("remaining"), 0, 100)
                or not number(w.get("duration_seconds"), 1)):
            errors.append("quota_stale_or_invalid")
        reset = timestamp(w.get("reset_at"))
        # Native unused windows can validly report no reset. Never invent one.
        if reset is None and w.get("remaining") != 100:
            errors.append("quota_reset_unknown")
        elif reset is not None and reset <= now:
            errors.append("quota_reset_requires_refresh")
        if not number(w.get("demand"), 0.000001, 100) or not number(w.get("uncertainty"), 0, 100):
            errors.append("checkpoint_demand_unknown")
    return list(dict.fromkeys(errors))


def candidates(snapshot, reservations=(), policy="concentrate", advised_model=None, *, now=None):
    if not isinstance(snapshot, dict) or snapshot.get("schema") != "jev-route-input/v1":
        raise ValueError("unsupported snapshot schema")
    now = time.time() if now is None else now
    if not fresh(snapshot.get("observed_at"), now):
        raise ValueError("snapshot timestamp stale or missing")
    task = snapshot.get("task")
    if not isinstance(task, dict):
        raise ValueError("task object required")
    if not all(isinstance(task.get(k), str) and task[k] for k in ("work_id", "segment_id")):
        raise ValueError("stable work and segment identity required")
    if task.get("origin") not in ("worker", "original_lead"):
        raise ValueError("task origin must be worker or original_lead")
    if task.get("held") is not False:
        raise ValueError("task hold must be explicitly clear")
    if not isinstance(task.get("required_capabilities", []), list) or not all(isinstance(x, str) for x in task.get("required_capabilities", [])):
        raise ValueError("required_capabilities must be a list")
    options = snapshot.get("options")
    if not isinstance(options, list) or not options:
        raise ValueError("at least one route option is required")
    if policy not in ("concentrate", "baseline"):
        raise ValueError("unknown policy")
    option_failures = [option_errors(option, task, now) for option in options]
    # Build shared meters only from structurally valid observations. Invalid
    # aliases of a pool fail closed for that pool; no optimistic alias bypass.
    pool_errors, pool_windows, signatures, ids = {}, {}, {}, {}
    for option, errors in zip(options, option_failures):
        if not isinstance(option, dict):
            continue
        route_id, pool = option.get("id"), option.get("pool")
        if isinstance(route_id, str):
            ids[route_id] = ids.get(route_id, 0) + 1
        if not isinstance(pool, str):
            continue
        windows = option.get("windows")
        quota_errors = [e for e in errors if e.startswith(("quota", "duplicate_quota", "checkpoint_demand"))]
        if quota_errors:
            pool_errors[pool] = "shared_pool_quota_unproved"
            continue
        if not isinstance(windows, list) or not windows:
            continue
        if not all(isinstance(w, dict) and isinstance(w.get("key"), str)
                   and number(w.get("duration_seconds"), 1)
                   and number(w.get("remaining"), 0, 100)
                   and number(w.get("demand"), 0.000001, 100)
                   and number(w.get("uncertainty"), 0, 100)
                   and timestamp(w.get("observed_at")) is not None for w in windows):
            pool_errors[pool] = "shared_pool_quota_unproved"
            continue
        model_key = (pool, option.get("model") if isinstance(option.get("model"), str) else "")
        signature = sorted((w["key"], w["duration_seconds"], timestamp(w.get("reset_at"))) for w in windows)
        if model_key in signatures and signature != signatures[model_key]:
            pool_errors[pool] = "shared_pool_windows_conflict"
        signatures[model_key] = signature
        for w in windows:
            pool_windows.setdefault((pool, w["key"]), []).append(w)
    held, active, machine_claims = {}, set(), {}
    for r in reservations:
        if not isinstance(r, dict) or not isinstance(r.get("demand"), dict):
            raise ValueError("reservation demand must be an object")
        # Completed inference can reach the ledger before the provider meter
        # refreshes. Keep that debit until a later native observation exists.
        debit_pending = r.get("state") in ("completed", "failed_native_terminal")
        if r.get("state") not in ACTIVE and not debit_pending:
            continue
        if not isinstance(r.get("pool"), str) or not isinstance(r.get("machine"), str):
            raise ValueError("active reservation identity required")
        if r.get("state") in ACTIVE:
            active.add(r["pool"])
            # Native telemetry can explicitly identify the running reservation.
            # Otherwise conservatively reserve one slot until reconciled.
            observed = any(isinstance(o, dict) and isinstance(o.get("machine"), dict)
                           and o["machine"].get("id") == r["machine"]
                           and r.get("work_key") is not None
                           and isinstance(o["machine"].get("observed_work_keys", []), list)
                           and r["work_key"] in o["machine"].get("observed_work_keys", []) for o in options)
            if not observed:
                machine_claims[r["machine"]] = machine_claims.get(r["machine"], 0) + 1
        for k, amount in r["demand"].items():
            if not isinstance(k, str) or not number(amount, 0):
                raise ValueError("reservation demand must be finite and non-negative")
            observations = pool_windows.get((r["pool"], k), [])
            reconciled = timestamp(r.get("reconciled_at"))
            if debit_pending and reconciled is not None and observations and all(timestamp(w["observed_at"]) > reconciled for w in observations):
                continue
            held[(r["pool"], k)] = held.get((r["pool"], k), 0) + amount
    admitted, denied = [], []
    for option, errors in zip(options, option_failures):
        if not isinstance(option, dict):
            denied.append({"id": None, "reasons": ["route_option_must_be_object"]})
            continue
        if isinstance(option.get("id"), str) and ids.get(option["id"], 0) > 1:
            errors.append("duplicate_route_identity")
        if isinstance(option.get("pool"), str) and option["pool"] in pool_errors:
            errors.append(pool_errors[option["pool"]])
        if errors:
            denied.append({"id": option.get("id"), "reasons": errors})
            continue
        machine = option["machine"]
        # An account may have several route aliases (for example, separate
        # machines or launchers). Treat the pool as one quota identity and use
        # the lowest observed remaining value across aliases. This prevents a
        # stale or optimistic alias from bypassing a stricter observation.
        remaining = {
            w["key"]: min(x["remaining"] for x in pool_windows[(option["pool"], w["key"])])
                         - held.get((option["pool"], w["key"]), 0)
            for w in option["windows"]
        }
        demand = {
            w["key"]: max(x["demand"] + x["uncertainty"]
                          for x in pool_windows[(option["pool"], w["key"])])
            for w in option["windows"]
        }
        if any(remaining[k] + 1e-9 < need for k, need in demand.items()):
            denied.append({"id": option["id"], "reasons": ["checkpoint_does_not_fit"]})
            continue
        slots = machine["new_slots"] - machine_claims.get(machine["id"], 0)
        if slots < 1:
            denied.append({"id": option["id"], "reasons": ["machine_slots_reserved"]})
            continue
        longest = max(w["duration_seconds"] for w in option["windows"])
        long_remaining = min(remaining[w["key"]] for w in option["windows"] if w["duration_seconds"] == longest)
        usable = min(remaining.values())
        resets = [timestamp(w.get("reset_at")) for w in option["windows"]]
        reset = min((x for x in resets if x is not None), default=float("inf"))
        model_rank = 0 if task.get("fixed_model") or option["model"] == advised_model else 1
        model_priority = option.get("model_priority", 100)
        if not number(model_priority):
            model_priority = 100
        physical = (machine["load_per_core"] >= 1, -slots, machine["load_per_core"], machine["id"])
        if policy == "baseline":
            key = (model_rank, model_priority, -usable, reset, physical, option["id"])
        else:
            key = (model_rank, model_priority, long_remaining >= 100,
                   option["pool"] not in active, long_remaining, usable, reset, physical, option["id"])
        admitted.append({"option": option, "key": key, "remaining": remaining, "demand": demand,
                         "long_remaining": long_remaining, "usable": usable,
                         "active_pool": option["pool"] in active})
    admitted.sort(key=lambda x: x["key"])
    selected = admitted[0] if admitted else None
    return {"schema": "jev-route-decision/v1", "policy": policy, "input_sha256": digest(snapshot),
            "selected": None if selected is None else {
                "option_id": selected["option"]["id"], "model": selected["option"]["model"],
                "route": selected["option"]["route"], "account": selected["option"]["account"],
                "pool": selected["option"]["pool"], "machine": selected["option"]["machine"]["id"],
                "demand": selected["demand"], "remaining_after_reservations": selected["remaining"],
                "prepare_successor": min(selected["remaining"][k] - need
                                         for k, need in selected["demand"].items()) <= 10,
                "reason": "reuse_active_fitting_pool" if selected["active_pool"] else "best_fitting_eligible_pool"},
            "eligible_count": len(admitted), "denied": denied,
            "advised_model": advised_model, "billing_mutations": 0, "execution_authorized": False}


class Ledger:
    """Single route-host ledger. Remote callers use the same host, never replicas."""
    def __init__(self, path):
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with self.connect() as db:
            db.execute("CREATE TABLE IF NOT EXISTS reservations (work_key TEXT PRIMARY KEY, task_hash TEXT NOT NULL, state TEXT NOT NULL, record TEXT NOT NULL)")
        self.path.chmod(0o600)

    @contextmanager
    def connect(self):
        db = sqlite3.connect(self.path, timeout=10)
        db.execute("PRAGMA synchronous=FULL")
        try:
            with db:
                yield db
        finally:
            db.close()

    def select(self, snapshot, *, claim=False, policy="concentrate", advised_model=None):
        task = snapshot.get("task") if isinstance(snapshot, dict) else None
        if not isinstance(task, dict) or not all(isinstance(task.get(k), str) and task[k] for k in ("work_id", "segment_id")):
            raise ValueError("stable work and segment identity required")
        key = digest([task.get("work_id"), task.get("segment_id")])
        task_hash = digest(task)
        with self.connect() as db:
            db.execute("BEGIN IMMEDIATE")
            old = db.execute("SELECT task_hash,record FROM reservations WHERE work_key=?", (key,)).fetchone()
            if old:
                if old[0] != task_hash:
                    raise ValueError("work identity reused with changed task; reconcile before another launch")
                record = json.loads(old[1])
                return {"idempotent": True, "reservation": record, "execution_authorized": False}
            rows = [json.loads(x[0]) for x in db.execute("SELECT record FROM reservations WHERE state != 'refused_before_start'")]
            result = candidates(snapshot, rows, policy, advised_model)
            if claim and result["selected"]:
                record = {**result["selected"], "work_id": task["work_id"], "segment_id": task["segment_id"],
                          "state": "reserved", "created_at": time.time(), "input_sha256": result["input_sha256"],
                          "work_key": key, "native_receipt": None}
                db.execute("INSERT INTO reservations VALUES (?,?,?,?)", (key, task_hash, "reserved", canonical(record)))
                result["reservation"] = record
            return result

    def reconcile(self, work_key, state, receipt):
        if state not in (*TERMINAL, "dispatched", "uncertain"):
            raise ValueError("invalid native state")
        if not isinstance(receipt, dict) or not receipt.get("source_sha256") or not receipt.get("source_path"):
            raise ValueError("native receipt evidence required")
        source = Path(receipt["source_path"])
        if hashlib.sha256(source.read_bytes()).hexdigest() != receipt["source_sha256"]:
            raise ValueError("native receipt source changed")
        with self.connect() as db:
            db.execute("BEGIN IMMEDIATE")
            row = db.execute("SELECT record FROM reservations WHERE work_key=?", (work_key,)).fetchone()
            if not row:
                raise ValueError("unknown reservation")
            record = json.loads(row[0])
            if receipt.get("work_id") != record["work_id"] or receipt.get("segment_id") != record["segment_id"]:
                raise ValueError("native receipt work identity mismatch")
            if record["state"] in TERMINAL:
                if state != record["state"]:
                    raise ValueError("terminal reservation cannot reopen")
                return record
            if state in TERMINAL:
                # Bind terminal release to the hashed native export itself.
                # A caller-supplied boolean cannot turn a transport error into
                # proof that the provider stopped or never started the work.
                try:
                    native = json.loads(source.read_text())
                except (ValueError, UnicodeError) as error:
                    raise ValueError("native terminal export required") from error
                identity = ("work_id", "segment_id", "account", "machine", "model", "route")
                if (not isinstance(native, dict) or native.get("schema") != "jev-native-receipt/v1"
                        or any(native.get(k) != record.get(k) for k in identity)
                        or native.get("state") != state
                        or not isinstance(native.get("native_id"), str) or not native["native_id"]
                        or native.get("terminal_observed") is not True
                        or timestamp(native.get("observed_at")) is None
                        or native.get("started") is not (state != "refused_before_start")):
                    raise ValueError("uncertain outcomes must remain reserved")
            record.update(state=state, native_receipt=receipt, reconciled_at=time.time())
            db.execute("UPDATE reservations SET state=?,record=? WHERE work_key=?", (state, canonical(record), work_key))
            return record