"""Delegation chains with capability attenuation."""

from __future__ import annotations

from typing import Any

from .errors import AttenuationViolationError
from .scope import match_scope
from .token import create_token
from .types import (
    APOADefinition,
    APOAToken,
    ChainVerificationResult,
    DelegationDefinition,
    ServiceAuthorization,
    SigningOptions,
)
from .utils import is_expired


def delegate(
    parent_token: APOAToken,
    child_def: DelegationDefinition,
    options: SigningOptions,
) -> APOAToken:
    """Create a delegated (attenuated) token from a parent token.

    - Inherits principal from parent (cannot be overridden)
    - Sets parent_token on child to parent's jti
    - Enforces attenuation rules
    - Additional rules can only be added, not removed
    """
    current_depth = _count_depth(parent_token)

    parent_def = parent_token.definition

    # Merge rules: parent rules + child-only additions
    parent_rules = parent_def.rules or []
    parent_rule_ids = {r.id for r in parent_rules}
    child_extra_rules = [r for r in (child_def.rules or []) if r.id not in parent_rule_ids]
    merged_rules = [*parent_rules, *child_extra_rules]

    # Track delegation depth
    child_depth = current_depth + 1
    child_metadata = {**(child_def.metadata or {}), "_delegationDepth": child_depth}

    # Inherit parent's false constraints into child services BEFORE attenuation
    # check. If parent says { signing: false }, the child MUST carry that
    # constraint even if the delegation definition omits it — otherwise the
    # child's authorize() would skip the check entirely (privilege escalation).
    # _verify_attenuation now treats omission as relaxation, so inheritance
    # must happen before verification or strict callers would fail here.
    inherited_services = []
    for child_svc in child_def.services:
        parent_svc = next((s for s in parent_def.services if s.service == child_svc.service), None)
        if parent_svc and parent_svc.constraints:
            inherited = {k: v for k, v in parent_svc.constraints.items() if v is False}
            if inherited:
                merged_constraints = {**inherited, **(child_svc.constraints or {})}
                inherited_services.append(ServiceAuthorization(
                    service=child_svc.service,
                    scopes=child_svc.scopes,
                    constraints=merged_constraints,
                    access_mode=child_svc.access_mode,
                    browser_config=child_svc.browser_config,
                    api_config=child_svc.api_config,
                ))
            else:
                inherited_services.append(child_svc)
        else:
            inherited_services.append(child_svc)

    # Now verify attenuation against the inherited (filled-in) child definition.
    inherited_child_def = DelegationDefinition(
        agent=child_def.agent,
        services=inherited_services,
        rules=merged_rules or None,
        expires=child_def.expires,
        metadata=child_def.metadata,
    )
    _verify_attenuation(parent_token, inherited_child_def, current_depth)

    full_definition = APOADefinition(
        principal=parent_def.principal,
        agent=child_def.agent,
        agent_provider=parent_def.agent_provider,
        services=inherited_services,
        rules=merged_rules if merged_rules else None,
        expires=child_def.expires or parent_def.expires,
        revocable=parent_def.revocable,
        delegatable=parent_def.delegatable,
        max_delegation_depth=parent_def.max_delegation_depth,
        metadata=child_metadata,
        legal=parent_def.legal,
    )

    # Pass parent_token_id so it's included in the signed JWT payload
    child_token = create_token(full_definition, options, parent_token_id=parent_token.jti)
    return child_token


def verify_chain(
    chain: list[APOAToken],
    revocation_store: Any | None = None,
) -> ChainVerificationResult:
    """Verify a full delegation chain.

    Checks structural integrity (attenuation, expiry, revocation, parent links)
    but does NOT verify cryptographic signatures. Each token MUST be validated
    via validate_token() before passing to verify_chain().
    """
    errors: list[str] = []
    failed_at: int | None = None

    if not chain:
        return ChainVerificationResult(valid=False, depth=0, errors=["Chain is empty"])

    if len(chain) == 1:
        _check_token_validity(chain[0], 0, errors, revocation_store)
        return ChainVerificationResult(
            valid=len(errors) == 0,
            depth=0,
            errors=errors,
            failed_at=0 if errors else None,
            root=chain[0],
            leaf=chain[0],
        )

    for i, token in enumerate(chain):
        errors_before = len(errors)

        _check_token_validity(token, i, errors, revocation_store)

        if i > 0:
            parent = chain[i - 1]
            _check_attenuation(parent, token, i, errors)

            if token.parent_token != parent.jti:
                errors.append(
                    f"Chain link {i}: parentToken '{token.parent_token}' does not match parent jti '{parent.jti}'"
                )

        if failed_at is None and len(errors) > errors_before:
            failed_at = i

    return ChainVerificationResult(
        valid=len(errors) == 0,
        depth=len(chain) - 1,
        errors=errors,
        failed_at=failed_at,
        root=chain[0],
        leaf=chain[-1],
    )


# --- Internal helpers ---


def _count_depth(token: APOAToken) -> int:
    stored = (token.definition.metadata or {}).get("_delegationDepth")
    if isinstance(stored, int):
        return stored
    return 1 if token.parent_token else 0


def _verify_attenuation(
    parent: APOAToken,
    child: DelegationDefinition,
    current_depth: int,
) -> None:
    """Verify attenuation rules. Raises AttenuationViolationError on violation."""
    parent_def = parent.definition
    parent_scopes = [s for svc in parent_def.services for s in svc.scopes]
    child_scopes = [s for svc in child.services for s in svc.scopes]

    # Delegation allowed?
    if parent_def.delegatable is False:
        raise AttenuationViolationError("Parent token does not allow delegation", parent_scopes, child_scopes)

    # Depth limit
    if parent_def.max_delegation_depth is not None:
        if current_depth >= parent_def.max_delegation_depth:
            raise AttenuationViolationError(
                f"Delegation depth {current_depth + 1} exceeds maxDelegationDepth {parent_def.max_delegation_depth}",
                parent_scopes,
                child_scopes,
            )

    # Expiration: child <= parent
    if child.expires is not None:
        from .utils import _to_datetime

        child_exp = _to_datetime(child.expires).timestamp()
        parent_exp = _to_datetime(parent_def.expires).timestamp()
        if child_exp > parent_exp:
            raise AttenuationViolationError("Child token expiration exceeds parent expiration", parent_scopes, child_scopes)

    # Per-service checks
    for child_svc in child.services:
        parent_svc = next((s for s in parent_def.services if s.service == child_svc.service), None)
        if parent_svc is None:
            raise AttenuationViolationError(
                f"Child requests service '{child_svc.service}' not in parent token",
                parent_scopes,
                child_scopes,
            )
        _verify_scope_subset(parent_svc, child_svc)
        _verify_constraints_not_relaxed(parent_svc, child_svc)

    # Rules: child can only add, not remove
    if parent_def.rules:
        parent_rule_ids = [r.id for r in parent_def.rules]
        child_rule_ids = {r.id for r in (child.rules or [])}
        for pid in parent_rule_ids:
            if pid not in child_rule_ids:
                raise AttenuationViolationError(
                    f"Child removes parent rule '{pid}'. Rules can only be added, not removed.",
                    parent_scopes,
                    child_scopes,
                )


def _verify_scope_subset(parent: ServiceAuthorization, child: ServiceAuthorization) -> None:
    parent_scopes = parent.scopes
    for child_scope in child.scopes:
        if not any(match_scope(ps, child_scope) for ps in parent_scopes):
            raise AttenuationViolationError(
                f"Child scope '{child_scope}' on service '{child.service}' is not covered by parent scopes [{', '.join(parent_scopes)}]",
                parent_scopes,
                child.scopes,
            )


def _verify_constraints_not_relaxed(parent: ServiceAuthorization, child: ServiceAuthorization) -> None:
    """A constraint set to False in the parent must remain False in the child.

    Setting it to True flips it (relaxation), and omitting it makes the child's
    authorize() skip the check entirely (silent relaxation). Both fail.
    """
    if not parent.constraints:
        return
    for key, parent_value in parent.constraints.items():
        if parent_value is False:
            child_value = (child.constraints or {}).get(key)
            if child_value is False:
                continue
            raise AttenuationViolationError(
                f"Child relaxes constraint '{key}' on service '{child.service}' (parent: false, child: true)"
                if child_value is True
                else f"Child omits constraint '{key}' on service '{child.service}' (parent: false, child: undefined)",
                parent.scopes,
                child.scopes,
            )


def _check_token_validity(
    token: APOAToken,
    index: int,
    errors: list[str],
    store: Any | None,
) -> None:
    # Expiration (strict, no clock skew)
    if is_expired(token.definition.expires, 0):
        errors.append(f"Chain link {index}: token '{token.jti}' has expired")

    if store is not None:
        record = store.check(token.jti)
        if record:
            errors.append(f"Chain link {index}: token '{token.jti}' has been revoked")


def _check_attenuation(
    parent: APOAToken,
    child: APOAToken,
    index: int,
    errors: list[str],
) -> None:
    for child_svc in child.definition.services:
        parent_svc = next((s for s in parent.definition.services if s.service == child_svc.service), None)
        if parent_svc is None:
            errors.append(f"Chain link {index}: service '{child_svc.service}' not in parent token")
            continue
        for child_scope in child_svc.scopes:
            if not any(match_scope(ps, child_scope) for ps in parent_svc.scopes):
                errors.append(f"Chain link {index}: scope '{child_scope}' on '{child_svc.service}' not covered by parent")

        # Constraints not relaxed: a parent False must remain False. True flips
        # it (relaxation), None makes the child's authorize() skip the check
        # (silent relaxation).
        if parent_svc.constraints:
            for key, parent_value in parent_svc.constraints.items():
                if parent_value is False:
                    child_value = (child_svc.constraints or {}).get(key)
                    if child_value is not False:
                        rendered = "undefined" if child_value is None else repr(child_value)
                        errors.append(
                            f"Chain link {index}: constraint '{key}' on '{child_svc.service}' relaxed by child (parent: false, child: {rendered})"
                        )

    # Child expiration <= parent
    from .utils import _to_datetime

    child_exp = _to_datetime(child.definition.expires).timestamp()
    parent_exp = _to_datetime(parent.definition.expires).timestamp()
    if child_exp > parent_exp:
        errors.append(f"Chain link {index}: child expiration exceeds parent expiration")


def get_delegation_ancestor_ids(input_obj: APOAToken | dict[str, Any]) -> list[str]:
    """Return ancestor token IDs referenced by a token-like input.

    Canonical SDK tokens carry the direct parent jti on ``APOAToken.parent_token``.
    Transport adapters and raw JWT payloads may also carry a ``delegationChain``
    array (snake_case ``delegation_chain`` is also accepted) with ancestor token
    IDs as strings or as ``{"parentTokenId": "..."}`` entries. This helper
    normalizes those forms into a single de-duplicated, order-preserving list so
    revocation checks can consistently include every known ancestor.

    Mirrors ``getDelegationAncestorIds`` in ``@apoa/core``.
    """
    ids: list[str] = []
    seen: set[str] = set()

    def push(value: Any) -> None:
        if isinstance(value, str) and value and value not in seen:
            seen.add(value)
            ids.append(value)

    if isinstance(input_obj, APOAToken):
        push(input_obj.parent_token)
        return ids

    if not isinstance(input_obj, dict):
        return ids

    nested_def = input_obj.get("definition") if isinstance(input_obj.get("definition"), dict) else None
    definition_view = nested_def if nested_def is not None else input_obj

    # Top-level parentToken on the input (camelCase and snake_case both honored).
    push(input_obj.get("parentToken"))
    push(input_obj.get("parent_token"))

    # If a nested definition exists, also look at its parentToken explicitly.
    if nested_def is not None:
        push(nested_def.get("parentToken"))
        push(nested_def.get("parent_token"))

    chain = definition_view.get("delegationChain")
    if chain is None:
        chain = definition_view.get("delegation_chain")

    if isinstance(chain, list):
        for link in chain:
            if isinstance(link, str):
                push(link)
            elif isinstance(link, dict):
                push(link.get("parentTokenId") or link.get("parent_token_id"))

    return ids
