"""Encrypted token storage for platform OAuth integrations."""

from __future__ import annotations

import json
import logging
import os
import re
import secrets
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any

from cryptography.fernet import Fernet, InvalidToken, MultiFernet

_TOKEN_FILE_RE = re.compile(r"^[a-z0-9_-]+$")
_CURRENT_SCHEMA_VERSION = 1
_PREV_KEY_HISTORY_LIMIT = 16
_ROTATION_JOURNAL_VERSION = 1
_ROTATION_JOURNAL_FILENAME = ".token_rotation.journal.json"
_ROTATION_TMP_PREFIX = ".rotate."
logger = logging.getLogger(__name__)


def _utc_now() -> datetime:
    return datetime.now(timezone.utc)


def _parse_iso(value: Any) -> datetime | None:
    if not value:
        return None
    try:
        parsed = datetime.fromisoformat(str(value).replace("Z", "+00:00"))
    except Exception:
        return None
    if parsed.tzinfo is None:
        return parsed.replace(tzinfo=timezone.utc)
    return parsed.astimezone(timezone.utc)


def _iso_utc(value: datetime) -> str:
    dt = value if value.tzinfo else value.replace(tzinfo=timezone.utc)
    return dt.astimezone(timezone.utc).isoformat()


def _coerce_expires_at(raw_expires_at: Any, raw_expires_in: Any, now: datetime) -> str | None:
    expires_at = _parse_iso(raw_expires_at)
    if expires_at is not None:
        return _iso_utc(expires_at)
    if raw_expires_in in (None, ""):
        return None
    try:
        expires_in = int(str(raw_expires_in))
    except Exception:
        return None
    if expires_in <= 0:
        return None
    return _iso_utc(now + timedelta(seconds=expires_in))


def _migrate_to_v1(payload: dict[str, Any]) -> dict[str, Any]:
    migrated = dict(payload)
    tokens = migrated.get("tokens")
    if not isinstance(tokens, dict):
        tokens = {
            key: value
            for key, value in migrated.items()
            if key not in {"_schema_version", "stored_at", "refreshed_at", "expires_at"}
        }
    migrated = {
        "_schema_version": 1,
        "stored_at": migrated.get("stored_at"),
        "refreshed_at": migrated.get("refreshed_at"),
        "expires_at": migrated.get("expires_at"),
        "tokens": dict(tokens),
    }
    return migrated


_MIGRATIONS: dict[int, Any] = {
    0: _migrate_to_v1,
}


class TokenStore:
    """Fernet-encrypted token store backed by files in `{state_dir}/tokens`."""

    def __init__(self, state_dir: Path):
        self.state_dir = Path(state_dir).resolve()
        self.tokens_dir = self.state_dir / "tokens"
        self.tokens_dir.mkdir(parents=True, exist_ok=True)

        self._key_path = self.state_dir / ".token_key"
        self._prev_key_path = self.state_dir / ".token_key.prev"
        self._rotation_journal_path = self.state_dir / _ROTATION_JOURNAL_FILENAME
        self._key, self._key_source, prev_keys = self._resolve_key()
        self._prev_keys = list(prev_keys)
        self._configure_crypto(self._key, prev_keys=prev_keys)
        startup_journal = self._load_rotation_journal()
        self._cleanup_stale_rotation_artifacts(journal=startup_journal)
        if (
            startup_journal is not None
            and str(startup_journal.get("status") or "").lower() == "completed"
        ):
            self._delete_rotation_journal()

    def save_tokens(self, integration: str, tokens: dict[str, Any]) -> None:
        if not isinstance(tokens, dict):
            raise ValueError("tokens payload must be a dict")
        token_data = dict(tokens)
        now = _utc_now()
        stored_at = _parse_iso(token_data.pop("stored_at", None)) or now
        refreshed_at = _parse_iso(token_data.pop("refreshed_at", None)) or now
        expires_at = _coerce_expires_at(
            token_data.pop("expires_at", None),
            token_data.get("expires_in"),
            now,
        )
        payload = {
            "_schema_version": _CURRENT_SCHEMA_VERSION,
            "stored_at": _iso_utc(stored_at),
            "refreshed_at": _iso_utc(refreshed_at),
            "expires_at": expires_at,
            "tokens": token_data,
        }
        self._write_payload(integration, payload)

    def load_tokens(self, integration: str) -> dict[str, Any] | None:
        path = self._token_path(integration)
        if not path.exists():
            return None
        try:
            encrypted = path.read_bytes()
            raw = self._decryptor.decrypt(encrypted)
            parsed = json.loads(raw.decode("utf-8"))
        except InvalidToken:
            logger.warning(
                "Failed to decrypt token payload for integration '%s' at %s; key mismatch or corrupt file.",
                integration,
                path,
            )
            return None
        except OSError:
            logger.warning("Failed to read token payload for integration '%s' at %s.", integration, path)
            return None
        except (ValueError, json.JSONDecodeError):
            logger.warning(
                "Token payload for integration '%s' at %s is not valid JSON.",
                integration,
                path,
            )
            return None
        if not isinstance(parsed, dict):
            return None

        payload = self._apply_migrations(parsed)
        tokens = payload.get("tokens")
        if not isinstance(tokens, dict):
            return None
        result = dict(tokens)
        result["stored_at"] = payload.get("stored_at")
        result["refreshed_at"] = payload.get("refreshed_at")
        result["expires_at"] = payload.get("expires_at")
        result["_schema_version"] = payload.get("_schema_version", _CURRENT_SCHEMA_VERSION)
        return result

    def delete_tokens(self, integration: str) -> bool:
        path = self._token_path(integration)
        if not path.exists():
            return False
        path.unlink()
        return True

    def list_integrations(self) -> list[str]:
        names: list[str] = []
        for path in self.tokens_dir.glob("*.enc"):
            if not path.is_file():
                continue
            names.append(path.stem.lower())
        return sorted(set(names))

    def needs_refresh(self, integration: str, max_age_seconds: int = 5400) -> bool:
        payload = self.load_tokens(integration)
        if payload is None:
            return False
        now = _utc_now()

        expires_at = _parse_iso(payload.get("expires_at"))
        if expires_at is not None:
            issued_at = _parse_iso(payload.get("refreshed_at")) or _parse_iso(payload.get("stored_at"))
            if issued_at is not None and expires_at > issued_at:
                ttl_seconds = (expires_at - issued_at).total_seconds()
                remaining_seconds = (expires_at - now).total_seconds()
                return remaining_seconds <= (ttl_seconds * 0.10)
            return expires_at <= now

        baseline = _parse_iso(payload.get("refreshed_at")) or _parse_iso(payload.get("stored_at"))
        if baseline is None:
            return True
        return (now - baseline).total_seconds() >= max_age_seconds

    def update_access_token(
        self,
        integration: str,
        new_access_token: str,
        issued_at: datetime,
        *,
        expires_in: int | None = None,
    ) -> None:
        payload = self.load_tokens(integration)
        if payload is None:
            raise ValueError(f"No tokens found for integration '{integration}'")
        payload["access_token"] = new_access_token
        payload["refreshed_at"] = _iso_utc(issued_at)
        if expires_in is not None:
            payload["expires_at"] = _iso_utc(issued_at + timedelta(seconds=int(expires_in)))
        self.save_tokens(integration, payload)

    def rotate_key(self, new_key_bytes: bytes | str | None = None) -> int:
        resumed = self._resume_key_rotation()
        if resumed is not None:
            return resumed

        old_key = self._key
        old_prev_keys = list(self._prev_keys)
        if new_key_bytes is None:
            new_key = Fernet.generate_key()
        else:
            new_key = self._validate_key(new_key_bytes)

        rotation_id = secrets.token_hex(8)
        file_entries: list[dict[str, Any]] = []
        for path in sorted(self.tokens_dir.glob("*.enc")):
            if not path.is_file():
                continue
            tmp_path = self._rotation_tmp_path(path, rotation_id)
            file_entries.append(
                {
                    "token_name": path.name,
                    "tmp_name": tmp_path.name,
                    "status": "pending",
                }
            )

        journal: dict[str, Any] = {
            "version": _ROTATION_JOURNAL_VERSION,
            "rotation_id": rotation_id,
            "started_at": _iso_utc(_utc_now()),
            "status": "in_progress",
            "new_key": new_key.decode("utf-8"),
            "old_key": old_key.decode("utf-8"),
            "old_prev_keys": [key.decode("utf-8") for key in old_prev_keys],
            "new_prev_keys": [],
            "key_committed": False,
            "files": file_entries,
        }
        self._write_rotation_journal(journal)
        return self._continue_key_rotation(journal)

    def _resolve_key(self) -> tuple[bytes, str, list[bytes]]:
        env_key = os.getenv("G_GREMLIN_TOKEN_ENCRYPTION_KEY", "").strip()
        if env_key:
            return self._validate_key(env_key), "env", []

        if self._key_path.exists():
            key = self._validate_key(self._key_path.read_text(encoding="utf-8").strip())
            prev_keys = self._read_prev_keys()
            return key, "file", prev_keys

        key = Fernet.generate_key()
        self._write_key_file(self._key_path, key)
        prev_keys = self._read_prev_keys()
        return key, "generated", prev_keys

    def _validate_key(self, raw: bytes | str) -> bytes:
        if isinstance(raw, bytes):
            key = raw.strip()
        else:
            key = raw.strip().encode("utf-8")
        try:
            Fernet(key)
        except Exception as exc:
            raise ValueError("Invalid Fernet key for token store") from exc
        return key

    def _write_key_file(self, path: Path, key: bytes) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp_path = path.parent / f"{path.name}.tmp"
        tmp_path.write_text(key.decode("utf-8"), encoding="utf-8")
        try:
            tmp_path.chmod(0o600)
        except Exception:
            pass
        os.replace(tmp_path, path)

    def _write_prev_keys(self, keys: list[bytes]) -> None:
        normalized = self._normalized_prev_keys(keys, current_key=None)
        payload = {"keys": [key.decode("utf-8") for key in normalized]}
        self._prev_key_path.parent.mkdir(parents=True, exist_ok=True)
        tmp_path = self._prev_key_path.parent / f"{self._prev_key_path.name}.tmp"
        tmp_path.write_text(
            json.dumps(payload, separators=(",", ":"), sort_keys=True),
            encoding="utf-8",
        )
        try:
            tmp_path.chmod(0o600)
        except Exception:
            pass
        os.replace(tmp_path, self._prev_key_path)

    def _token_path(self, integration: str) -> Path:
        name = str(integration or "").strip().lower()
        if not name or not _TOKEN_FILE_RE.match(name):
            raise ValueError(f"Invalid integration name: {integration!r}")
        return self.tokens_dir / f"{name}.enc"

    def _write_payload(self, integration: str, payload: dict[str, Any]) -> None:
        path = self._token_path(integration)
        serialized = json.dumps(payload, separators=(",", ":"), sort_keys=True).encode("utf-8")
        encrypted = self._encryptor.encrypt(serialized)
        tmp_path = path.parent / f"{path.name}.tmp"
        tmp_path.write_bytes(encrypted)
        try:
            tmp_path.chmod(0o600)
        except Exception:
            pass
        os.replace(tmp_path, path)

    def _apply_migrations(self, payload: dict[str, Any]) -> dict[str, Any]:
        try:
            version = int(payload.get("_schema_version", 0))
        except Exception:
            version = 0
        if version > _CURRENT_SCHEMA_VERSION:
            raise ValueError(f"Unsupported token schema version: {version}")

        migrated = dict(payload)
        while version < _CURRENT_SCHEMA_VERSION:
            migrate_fn = _MIGRATIONS.get(version)
            if migrate_fn is None:
                raise ValueError(f"Missing token migration for version {version}")
            migrated = migrate_fn(migrated)
            version = int(migrated.get("_schema_version", version + 1))
        return migrated

    def _read_prev_keys(self) -> list[bytes]:
        if not self._prev_key_path.exists():
            return []
        try:
            raw = self._prev_key_path.read_text(encoding="utf-8").strip()
        except Exception:
            return []

        if not raw:
            return []

        parsed: Any = None
        try:
            parsed = json.loads(raw)
        except Exception:
            parsed = None

        candidates: list[Any]
        if isinstance(parsed, dict) and isinstance(parsed.get("keys"), list):
            candidates = list(parsed.get("keys") or [])
        elif isinstance(parsed, list):
            candidates = list(parsed)
        else:
            # Backward compatibility for legacy plain-text single-key file.
            candidates = [raw]

        validated: list[bytes] = []
        for candidate in candidates:
            if candidate in (None, ""):
                continue
            try:
                validated.append(self._validate_key(candidate))
            except Exception:
                continue
        return self._normalized_prev_keys(validated, current_key=self._key if hasattr(self, "_key") else None)

    def _normalized_prev_keys(self, keys: list[bytes], *, current_key: bytes | None) -> list[bytes]:
        seen: set[bytes] = set()
        ordered: list[bytes] = []
        for key in keys:
            if not isinstance(key, bytes):
                continue
            if current_key is not None and key == current_key:
                continue
            if key in seen:
                continue
            seen.add(key)
            ordered.append(key)
        return ordered[:_PREV_KEY_HISTORY_LIMIT]

    def _configure_crypto(self, key: bytes, *, prev_keys: list[bytes]) -> None:
        current = Fernet(key)
        decryptors: list[Fernet] = [current]
        normalized_prev = self._normalized_prev_keys(prev_keys, current_key=key)
        for prev_key in normalized_prev:
            decryptors.append(Fernet(prev_key))
        self._encryptor = current
        self._decryptor = MultiFernet(decryptors) if len(decryptors) > 1 else current

    def _rotation_tmp_path(self, token_path: Path, rotation_id: str) -> Path:
        safe_rotation_id = re.sub(r"[^a-zA-Z0-9_-]", "", str(rotation_id)) or "rotation"
        return token_path.parent / f"{_ROTATION_TMP_PREFIX}{safe_rotation_id}.{token_path.name}"

    def _resume_key_rotation(self) -> int | None:
        journal = self._load_rotation_journal()
        if journal is None:
            return None

        files = journal.get("files") if isinstance(journal.get("files"), list) else []
        if str(journal.get("status") or "").lower() == "completed":
            self._cleanup_stale_rotation_artifacts(journal=journal)
            self._delete_rotation_journal()
            return 0
        if files and all(str(entry.get("status") or "").lower() in {"applied", "missing"} for entry in files):
            self._cleanup_stale_rotation_artifacts(journal=journal)
            self._delete_rotation_journal()
            return 0
        return self._continue_key_rotation(journal)

    def _continue_key_rotation(self, journal: dict[str, Any]) -> int:
        journal_payload = self._coerce_rotation_journal(journal)
        new_key = self._validate_key(journal_payload["new_key"])
        old_key = self._validate_key(journal_payload["old_key"])

        old_prev_keys: list[bytes] = []
        for candidate in journal_payload.get("old_prev_keys", []):
            try:
                old_prev_keys.append(self._validate_key(candidate))
            except Exception:
                continue

        new_prev_keys_raw = journal_payload.get("new_prev_keys", [])
        new_prev_keys: list[bytes] = []
        for candidate in new_prev_keys_raw:
            try:
                new_prev_keys.append(self._validate_key(candidate))
            except Exception:
                continue
        if not new_prev_keys:
            new_prev_keys = self._normalized_prev_keys([old_key, *old_prev_keys], current_key=new_key)
            journal_payload["new_prev_keys"] = [key.decode("utf-8") for key in new_prev_keys]

        new_fernet = Fernet(new_key)
        rotation_id = str(journal_payload.get("rotation_id") or "")

        # Stage encrypted payloads under a rotation-id-specific temp name.
        for entry in journal_payload["files"]:
            status = str(entry.get("status") or "pending").lower()
            if status in {"applied", "missing"}:
                continue

            token_name = str(entry.get("token_name") or "").strip()
            if not token_name:
                continue
            token_path = self.tokens_dir / token_name
            tmp_name = str(entry.get("tmp_name") or "").strip()
            if not tmp_name:
                tmp_name = self._rotation_tmp_path(token_path, rotation_id).name
                entry["tmp_name"] = tmp_name
            tmp_path = self.tokens_dir / tmp_name

            if status == "staged" and tmp_path.exists():
                continue
            if not token_path.exists():
                entry["status"] = "missing"
                self._write_rotation_journal(journal_payload)
                continue

            decrypted = self._decryptor.decrypt(token_path.read_bytes())
            encrypted = new_fernet.encrypt(decrypted)
            tmp_path.write_bytes(encrypted)
            try:
                tmp_path.chmod(0o600)
            except Exception:
                pass
            entry["status"] = "staged"
            self._write_rotation_journal(journal_payload)

        if not bool(journal_payload.get("key_committed")):
            if self._key_source != "env":
                # Persist a fallback key chain before switching current key.
                # This ensures mixed old/new token payloads remain decryptable
                # after any mid-commit crash.
                self._write_prev_keys(new_prev_keys)
                self._write_key_file(self._key_path, new_key)
            journal_payload["key_committed"] = True
            self._write_rotation_journal(journal_payload)

        self._key = new_key
        self._prev_keys = list(new_prev_keys)
        self._configure_crypto(new_key, prev_keys=new_prev_keys)

        rotated = 0
        for entry in journal_payload["files"]:
            status = str(entry.get("status") or "pending").lower()
            if status in {"applied", "missing"}:
                continue

            token_name = str(entry.get("token_name") or "").strip()
            if not token_name:
                continue
            token_path = self.tokens_dir / token_name
            tmp_name = str(entry.get("tmp_name") or "").strip()
            if not tmp_name:
                tmp_name = self._rotation_tmp_path(token_path, rotation_id).name
                entry["tmp_name"] = tmp_name
            tmp_path = self.tokens_dir / tmp_name

            if not tmp_path.exists():
                if not token_path.exists():
                    entry["status"] = "missing"
                    self._write_rotation_journal(journal_payload)
                    continue
                decrypted = self._decryptor.decrypt(token_path.read_bytes())
                encrypted = new_fernet.encrypt(decrypted)
                tmp_path.write_bytes(encrypted)
                try:
                    tmp_path.chmod(0o600)
                except Exception:
                    pass

            os.replace(tmp_path, token_path)
            rotated += 1
            entry["status"] = "applied"
            self._write_rotation_journal(journal_payload)

        journal_payload["status"] = "completed"
        journal_payload["completed_at"] = _iso_utc(_utc_now())
        self._write_rotation_journal(journal_payload)
        self._cleanup_stale_rotation_artifacts(journal=journal_payload)
        self._delete_rotation_journal()
        return rotated

    def _coerce_rotation_journal(self, payload: dict[str, Any]) -> dict[str, Any]:
        if not isinstance(payload, dict):
            raise ValueError("Rotation journal payload must be an object")

        try:
            version = int(payload.get("version", 0))
        except Exception as exc:
            raise ValueError("Rotation journal version is invalid") from exc
        if version != _ROTATION_JOURNAL_VERSION:
            raise ValueError(f"Unsupported rotation journal version: {version}")

        rotation_id = str(payload.get("rotation_id") or "").strip()
        new_key = str(payload.get("new_key") or "").strip()
        old_key = str(payload.get("old_key") or "").strip()
        if not rotation_id or not new_key or not old_key:
            raise ValueError("Rotation journal missing required key metadata")

        files_raw = payload.get("files")
        if not isinstance(files_raw, list):
            raise ValueError("Rotation journal missing files list")

        normalized_files: list[dict[str, str]] = []
        for entry in files_raw:
            if not isinstance(entry, dict):
                continue
            token_name = str(entry.get("token_name") or "").strip()
            if not token_name:
                continue
            tmp_name = str(entry.get("tmp_name") or "").strip()
            status = str(entry.get("status") or "pending").strip().lower()
            if status not in {"pending", "staged", "applied", "missing"}:
                status = "pending"
            normalized_files.append(
                {
                    "token_name": token_name,
                    "tmp_name": tmp_name,
                    "status": status,
                }
            )

        normalized_old_prev = [
            str(item).strip()
            for item in (payload.get("old_prev_keys") or [])
            if str(item).strip()
        ]
        normalized_new_prev = [
            str(item).strip()
            for item in (payload.get("new_prev_keys") or [])
            if str(item).strip()
        ]

        return {
            "version": _ROTATION_JOURNAL_VERSION,
            "rotation_id": rotation_id,
            "started_at": str(payload.get("started_at") or _iso_utc(_utc_now())),
            "status": str(payload.get("status") or "in_progress").strip().lower(),
            "new_key": new_key,
            "old_key": old_key,
            "old_prev_keys": normalized_old_prev,
            "new_prev_keys": normalized_new_prev,
            "key_committed": bool(payload.get("key_committed")),
            "files": normalized_files,
            "completed_at": payload.get("completed_at"),
        }

    def _write_rotation_journal(self, payload: dict[str, Any]) -> None:
        normalized = self._coerce_rotation_journal(payload)
        self._rotation_journal_path.parent.mkdir(parents=True, exist_ok=True)
        tmp_path = self._rotation_journal_path.parent / f"{self._rotation_journal_path.name}.tmp"
        tmp_path.write_text(
            json.dumps(normalized, separators=(",", ":"), sort_keys=True),
            encoding="utf-8",
        )
        try:
            tmp_path.chmod(0o600)
        except Exception:
            pass
        os.replace(tmp_path, self._rotation_journal_path)
        payload.clear()
        payload.update(normalized)

    def _load_rotation_journal(self) -> dict[str, Any] | None:
        if not self._rotation_journal_path.exists():
            return None
        try:
            raw = self._rotation_journal_path.read_text(encoding="utf-8")
            parsed = json.loads(raw)
            if not isinstance(parsed, dict):
                raise ValueError("Rotation journal must be a JSON object")
            return self._coerce_rotation_journal(parsed)
        except Exception as exc:
            logger.warning("Ignoring invalid token rotation journal at %s: %s", self._rotation_journal_path, exc)
            self._delete_rotation_journal()
            return None

    def _delete_rotation_journal(self) -> None:
        if not self._rotation_journal_path.exists():
            return
        try:
            self._rotation_journal_path.unlink()
        except OSError:
            pass

    def _cleanup_stale_rotation_artifacts(self, journal: dict[str, Any] | None = None) -> None:
        active_tmp_names: set[str] = set()
        if isinstance(journal, dict):
            files = journal.get("files")
            if isinstance(files, list):
                for entry in files:
                    if not isinstance(entry, dict):
                        continue
                    tmp_name = str(entry.get("tmp_name") or "").strip()
                    status = str(entry.get("status") or "").strip().lower()
                    if tmp_name and status in {"pending", "staged"}:
                        active_tmp_names.add(tmp_name)

        if not self.tokens_dir.exists():
            return

        for path in self.tokens_dir.iterdir():
            if not path.is_file():
                continue
            name = path.name
            is_rotation_tmp = name.startswith(_ROTATION_TMP_PREFIX) or name.endswith(".rotate")
            if not is_rotation_tmp:
                continue
            if name in active_tmp_names:
                continue
            try:
                path.unlink()
            except OSError:
                pass


__all__ = ["TokenStore"]
