# retoor <retoor@molodetz.nl>

from __future__ import annotations

import contextvars
import json
import logging
import uuid
from collections import OrderedDict
from typing import Optional

logger = logging.getLogger("devii.chunks")

STORE_MAX_ENTRIES = 16
STORE_MAX_CHARS = 8_000_000

_active_store: contextvars.ContextVar = contextvars.ContextVar(
    "devii_chunk_store", default=None
)


class ChunkStore:
    def __init__(self) -> None:
        self._items: "OrderedDict[str, str]" = OrderedDict()
        self._resources: dict[str, str] = {}

    def put(self, text: str, resource_key: Optional[str] = None) -> str:
        chunk_id = uuid.uuid4().hex[:12]
        self._items[chunk_id] = text[:STORE_MAX_CHARS]
        self._items.move_to_end(chunk_id)
        if resource_key:
            self._resources[resource_key] = chunk_id
        while len(self._items) > STORE_MAX_ENTRIES:
            evicted, _ = self._items.popitem(last=False)
            self._drop_resource_for(evicted)
        logger.debug("Stored chunkable result %s (%d chars)", chunk_id, len(text))
        return chunk_id

    def get(self, chunk_id: str) -> Optional[str]:
        text = self._items.get(chunk_id)
        if text is not None:
            self._items.move_to_end(chunk_id)
        return text

    def get_by_resource(self, resource_key: str) -> Optional[tuple[str, str]]:
        chunk_id = self._resources.get(resource_key)
        if chunk_id is None:
            return None
        text = self._items.get(chunk_id)
        if text is None:
            del self._resources[resource_key]
            return None
        self._items.move_to_end(chunk_id)
        return chunk_id, text

    def invalidate_resources(self, prefix: str = "") -> None:
        if not prefix:
            self._resources.clear()
            return
        for key in [key for key in self._resources if key.startswith(prefix)]:
            del self._resources[key]

    def _drop_resource_for(self, chunk_id: str) -> None:
        for key in [key for key, value in self._resources.items() if value == chunk_id]:
            del self._resources[key]


def get_store() -> Optional[ChunkStore]:
    return _active_store.get()


def set_store(store: ChunkStore) -> contextvars.Token:
    return _active_store.set(store)


def reset_store(token: contextvars.Token) -> None:
    _active_store.reset(token)


def chunk_envelope(text: str, chunk_id: str, max_chars: int) -> str:
    total = len(text)
    shown = text[:max_chars]
    return json.dumps(
        {
            "truncated": True,
            "chunk_id": chunk_id,
            "total_chars": total,
            "shown_chars": len(shown),
            "remaining_chars": total - len(shown),
            "next_offset": len(shown),
            "hint": (
                f"Output truncated. Call read_more with chunk_id='{chunk_id}' and "
                f"offset={len(shown)} to read the next part (repeat until remaining_chars is 0). "
                "Do not re-run the original tool; the full content is already cached."
            ),
            "content": shown,
        },
        ensure_ascii=False,
    )


def wrap_if_large(
    result: str, max_chars: int, resource_key: Optional[str] = None
) -> str:
    if len(result) <= max_chars:
        return result
    store = get_store()
    if store is None:
        total = len(result)
        shown = result[:max_chars]
        return json.dumps(
            {
                "truncated": True,
                "total_chars": total,
                "shown_chars": len(shown),
                "remaining_chars": total - len(shown),
                "hint": "Output truncated; no chunk store is available to fetch the rest.",
                "content": shown,
            },
            ensure_ascii=False,
        )
    chunk_id = store.put(result, resource_key)
    return chunk_envelope(result, chunk_id, max_chars)


def serve_resource(resource_key: str, max_chars: int) -> Optional[str]:
    store = get_store()
    if store is None:
        return None
    hit = store.get_by_resource(resource_key)
    if hit is None:
        return None
    chunk_id, text = hit
    return chunk_envelope(text, chunk_id, max_chars)
