"""
ID Wispera Session Token Manager
Manages short-lived session tokens for headless / CI environments.

Sidecar file schema (cross-SDK compatible — TS, Python, Go all read/write this):
  {
    "version": 1,
    "entries": [
      {
        "tokenHash":            "<SHA-256 hex>",
        "salt":                 "<32 bytes hex>",
        "iv":                   "<12 bytes hex>",
        "encryptedPassphrase":  "<AES-256-GCM ciphertext, base64url>",
        "createdAt":            "<ISO 8601>",
        "expiresAt":            "<ISO 8601 | null>",
        "label":                "<user-provided label>"
      }
    ]
  }
"""

import base64
import hashlib
import json
import os
import secrets
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes

# ──────────────────────────────────────────────────────────────────────────────
# Constants
# ──────────────────────────────────────────────────────────────────────────────

TOKEN_PREFIX = "idw_token_"
TOKEN_RANDOM_BYTES = 64
SIDECAR_VERSION = 1
PBKDF2_ITERATIONS = 100_000
SALT_BYTES = 32
IV_BYTES = 12
SESSION_TOKEN_ENV_VAR = "IDW_SESSION_TOKEN"

# ──────────────────────────────────────────────────────────────────────────────
# Types
# ──────────────────────────────────────────────────────────────────────────────


@dataclass
class SidecarEntry:
    token_hash: str
    salt: str
    iv: str
    encrypted_passphrase: str
    created_at: str
    expires_at: Optional[str]
    label: str


@dataclass
class TokenInfo:
    token_hash: str
    label: str
    created_at: str
    expires_at: Optional[str]


# ──────────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────────


def _sha256_hex(data: str) -> str:
    return hashlib.sha256(data.encode("utf-8")).hexdigest()


def _base64url_encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _base64url_decode(data: str) -> bytes:
    # Pad back to a multiple of 4
    padded = data + "=" * (-len(data) % 4)
    return base64.urlsafe_b64decode(padded)


def _derive_key(token: str, salt: bytes) -> bytes:
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=PBKDF2_ITERATIONS,
    )
    return kdf.derive(token.encode("utf-8"))


def _encrypt_passphrase(passphrase: str, token: str, salt: bytes, iv: bytes) -> str:
    key = _derive_key(token, salt)
    aesgcm = AESGCM(key)
    ciphertext = aesgcm.encrypt(iv, passphrase.encode("utf-8"), None)
    return _base64url_encode(ciphertext)


def _decrypt_passphrase(encrypted_b64: str, token: str, salt: bytes, iv: bytes) -> str:
    key = _derive_key(token, salt)
    aesgcm = AESGCM(key)
    plaintext = aesgcm.decrypt(iv, _base64url_decode(encrypted_b64), None)
    return plaintext.decode("utf-8")


# ──────────────────────────────────────────────────────────────────────────────
# Manager
# ──────────────────────────────────────────────────────────────────────────────


class SessionTokenManager:
    """Manage session token entries in the sidecar file."""

    def __init__(self, sidecar_path: Optional[str] = None) -> None:
        if sidecar_path is None:
            sidecar_path = os.path.join(
                os.path.expanduser("~"), ".id-wispera", "sessions.json"
            )
        self._path = Path(sidecar_path)

    # ── sidecar I/O ──────────────────────────────────────────────────────

    def _read_sidecar(self) -> dict:
        try:
            data = json.loads(self._path.read_text("utf-8"))
            if data.get("version") != SIDECAR_VERSION or not isinstance(
                data.get("entries"), list
            ):
                return {"version": SIDECAR_VERSION, "entries": []}
            return data
        except (FileNotFoundError, json.JSONDecodeError, OSError):
            return {"version": SIDECAR_VERSION, "entries": []}

    def _write_sidecar(self, data: dict) -> None:
        self._path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self._path.with_suffix(".json.tmp")
        tmp.write_text(json.dumps(data, indent=2), "utf-8")
        # Restrictive permissions (owner read/write only)
        tmp.chmod(0o600)
        tmp.replace(self._path)

    # ── public API ───────────────────────────────────────────────────────

    @staticmethod
    def generate_token() -> str:
        """Generate a new CSPRNG token string."""
        random_bytes = secrets.token_bytes(TOKEN_RANDOM_BYTES)
        return TOKEN_PREFIX + _base64url_encode(random_bytes)

    def create(
        self,
        passphrase: str,
        *,
        label: str = "default",
        expires_in_seconds: Optional[int] = None,
    ) -> str:
        """
        Create a session token entry.

        Returns the plaintext token (shown to the user once — never stored).
        """
        token = self.generate_token()
        token_hash = _sha256_hex(token)
        salt = secrets.token_bytes(SALT_BYTES)
        iv = secrets.token_bytes(IV_BYTES)
        encrypted = _encrypt_passphrase(passphrase, token, salt, iv)

        now = datetime.now(timezone.utc)
        expires_at: Optional[str] = None
        if expires_in_seconds is not None:
            from datetime import timedelta

            expires_at = (now + timedelta(seconds=expires_in_seconds)).isoformat()

        entry = {
            "tokenHash": token_hash,
            "salt": salt.hex(),
            "iv": iv.hex(),
            "encryptedPassphrase": encrypted,
            "createdAt": now.isoformat(),
            "expiresAt": expires_at,
            "label": label,
        }

        sidecar = self._read_sidecar()
        sidecar["entries"].append(entry)
        self._write_sidecar(sidecar)

        return token

    def resolve(self, token: str) -> str:
        """
        Resolve a session token to the vault passphrase.

        Raises:
            ValueError: If the token is invalid, expired, or not found.
        """
        if not token.startswith(TOKEN_PREFIX):
            raise ValueError("Invalid session token format.")

        token_hash = _sha256_hex(token)
        sidecar = self._read_sidecar()

        entry = None
        for e in sidecar["entries"]:
            if e["tokenHash"] == token_hash:
                entry = e
                break

        if entry is None:
            raise ValueError(
                "Session token not found. It may have been revoked or the sidecar "
                "file is missing. Use your passphrase directly, or create a new "
                "token with 'idw auth token create'."
            )

        # Check expiry
        if entry.get("expiresAt"):
            expires = datetime.fromisoformat(entry["expiresAt"])
            if expires < datetime.now(timezone.utc):
                raise ValueError(
                    "Session token has expired. "
                    "Create a new token with 'idw auth token create'."
                )

        try:
            return _decrypt_passphrase(
                entry["encryptedPassphrase"],
                token,
                bytes.fromhex(entry["salt"]),
                bytes.fromhex(entry["iv"]),
            )
        except Exception as exc:
            raise ValueError(
                "Failed to decrypt passphrase from session token. The sidecar "
                "entry may be corrupt. Use your passphrase directly, or create "
                "a new token with 'idw auth token create'."
            ) from exc

    def revoke(self, hash_prefix: str) -> bool:
        """Revoke a session token by its hash prefix (first 8+ hex chars)."""
        sidecar = self._read_sidecar()
        before = len(sidecar["entries"])
        sidecar["entries"] = [
            e
            for e in sidecar["entries"]
            if not e["tokenHash"].startswith(hash_prefix)
        ]
        if len(sidecar["entries"]) == before:
            return False
        self._write_sidecar(sidecar)
        return True

    def list(self) -> list[TokenInfo]:
        """List all session tokens (metadata only — never returns secrets)."""
        sidecar = self._read_sidecar()
        return [
            TokenInfo(
                token_hash=e["tokenHash"],
                label=e.get("label", ""),
                created_at=e["createdAt"],
                expires_at=e.get("expiresAt"),
            )
            for e in sidecar["entries"]
        ]

    def prune_expired(self) -> int:
        """Remove all expired entries. Returns the number removed."""
        sidecar = self._read_sidecar()
        before = len(sidecar["entries"])
        now = datetime.now(timezone.utc)
        sidecar["entries"] = [
            e
            for e in sidecar["entries"]
            if not e.get("expiresAt")
            or datetime.fromisoformat(e["expiresAt"]) >= now
        ]
        removed = before - len(sidecar["entries"])
        if removed > 0:
            self._write_sidecar(sidecar)
        return removed

    @staticmethod
    def get_token_from_env() -> Optional[str]:
        """Read the token from the environment variable."""
        return os.environ.get(SESSION_TOKEN_ENV_VAR)
