"""Lab 01 — The resource graph, the admission gate, the reachability prover.

The phase's one question: **how do you make a topology claim you can prove?**

"We configured a private endpoint" and "no traffic leaves the region" are different
statements. The first is a configuration; the second is a property of the whole topology,
and it is only true if *every* path is checked.

Work top to bottom; each section's tests pass once that section is done:

  1.  the resource graph
  2.  plan
  3.  apply
  4.  drift
  5.  the admission gate
  6.  NSG evaluation
  7.  topology construction
  8.  the reachability prover
  9.  GPU scheduling

    pytest                     # against your work
    LAB_MODULE=solution pytest # against the reference

Determinism rules: topological order sorts its frontier, path search sorts its output,
and nothing calls a clock or a network.
"""

from __future__ import annotations

import hashlib
import ipaddress
import json
from collections import deque
from dataclasses import dataclass, field, replace
from enum import Enum
from typing import (Any, Callable, Dict, FrozenSet, Iterable, List, Mapping, Optional,
                    Sequence, Set, Tuple)

# ======================================================================================
# 1. The resource graph
# ======================================================================================


class GraphError(Exception):
    pass


@dataclass(frozen=True)
class Resource:
    """One declared resource. ``address`` is Terraform's ``type.name`` form."""

    address: str
    resource_type: str
    attributes: Mapping[str, Any] = field(default_factory=dict)
    depends_on: Tuple[str, ...] = ()

    def digest(self) -> str:
        raise NotImplementedError


class ResourceGraph:
    """The declared configuration as a DAG.

    Terraform's real contribution is not the language — it is that infrastructure is a
    **graph**, so ordering, parallelism and blast radius all fall out of one structure
    rather than being managed by hand in a script.
    """

    def __init__(self) -> None:
        self._resources: Dict[str, Resource] = {}

    def add(self, resource: Resource) -> Resource:
        """TODO: reject a duplicate address."""
        raise NotImplementedError

    def get(self, address: str) -> Resource:
        """TODO: raise ``GraphError`` naming the unknown address."""
        raise NotImplementedError

    def addresses(self) -> List[str]:
        raise NotImplementedError

    def validate(self) -> None:
        """TODO: every dependency exists, and there is no cycle.

        Both at CONSTRUCTION, not during apply. A cycle discovered halfway through an
        apply leaves the estate in a state neither the config nor the previous state
        describes.
        """
        raise NotImplementedError

    def order(self) -> List[str]:
        """TODO: a DETERMINISTIC topological order — Kahn's algorithm with a **sorted**
        frontier. Raise ``GraphError`` naming the cycle members.

        The sort is what makes it deterministic. Without it, two runs over the same graph
        produce different orders, and a plan that reorders between runs is a plan nobody
        can review.
        """
        raise NotImplementedError

    def transitive_dependents(self, address: str) -> List[str]:
        """TODO: everything downstream, sorted — the blast radius of a change, and the
        set that must not be applied when this resource fails."""
        raise NotImplementedError


# ======================================================================================
# 2. State, plan and apply
# ======================================================================================


class Action(str, Enum):
    NOOP = "noop"
    CREATE = "create"
    UPDATE = "update"
    DELETE = "delete"
    REPLACE = "replace"


#: Attributes that cannot be changed in place. Changing one forces a **replace** —
#: destroy then create — which for a stateful resource is data loss, and it is the single
#: most dangerous thing a plan can silently contain.
FORCE_NEW: Mapping[str, FrozenSet[str]] = {
    "azurerm_subnet": frozenset({"address_prefix", "virtual_network"}),
    "azurerm_kubernetes_cluster": frozenset({"location", "resource_group"}),
    "azurerm_storage_account": frozenset({"location", "account_kind"}),
    "azurerm_private_endpoint": frozenset({"subnet", "target"}),
}


@dataclass(frozen=True)
class Change:
    address: str
    action: Action
    before: Optional[Mapping[str, Any]]
    after: Optional[Mapping[str, Any]]
    changed_attributes: Tuple[str, ...] = ()
    force_new_because: Tuple[str, ...] = ()

    @property
    def destructive(self) -> bool:
        raise NotImplementedError


@dataclass(frozen=True)
class Plan:
    changes: Tuple[Change, ...]

    @property
    def empty(self) -> bool:
        """TODO: true when every change is a NOOP."""
        raise NotImplementedError

    def by_action(self, action: Action) -> List[Change]:
        raise NotImplementedError

    @property
    def destructive_changes(self) -> List[Change]:
        raise NotImplementedError

    def summary(self) -> str:
        """TODO: ``"N to add, N to change, N to replace, N to destroy"``."""
        raise NotImplementedError


class State:
    """What we believe exists. Terraform's state file, in a dict.

    State is the source of most Terraform pain, and the reason is structural: it is a
    *third* thing, alongside the config and reality, and any two of them can disagree.
    """

    def __init__(self) -> None:
        self._resources: Dict[str, Mapping[str, Any]] = {}

    def get(self, address: str) -> Optional[Mapping[str, Any]]:
        raise NotImplementedError

    def set(self, address: str, attributes: Mapping[str, Any]) -> None:
        raise NotImplementedError

    def remove(self, address: str) -> None:
        raise NotImplementedError

    def addresses(self) -> List[str]:
        raise NotImplementedError

    def snapshot(self) -> Dict[str, Mapping[str, Any]]:
        raise NotImplementedError


def diff_attributes(before: Mapping[str, Any], after: Mapping[str, Any]) -> List[str]:
    """TODO: every key whose value differs — added, removed or changed. Sorted."""
    raise NotImplementedError


def plan(graph: ResourceGraph, state: State) -> Plan:
    """TODO: diff desired against believed-actual, in dependency order.

      * not in state          -> CREATE;
      * no attribute changed  -> NOOP;
      * a ``FORCE_NEW`` attribute changed -> **REPLACE**, naming which;
      * otherwise             -> UPDATE, naming the changed attributes;
      * in state, not in the config -> DELETE, in **reverse** order so a dependent goes
        before the thing it depends on.

    Validate the graph first.

    Note what a plan is *not*: it is not a diff against reality. It is a diff against
    **state**, which is why drift detection (§4) is a separate operation — and why a plan
    can be empty while the estate is wrong.
    """
    raise NotImplementedError


@dataclass(frozen=True)
class ApplyResult:
    applied: Tuple[str, ...]
    failed: Tuple[str, ...]
    skipped: Tuple[str, ...]
    errors: Mapping[str, str]

    @property
    def ok(self) -> bool:
        raise NotImplementedError


#: A provider takes (address, action, attributes) and either returns the realized
#: attributes or raises. Injected so a failure is scriptable and the tests deterministic.
Provider = Callable[[str, Action, Optional[Mapping[str, Any]]], Mapping[str, Any]]


def apply(plan_: Plan, graph: ResourceGraph, state: State,
          provider: Provider) -> ApplyResult:
    """TODO: apply in dependency order; a failure **skips every dependent**.

    That skip is the whole safety property. Creating a private endpoint whose subnet
    failed to create produces a resource in an undefined state and a state file that
    disagrees with reality — which is much worse than stopping.

    Write state **per resource, as it succeeds**. A crash mid-apply must leave a state
    file describing what actually got built; that is the only recoverable outcome.

    Handle DELETEs after the creates and updates.
    """
    raise NotImplementedError


# ======================================================================================
# 3. Drift
# ======================================================================================


@dataclass(frozen=True)
class Drift:
    address: str
    attribute: str
    in_state: Any
    in_reality: Any

    def __str__(self) -> str:
        raise NotImplementedError


@dataclass(frozen=True)
class DriftReport:
    drifted: Tuple[Drift, ...]
    missing: Tuple[str, ...]        # in state, not in reality
    unmanaged: Tuple[str, ...]      # in reality, not in state

    @property
    def clean(self) -> bool:
        raise NotImplementedError


def detect_drift(state: State, reality: Mapping[str, Mapping[str, Any]]) -> DriftReport:
    """TODO: compare state against reality, **per attribute**, in three categories.

    Three categories calling for three different responses:

      * **drifted** — somebody changed it out of band. Revert, or absorb it into the
        config; doing neither is how drift accumulates.
      * **missing** — state believes it exists and it does not.
      * **unmanaged** — it exists and nothing manages it. The most dangerous category,
        because it will be destroyed by nothing and audited by nobody.
    """
    raise NotImplementedError


# ======================================================================================
# 4. The admission gate
# ======================================================================================


class Effect(str, Enum):
    ALLOW = "allow"
    DENY = "deny"


@dataclass(frozen=True)
class AdmissionDecision:
    effect: Effect
    address: str
    rule: str
    reason: str

    @property
    def allowed(self) -> bool:
        raise NotImplementedError


#: A rule reads a resource and returns a reason to deny, or None.
AdmissionRule = Callable[[Resource], Optional[str]]


@dataclass(frozen=True)
class Policy:
    name: str
    applies_to: Tuple[str, ...]         # resource types; empty = all
    check: AdmissionRule
    description: str = ""

    def matches(self, resource: Resource) -> bool:
        """TODO: an empty ``applies_to`` matches every type."""
        raise NotImplementedError


def requires_private_endpoint(resource: Resource) -> Optional[str]:
    """TODO: restricted data must not be reachable from the public network.

    **Two** separate conditions: a private endpoint present AND public access disabled.
    A private endpoint alone does not close the public path — the service keeps its
    public FQDN and its public listener, and traffic simply *may* use the private route.
    That distinction is the single most common misconfiguration in this area, and it is
    invisible in a diagram.
    """
    raise NotImplementedError


def forbids_public_ip(resource: Resource) -> Optional[str]:
    raise NotImplementedError


def requires_signed_image(resource: Resource) -> Optional[str]:
    """TODO: an image needs a signature AND a signer in ``ALLOWED_SIGNERS``."""
    raise NotImplementedError


def requires_digest_pin(resource: Resource) -> Optional[str]:
    """TODO: the image reference must contain ``@sha256:``.

    A tag is mutable. ``:latest`` today is not ``:latest`` tomorrow, and the signature
    you verified was for a different artifact.
    """
    raise NotImplementedError


def requires_resource_limits(resource: Resource) -> Optional[str]:
    raise NotImplementedError


def requires_owner_tag(resource: Resource) -> Optional[str]:
    raise NotImplementedError


def forbids_wildcard_egress(resource: Resource) -> Optional[str]:
    raise NotImplementedError


ALLOWED_SIGNERS = frozenset({"platform-ci", "security-team"})


#: TODO: wire the rules above to resource types. The tests pin these rule NAMES:
#:   restricted-data-private-endpoint · no-public-ip · signed-images ·
#:   digest-pinned-images · resource-limits · owner-tag (applies to ALL types) ·
#:   no-wildcard-egress
DEFAULT_POLICIES: Tuple[Policy, ...] = ()


class AdmissionGate:
    """Deny by default. Every deny names its rule."""

    def __init__(self, policies: Sequence[Policy] = DEFAULT_POLICIES) -> None:
        raise NotImplementedError

    def evaluate(self, resource: Resource) -> List[AdmissionDecision]:
        """TODO: every matching policy; collect every denial.

        **A policy that RAISES is a deny.** A policy engine that fails open when its own
        code breaks is a policy engine that will fail open during an incident.

        Return a single ALLOW when nothing objected.
        """
        raise NotImplementedError

    def admit(self, graph: ResourceGraph) -> Tuple[bool, List[AdmissionDecision]]:
        """TODO: evaluate the whole graph; return (ok, every denial)."""
        raise NotImplementedError


# ======================================================================================
# 5. The network topology
# ======================================================================================


@dataclass(frozen=True)
class Subnet:
    name: str
    vnet: str
    cidr: str
    nat_gateway: bool = False
    route_table: Optional[str] = None

    @property
    def network(self) -> ipaddress.IPv4Network:
        raise NotImplementedError


@dataclass(frozen=True)
class VNet:
    name: str
    region: str
    cidr: str


@dataclass(frozen=True)
class Peering:
    """Peerings are **directional** in Azure, and they are not transitive.

    Both facts cause real outages. A peering configured one way carries no traffic; and
    A↔B plus B↔C does *not* give A↔C — you need a hub with routing, or a third peering.
    """

    source: str
    target: str
    allow_forwarded_traffic: bool = False


class Direction(str, Enum):
    INBOUND = "inbound"
    OUTBOUND = "outbound"


@dataclass(frozen=True)
class NsgRule:
    """Priority-ordered, **first-match-wins** — like a firewall, unlike a policy engine.

    A genuine trap for anyone arriving from Phase 09: there, deny beats allow regardless
    of order. Here, a low-priority allow at 100 beats a deny at 200, and the rule set's
    meaning depends entirely on the numbers.
    """

    name: str
    priority: int
    direction: Direction
    action: Effect
    source: str                  # a CIDR, a subnet name, or "*"
    destination: str
    ports: Tuple[int, ...] = ()  # empty = any

    def matches(self, source: str, destination: str, port: int) -> bool:
        """TODO: source, destination and port. ``"*"`` matches anything; a CIDR matches
        an address inside it; an empty port tuple matches any port."""
        raise NotImplementedError


@dataclass(frozen=True)
class PrivateEndpoint:
    """A private endpoint is a NIC in your subnet with a private IP for a PaaS service.

    The half that breaks: **DNS**. Without a private DNS zone linked to the VNet, the
    service's FQDN still resolves to its public IP, so traffic takes the public path and
    everything appears to work — until an auditor asks, or until public access is
    disabled and it stops working with no obvious cause.
    """

    name: str
    subnet: str
    target_service: str
    private_ip: str
    dns_zone_linked: bool = True


@dataclass(frozen=True)
class Service:
    name: str
    region: str
    public_network_access: bool = True
    data_classification: str = "internal"


@dataclass(frozen=True)
class Workload:
    name: str
    subnet: str
    identity: str = ""


@dataclass(frozen=True)
class EgressPolicy:
    subnet: str
    allowed_fqdns: FrozenSet[str] = frozenset()
    allow_internet: bool = False


class Topology:
    """VNets, subnets, peerings, NSGs, private endpoints, services, egress."""

    def __init__(self) -> None:
        self.vnets: Dict[str, VNet] = {}
        self.subnets: Dict[str, Subnet] = {}
        self.peerings: List[Peering] = []
        self.nsg_rules: Dict[str, List[NsgRule]] = {}     # subnet -> rules
        self.private_endpoints: Dict[str, PrivateEndpoint] = {}
        self.services: Dict[str, Service] = {}
        self.workloads: Dict[str, Workload] = {}
        self.egress: Dict[str, EgressPolicy] = {}

    def add_vnet(self, vnet: VNet) -> None:
        raise NotImplementedError

    def add_subnet(self, subnet: Subnet) -> None:
        """TODO: reject a subnet whose VNet does not exist."""
        raise NotImplementedError

    def add_peering(self, peering: Peering) -> None:
        """TODO: reject a peering referencing an unknown VNet."""
        raise NotImplementedError

    def add_nsg_rule(self, rule: NsgRule, subnet: str) -> None:
        raise NotImplementedError

    def add_private_endpoint(self, endpoint: PrivateEndpoint) -> None:
        raise NotImplementedError

    def add_service(self, service: Service) -> None:
        raise NotImplementedError

    def add_workload(self, workload: Workload) -> None:
        raise NotImplementedError

    def set_egress(self, policy: EgressPolicy) -> None:
        raise NotImplementedError

    def peered(self, source_vnet: str) -> List[str]:
        """TODO: DIRECTLY peered targets only. Directional, and NOT transitive."""
        raise NotImplementedError

    def vnet_of(self, subnet: str) -> str:
        raise NotImplementedError

    def region_of_subnet(self, subnet: str) -> str:
        raise NotImplementedError

    def nsg_verdict(self, subnet: str, direction: Direction, source: str,
                    destination: str, port: int) -> Tuple[Effect, str]:
        """TODO: **first match by priority wins.** Return ``(effect, rule_name)``, and
        ``(DENY, "implicit-deny")`` when nothing matches."""
        raise NotImplementedError


# ======================================================================================
# 6. The reachability prover
# ======================================================================================


@dataclass(frozen=True)
class Hop:
    kind: str      # "subnet" | "intra-vnet" | "peering" | "private-endpoint" | "egress" | "service"
    name: str
    detail: str = ""

    def __str__(self) -> str:
        raise NotImplementedError


@dataclass(frozen=True)
class Path:
    hops: Tuple[Hop, ...]

    def render(self) -> str:
        """TODO: ``"subnet:aks -> peering:a->b -> ..."``."""
        raise NotImplementedError

    @property
    def leaves_region(self) -> bool:
        """TODO: true when any hop is an egress hop."""
        raise NotImplementedError


@dataclass(frozen=True)
class ReachabilityResult:
    """A result is not a boolean.

    A control that says "denied" is an opinion. A control that says "denied, and here is
    the path" is a finding somebody can fix — and one that says "allowed, via this path,
    which leaves the region" is the counter-example that makes a residency claim
    falsifiable.
    """

    reachable: bool
    paths: Tuple[Path, ...]
    blocked_by: Tuple[str, ...] = ()

    @property
    def counter_example(self) -> Optional[Path]:
        raise NotImplementedError

    @property
    def any_path_leaves_region(self) -> bool:
        raise NotImplementedError


class ReachabilityProver:
    """Search the topology for every path from a workload to a service.

    The naive check — "is there a private endpoint in my subnet?" — misses three things,
    and each is how a residency claim turns out to be false:

      * the endpoint is in **another subnet of the same VNet** (Azure routes intra-VNet
        implicitly, so this is reachable and the naive check says no);
      * a **peering** into a VNet that has the endpoint, or has internet egress;
      * a **DNS** gap, where the endpoint exists but the FQDN still resolves publicly, so
        traffic takes the public path regardless.
    """

    def __init__(self, topology: Topology, *, max_depth: int = 6) -> None:
        raise NotImplementedError

    def reach(self, workload: str, service: str, *,
              port: int = 443) -> ReachabilityResult:
        """TODO: BFS from the workload's subnet. Raise ``GraphError`` for an unknown
        workload or service. At each subnet, try in order:

        1. **a private endpoint here for the target** — subject to the outbound NSG, and
           only if its DNS zone is linked. An unlinked zone is a ``blocked_by`` reason,
           not a path;
        2. **egress to the public endpoint** — only if the service permits public access,
           the egress policy allows it, and the NSG allows outbound;
        3. **other subnets in the same VNet** — implicit routing, subject to the outbound
           NSG here and the inbound NSG there;
        4. **peered VNets' subnets** — same NSG checks.

        Record a reason in ``blocked_by`` every time a hop is refused; that list is what
        makes an unreachable answer debuggable. Sort the paths shortest-first and the
        reasons deterministically.
        """
        raise NotImplementedError

    def prove_no_egress(self, workload: str, service: str) -> ReachabilityResult:
        """TODO: the residency question — is there ANY path that leaves the region?

        Stated as a proof obligation rather than a configuration check, which is the
        point of the whole phase.
        """
        raise NotImplementedError


# ======================================================================================
# 7. GPU scheduling
# ======================================================================================


@dataclass(frozen=True)
class NodePool:
    """A GPU node pool. Every field here breaks a normal autoscaling assumption.

    ``startup_seconds`` is the one people underestimate: a GPU node is a multi-gigabyte
    image pull plus driver initialization plus engine warm-up. Reactive autoscaling
    against a six-minute start time is not autoscaling, it is a six-minute outage with a
    graph.
    """

    name: str
    gpu_type: str
    gpus_per_node: int
    node_count: int
    max_nodes: int
    startup_seconds: int = 360
    taints: Tuple[str, ...] = ()
    mig_profile: Optional[str] = None       # e.g. "1g.10gb"

    @property
    def slices_per_gpu(self) -> int:
        """TODO: from ``MIG_PROFILES``; 1 when there is no profile or it is unknown.

        MIG partitions one physical GPU into isolated instances. The trade: more,
        smaller, *hardware-isolated* instances — good for many small models, useless for
        one large one, because a slice's memory is a hard ceiling and a 70B model does not
        fit in a 10 GB slice however many you have.
        """
        raise NotImplementedError

    @property
    def total_slices(self) -> int:
        raise NotImplementedError


#: profile -> (instances per GPU, GiB per instance) for an 80 GiB A100/H100.
MIG_PROFILES: Mapping[str, Tuple[int, int]] = {
    "1g.10gb": (7, 10), "2g.20gb": (3, 20), "3g.40gb": (2, 40), "7g.80gb": (1, 80),
}


@dataclass(frozen=True)
class Workload3D:
    name: str
    gpus_required: int
    gpu_type: str
    gang: bool = False              # all-or-nothing (tensor parallelism)
    toleration: Tuple[str, ...] = ()


@dataclass(frozen=True)
class Placement:
    workload: str
    node_pool: Optional[str]
    nodes_used: int
    reason: str

    @property
    def placed(self) -> bool:
        raise NotImplementedError


def schedule(workloads: Sequence[Workload3D],
             pools: Sequence[NodePool]) -> List[Placement]:
    """TODO: place workloads onto pools, honouring taints and **gang scheduling**.

    Order by ``(-gpus_required, name)`` — largest first. A gang of 8 placed after four
    gangs of 2 may find no contiguous capacity, even though the total was sufficient.

    For each workload: no pool with that GPU type -> unplaced, naming the type. Pools
    tainted beyond the workload's tolerations -> unplaced, naming the taint. A **gang**
    needs whole nodes: ``ceil(gpus_required / gpus_per_node)`` of them, all free. A
    non-gang workload just needs the GPUs.

    Always give a reason, placed or not. Return sorted by workload name.

    Gang scheduling is the property that breaks the normal model: a tensor-parallel
    deployment across 8 GPUs needs all 8 *simultaneously*, or it makes no progress at
    all. Kubernetes' default scheduler places pods one at a time and will happily place
    6 of 8 and leave them idle — **holding** the GPUs — while waiting for two more that
    may never arrive. That is a deadlock, and it is why Volcano and Kueue exist.
    """
    raise NotImplementedError


def scale_out_delay(pool: NodePool, *, engine_warmup_seconds: int = 120) -> int:
    """TODO: node start plus engine warm-up.

    This is the number that makes reactive autoscaling wrong for GPU inference and makes
    a **warm pool** the answer.
    """
    raise NotImplementedError


# ======================================================================================
# Worked example
# ======================================================================================


def main() -> None:  # pragma: no cover
    """TODO (optional): once the tests pass, build the seven-section demo.

    Compare against ``python solution.py`` — but only after your own runs.
    """
    raise NotImplementedError


if __name__ == "__main__":  # pragma: no cover
    main()