"""Cloud Locker — per-user S3 storage with local fallback.

Index architecture
------------------
The session index is stored in two layers:

1. ``index.json`` — compacted manifest of all entries, rebuilt lazily.
2. ``_delta/{session_id}.json`` — one tiny object written per ``persist()``
   call, with no read step required.

On ``persist()``: upload session + write one ~200-byte delta.  No download
of the full index needed.

On ``load_index()``: download ``index.json``, list the ``_delta/`` prefix,
fetch any pending deltas, merge, compact (re-upload ``index.json`` and delete
the delta objects), then return.

This makes writes O(1) in data transferred.  Reads are a single GET when no
deltas are pending; O(N) GETs only after N sessions have been written since
the last read.  For a research CLI this is the right trade-off.

Serialization
-------------
All JSON encode/decode uses msgspec for 3-5× faster throughput compared to
the stdlib ``json`` module.
"""

from __future__ import annotations

import contextlib
from pathlib import Path
from typing import Any

import msgspec
import msgspec.json
from botocore.exceptions import BotoCoreError, ClientError
from filelock import FileLock, Timeout

from alice_cli.aws.session import get_session
from alice_cli.console import err_console
from alice_cli.constants import FILELOCK_TIMEOUT
from alice_cli.errors import LockerError
from alice_cli.logging import get_logger
from alice_cli.models import CLIConfig, IndexEntry, LockerIndex, SessionRecord, TokenUsage
from alice_cli.retry import with_s3_retry
from alice_cli.session_record import build_index_entry, deserialize_record, serialize_record

log = get_logger(__name__)

_LOCKER_DIR = Path.home() / ".alice" / "locker"

# Mapping from session type to the S3 category prefix used when storing records.
# Used by backfill_tokens() to skip the O(C) category scan.
_TYPE_TO_CATEGORY: dict[str, str] = {
    "chat": "chat",
    "invoke": "invoke",
    "batch": "batch",
    "compare": "compare",
    "dialog": "dialog",
    "summarize": "invoke",  # summarize records are stored under "invoke"
    "cite-process": "cite",
    "cite-lookup": "cite",
}

# When this many delta objects accumulate, trigger proactive compaction on the
# next persist() call to prevent an O(N) read spike.
_MAX_PENDING_DELTAS = 50


class CloudLocker:
    """Per-user S3 storage with local fallback."""

    def __init__(self, config: CLIConfig, username: str, bucket: str) -> None:
        self._config = config
        self._username = username
        self._bucket = bucket
        self._namespace = config.namespace
        self._environment = config.environment

        # Build the boto3 S3 client.  The Bedrock API keys are scoped to
        # model invocation and do NOT have S3 access, so prefer the SSO
        # profile for bucket operations.  Only fall back to API keys when
        # no profile is available at all.
        session = get_session(
            profile=config.profile,
            region=config.region,
            access_key=config.bedrock_access_key if config.bedrock_secret_key else None,
            secret_key=config.bedrock_secret_key,
        )
        self._s3 = session.client("s3")

    # ── S3 operations with retry ──────────────────────────────────────

    @with_s3_retry
    def _put_object_with_retry(self, **kwargs: Any) -> Any:
        """Put an object to S3 with automatic retry on transient errors."""
        return self._s3.put_object(**kwargs)

    @with_s3_retry
    def _get_object_with_retry(self, **kwargs: Any) -> Any:
        """Get an object from S3 with automatic retry on transient errors."""
        return self._s3.get_object(**kwargs)

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

    def persist(self, record: SessionRecord, category: str) -> None:
        """Write session object + append index delta. Falls back to local on S3 failure.

        The index is updated by writing a single small delta file rather than
        downloading and re-uploading the full index.  This makes writes O(1).
        """
        log.debug(
            "persist_starting",
            session_id=record.id,
            category=category,
            username=self._username,
            bucket=self._bucket,
        )
        body = serialize_record(record).encode("utf-8")
        s3_key = self._s3_key(category, record.id)

        try:
            # 1. Upload the session object
            self._put_object_with_retry(
                Bucket=self._bucket,
                Key=s3_key,
                Body=body,
                ContentType="application/json",
                ServerSideEncryption="AES256",
            )
            log.info(
                "session_persisted_s3",
                session_id=record.id,
                category=category,
                s3_key=s3_key,
                size_bytes=len(body),
            )

            # 2. Write a single tiny delta entry — no GET of the full index.
            # If too many deltas have accumulated, trigger compaction first to
            # prevent a large O(N) read spike on the next load_index() call.
            entry = build_index_entry(record)
            pending_count = len(self._list_delta_keys())
            if pending_count >= _MAX_PENDING_DELTAS:
                log.info(
                    "proactive_compaction_triggered",
                    pending_deltas=pending_count,
                    threshold=_MAX_PENDING_DELTAS,
                )
                with contextlib.suppress(ClientError, BotoCoreError):
                    self._download_index()  # triggers compaction as a side-effect
            self._write_delta(entry)
            log.debug("delta_written", session_id=record.id)

            # 3. Sync any pending local sessions
            synced = self.sync_pending()
            if synced > 0:
                log.info("pending_sessions_synced", count=synced)

        except (ClientError, BotoCoreError) as exc:
            log.warning(
                "s3_persist_failed_using_local",
                session_id=record.id,
                category=category,
                error=str(exc),
            )
            err_console.print(
                f"[yellow]⚠ Cloud Locker unavailable: {exc}. Session saved locally.[/yellow]"
            )
            self._write_local_with_lock(record, category)

    def load_index(self) -> LockerIndex:
        """Download and parse the index manifest, applying any pending deltas.

        Falls back to the local index on S3 failure.
        """
        try:
            return self._download_index()
        except (ClientError, BotoCoreError):
            return self._load_local_index()

    def load_session(self, session_id: str, category: str) -> SessionRecord:
        """Download and deserialize a single session object."""
        log.debug("load_session_starting", session_id=session_id, category=category)
        s3_key = self._s3_key(category, session_id)
        try:
            resp = self._get_object_with_retry(Bucket=self._bucket, Key=s3_key)
            body = resp["Body"].read().decode("utf-8")
            log.info("session_loaded_s3", session_id=session_id, category=category)
            return deserialize_record(body)
        except (ClientError, BotoCoreError) as exc:
            log.debug("s3_load_failed_trying_local", session_id=session_id, error=str(exc))
            # Try local fallback
            local = self._local_path(category, session_id)
            if local.exists():
                log.info("session_loaded_local", session_id=session_id, category=category)
                return deserialize_record(local.read_text(encoding="utf-8"))
            log.error("session_not_found", session_id=session_id, category=category)
            raise LockerError(f"Session {session_id} not found in S3 or local cache.") from exc

    def sync_pending(self) -> int:
        """Upload locally cached sessions to S3. Returns count synced."""
        if not _LOCKER_DIR.exists():
            return 0

        synced = 0
        # Walk category subdirectories (chat, invoke, batch, compare, dialog)
        for category_dir in _LOCKER_DIR.iterdir():
            if not category_dir.is_dir():
                continue
            category = category_dir.name
            for session_file in category_dir.glob("*.json"):
                session_id = session_file.stem
                s3_key = self._s3_key(category, session_id)
                try:
                    body = session_file.read_bytes()
                    self._put_object_with_retry(
                        Bucket=self._bucket,
                        Key=s3_key,
                        Body=body,
                        ContentType="application/json",
                        ServerSideEncryption="AES256",
                    )
                    # Upload succeeded — write delta for index entry
                    record = deserialize_record(body.decode("utf-8"))
                    entry = build_index_entry(record)
                    self._write_delta(entry)

                    # Remove the local file after successful sync
                    session_file.unlink()
                    synced += 1
                except (ClientError, BotoCoreError):
                    # S3 still unreachable — leave local file for next attempt
                    pass

        # Clean up empty category directories
        for category_dir in _LOCKER_DIR.iterdir():
            if category_dir.is_dir() and not any(category_dir.iterdir()):
                category_dir.rmdir()

        return synced

    # ── S3 key construction ───────────────────────────────────────────

    def _s3_key(self, category: str, session_id: str) -> str:
        """Build: {namespace}/{environment}/{username}/{category}/{session_id}.json"""
        return (
            f"{self._namespace}/{self._environment}/{self._username}/{category}/{session_id}.json"
        )

    def _index_key(self) -> str:
        """Build: {namespace}/{environment}/{username}/index.json"""
        return f"{self._namespace}/{self._environment}/{self._username}/index.json"

    def _delta_key(self, session_id: str) -> str:
        """Build: {namespace}/{environment}/{username}/_delta/{session_id}.json"""
        return f"{self._namespace}/{self._environment}/{self._username}/_delta/{session_id}.json"

    def _delta_prefix(self) -> str:
        """Build: {namespace}/{environment}/{username}/_delta/"""
        return f"{self._namespace}/{self._environment}/{self._username}/_delta/"

    # ── local path construction ───────────────────────────────────────

    def _local_path(self, category: str, session_id: str) -> Path:
        """Build: ~/.alice/locker/{category}/{session_id}.json"""
        return _LOCKER_DIR / category / f"{session_id}.json"

    def _local_index_path(self) -> Path:
        """Return: ~/.alice/locker/index.json"""
        return _LOCKER_DIR / "index.json"

    # ── local fallback with filelock ──────────────────────────────────

    def _write_local_with_lock(self, record: SessionRecord, category: str) -> None:
        """Write session + update local index under filelock."""
        # Write the session file
        local = self._local_path(category, record.id)
        local.parent.mkdir(parents=True, exist_ok=True)
        local.write_text(serialize_record(record), encoding="utf-8")

        # Update the local index under a file lock
        index_path = self._local_index_path()
        lock_path = index_path.with_suffix(".json.lock")
        index_path.parent.mkdir(parents=True, exist_ok=True)

        try:
            with FileLock(lock_path, timeout=FILELOCK_TIMEOUT):
                index = self._load_local_index()
                entry = build_index_entry(record)
                index.entries.append(entry)
                index_path.write_bytes(msgspec.json.encode(_serialize_index(index)))
        except Timeout as exc:
            raise LockerError(
                "Could not acquire lock on local index — "
                "another alice process may be writing. Try again shortly."
            ) from exc

    # ── S3 delta helpers (incremental index) ─────────────────────────

    def _write_delta(self, entry: IndexEntry) -> None:
        """Write a single index entry as a tiny delta object.

        Uses the session ID as the object key so the write is idempotent —
        re-persisting the same session replaces its delta rather than
        creating a duplicate.  No read of the full index is needed.
        """
        body = msgspec.json.encode(_serialize_entry(entry))
        self._put_object_with_retry(
            Bucket=self._bucket,
            Key=self._delta_key(entry.id),
            Body=body,
            ContentType="application/json",
            ServerSideEncryption="AES256",
        )

    def _list_delta_keys(self) -> list[str]:
        """Return all delta object keys for this user."""
        prefix = self._delta_prefix()
        keys: list[str] = []
        paginator = self._s3.get_paginator("list_objects_v2")
        for page in paginator.paginate(Bucket=self._bucket, Prefix=prefix):
            for obj in page.get("Contents", []):
                keys.append(obj["Key"])
        return keys

    # ── S3 index helpers ──────────────────────────────────────────────

    def _download_index(self) -> LockerIndex:
        """Download index.json, apply any pending deltas, and compact.

        Compaction: if deltas are found, the merged index is re-uploaded as
        index.json and the delta objects are deleted, so the next read is fast.
        """
        # 1. Load the compacted base index
        try:
            resp = self._get_object_with_retry(Bucket=self._bucket, Key=self._index_key())
            index = _parse_index(resp["Body"].read().decode("utf-8"))
        except ClientError as exc:
            if exc.response["Error"]["Code"] == "NoSuchKey":
                index = LockerIndex(entries=[])
            else:
                raise

        # 2. Find any pending delta objects
        try:
            delta_keys = self._list_delta_keys()
        except (ClientError, BotoCoreError):
            delta_keys = []

        if not delta_keys:
            return index

        # 3. Download deltas and merge into the index
        existing_ids = {e.id for e in index.entries}
        for key in delta_keys:
            try:
                resp = self._get_object_with_retry(Bucket=self._bucket, Key=key)
                entry = _parse_entry(resp["Body"].read())
                if entry.id not in existing_ids:
                    index.entries.append(entry)
                    existing_ids.add(entry.id)
            except Exception:  # noqa: BLE001 — best-effort per delta
                log.warning("delta_apply_failed", key=key)

        # 4. Compact: re-upload merged index and delete processed deltas
        try:
            self._upload_index(index)
            for key in delta_keys:
                with contextlib.suppress(ClientError, BotoCoreError):
                    self._s3.delete_object(Bucket=self._bucket, Key=key)
            log.debug(
                "index_compacted", delta_count=len(delta_keys), total_entries=len(index.entries)
            )
        except (ClientError, BotoCoreError):
            pass  # Return merged index even if compaction fails

        return index

    def _upload_index(self, index: LockerIndex) -> None:
        """Serialize and upload the index to S3."""
        body = msgspec.json.encode(_serialize_index(index))
        self._put_object_with_retry(
            Bucket=self._bucket,
            Key=self._index_key(),
            Body=body,
            ContentType="application/json",
            ServerSideEncryption="AES256",
        )

    def backfill_tokens(self, entries: list[IndexEntry]) -> tuple[list[IndexEntry], int]:
        """Backfill token data for index entries missing it.

        Loads the full session record from S3/local for each entry whose
        ``tokens`` field is ``None`` and copies the token data across.
        Returns the updated list and the count of entries that were filled.

        Uses the entry's ``type`` field to determine the storage category directly
        (O(1) per entry) rather than trying all categories (O(C) per entry).

        The S3 index is re-uploaded if any entries were patched.
        """
        filled = 0

        for entry in entries:
            if entry.tokens is not None:
                continue
            category = _TYPE_TO_CATEGORY.get(entry.type)
            if not category:
                log.debug("backfill_skip_unknown_type", entry_id=entry.id, entry_type=entry.type)
                continue
            try:
                record = self.load_session(entry.id, category)
                if record.tokens is not None:
                    entry.tokens = record.tokens
                    filled += 1
            except Exception:
                continue

        # Persist the patched index so future calls don't need to backfill again
        if filled:
            try:
                index = self._download_index()
                patched_map = {e.id: e.tokens for e in entries if e.tokens is not None}
                for ie in index.entries:
                    if ie.tokens is None and ie.id in patched_map:
                        ie.tokens = patched_map[ie.id]
                self._upload_index(index)
            except (ClientError, BotoCoreError):
                pass  # Best-effort

        return entries, filled

    # ── local index helpers ───────────────────────────────────────────

    def _load_local_index(self) -> LockerIndex:
        """Load and parse the local index. Returns empty on missing/corrupt."""
        index_path = self._local_index_path()
        if not index_path.exists():
            return LockerIndex(entries=[])
        try:
            raw = index_path.read_text(encoding="utf-8")
            return _parse_index(raw)
        except (msgspec.DecodeError, KeyError, TypeError):
            # Corrupted index — rename to .bak and start fresh
            bak = index_path.with_suffix(".json.bak")
            index_path.rename(bak)
            return LockerIndex(entries=[])


# ── module-level helpers ──────────────────────────────────────────────


def _serialize_entry(entry: IndexEntry) -> dict[str, Any]:
    """Serialize a single IndexEntry to a JSON-serializable dict."""
    return {
        "id": entry.id,
        "type": entry.type,
        "timestamp": entry.timestamp,
        "model": entry.model,
        "prompt_preview": entry.prompt_preview,
        "tokens": (
            {
                "input": entry.tokens.input,
                "output": entry.tokens.output,
                "total": entry.tokens.total,
            }
            if entry.tokens is not None
            else None
        ),
    }


def _parse_entry(raw: bytes) -> IndexEntry:
    """Deserialize a delta object (bytes) back to an IndexEntry."""
    item: dict[str, Any] = msgspec.json.decode(raw)
    tokens_raw = item.get("tokens")
    tokens = (
        TokenUsage(
            input=tokens_raw["input"],
            output=tokens_raw["output"],
            total=tokens_raw["total"],
        )
        if tokens_raw is not None
        else None
    )
    return IndexEntry(
        id=item["id"],
        type=item["type"],
        timestamp=item["timestamp"],
        model=item["model"],
        prompt_preview=item["prompt_preview"],
        tokens=tokens,
    )


def _parse_index(raw: str) -> LockerIndex:
    """Parse a JSON string into a LockerIndex."""
    data: dict[str, Any] = msgspec.json.decode(raw.encode("utf-8"))
    entries: list[IndexEntry] = []
    for item in data.get("entries", []):
        tokens_raw = item.get("tokens")
        tokens = (
            TokenUsage(
                input=tokens_raw["input"],
                output=tokens_raw["output"],
                total=tokens_raw["total"],
            )
            if tokens_raw is not None
            else None
        )
        entries.append(
            IndexEntry(
                id=item["id"],
                type=item["type"],
                timestamp=item["timestamp"],
                model=item["model"],
                prompt_preview=item["prompt_preview"],
                tokens=tokens,
            )
        )
    return LockerIndex(entries=entries)


def _serialize_index(index: LockerIndex) -> dict[str, Any]:
    """Convert a LockerIndex to a JSON-serializable dict."""
    return {"entries": [_serialize_entry(e) for e in index.entries]}
