"""ECT payload and claim types per draft-nennemann-wimse-ect-01 Section 4."""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Any

# Preferred typ per -01; legacy accepted for backward compatibility.
ECT_TYPE = "exec+jwt"
ECT_TYPE_LEGACY = "wimse-exec+jwt"


def _audience_serialize(aud: list[str]) -> str | list[str]:
    if len(aud) == 1:
        return aud[0]
    return aud


def _audience_deserialize(raw: Any) -> list[str]:
    if isinstance(raw, list):
        return [str(x) for x in raw]
    if isinstance(raw, str):
        return [raw]
    raise ValueError("aud must be string or array of strings")


@dataclass
class Payload:
    """ECT JWT claims per Section 4. Task identity is jti only; no separate tid per spec."""

    iss: str
    aud: list[str]
    iat: int
    exp: int
    jti: str
    exec_act: str
    pred: list[str]  # predecessor jti values (renamed from par in -01)
    wid: str = ""
    inp_hash: str = ""
    out_hash: str = ""
    inp_classification: str = ""
    ext: dict[str, Any] = field(default_factory=dict)

    def to_claims(self) -> dict[str, Any]:
        """Export as JWT claims. Policy and compensation in ext per -01 spec."""
        out: dict[str, Any] = {
            "iss": self.iss,
            "aud": _audience_serialize(self.aud),
            "iat": self.iat,
            "exp": self.exp,
            "jti": self.jti,
            "exec_act": self.exec_act,
            "pred": self.pred,
        }
        if self.wid:
            out["wid"] = self.wid
        if self.inp_hash:
            out["inp_hash"] = self.inp_hash
        if self.out_hash:
            out["out_hash"] = self.out_hash
        if self.inp_classification:
            out["inp_classification"] = self.inp_classification
        if self.ext:
            out["ect_ext"] = dict(self.ext)
        return out

    @classmethod
    def from_claims(cls, claims: dict[str, Any]) -> Payload:
        """Build Payload from JWT claims. Policy claims read from ext per -01 spec."""
        ext = claims.get("ect_ext") or {}
        return cls(
            iss=claims["iss"],
            aud=_audience_deserialize(claims["aud"]),
            iat=int(claims["iat"]),
            exp=int(claims["exp"]),
            jti=claims["jti"],
            exec_act=claims["exec_act"],
            pred=claims.get("pred") or [],
            wid=claims.get("wid", ""),
            inp_hash=claims.get("inp_hash", ""),
            out_hash=claims.get("out_hash", ""),
            inp_classification=claims.get("inp_classification", ""),
            ext=ext,
        )

    def contains_audience(self, verifier_id: str) -> bool:
        return verifier_id in self.aud

    def compensation_required(self) -> bool:
        """Per spec: compensation_required is in ext."""
        if not self.ext:
            return False
        return bool(self.ext.get("compensation_required"))

    def has_policy_claims(self) -> bool:
        """True if both pol and pol_decision are present in ext (per -01, moved to extension)."""
        if not self.ext:
            return False
        return bool(self.ext.get("pol")) and bool(self.ext.get("pol_decision"))

    def pol_decision(self) -> str:
        """Return pol_decision from ext, or empty string."""
        if not self.ext:
            return ""
        return str(self.ext.get("pol_decision", ""))
