from __future__ import annotations

from datetime import datetime, timezone
from typing import Any, Dict, Iterable, List, Optional, Sequence, cast

from aiel_sdk.errors import RuntimeDependencyError

try:
    from langgraph.store.base import (
        BaseStore,
        GetOp,
        Item,
        ListNamespacesOp,
        MatchCondition,
        Op,
        PutOp,
        Result,
        SearchItem,
        SearchOp,
    )
except Exception as e:  # pragma: no cover
    raise RuntimeDependencyError("langgraph", "pip install 'aiel-sdk[langgraph]'") from e

from ...memory import RecallQuery
from ...memory.saver import Saver


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


def _parse_dt(value: Any) -> datetime:
    if isinstance(value, datetime):
        return value
    if isinstance(value, str):
        try:
            dt = datetime.fromisoformat(value.replace("Z", "+00:00"))
            if dt.tzinfo is None:
                return dt.replace(tzinfo=timezone.utc)
            return dt
        except ValueError:
            return _now_utc()
    return _now_utc()


class AIELLangGraphStore(BaseStore):
    """
    LangGraph BaseStore backed by aiel_sdk.memory.Saver recall operations.

    Notes:
    - Uses saver.recall_* as the long-term memory substrate.
    - `query` argument in SearchOp is mapped to saver.recall_search(query=...).
      If query is a string, it is forwarded as {"text": "..."} for backends that
      support text search through a query object.
    - Namespace listing is best-effort because Saver does not expose a dedicated
      list_namespaces primitive.
    """

    def __init__(self, saver: Saver) -> None:
        self._saver = saver

    @staticmethod
    def _match_filter(value: Dict[str, Any], filter_values: Optional[Dict[str, Any]]) -> bool:
        if not filter_values:
            return True
        for k, expected in filter_values.items():
            if value.get(k) != expected:
                return False
        return True

    @staticmethod
    def _ns_matches(namespace: tuple[str, ...], condition: MatchCondition) -> bool:
        path = condition.path
        if condition.match_type == "prefix":
            if len(namespace) < len(path):
                return False
            for idx, token in enumerate(path):
                if token != "*" and namespace[idx] != token:
                    return False
            return True
        if condition.match_type == "suffix":
            if len(namespace) < len(path):
                return False
            offset = len(namespace) - len(path)
            for idx, token in enumerate(path):
                if token != "*" and namespace[offset + idx] != token:
                    return False
            return True
        return False

    @staticmethod
    def _apply_namespace_filters(
        namespaces: List[tuple[str, ...]],
        op: ListNamespacesOp,
    ) -> List[tuple[str, ...]]:
        out = list(namespaces)
        if op.match_conditions:
            out = [
                ns
                for ns in out
                if all(AIELLangGraphStore._ns_matches(ns, cond) for cond in op.match_conditions)
            ]
        if op.max_depth is not None and op.max_depth > 0:
            out = [ns[: op.max_depth] for ns in out]
        out = sorted(set(out))
        return out[op.offset : op.offset + op.limit]

    def _op_get(self, op: GetOp) -> Item | None:
        record = self._saver.recall_get(namespace=list(op.namespace), key=op.key)
        if not record:
            return None
        ts = _parse_dt(record.updated_at)
        return Item(
            namespace=tuple(record.namespace),
            key=record.key,
            value=dict(record.value or {}),
            created_at=ts,
            updated_at=ts,
        )

    def _op_put(self, op: PutOp) -> None:
        if op.value is None:
            self._saver.recall_forget_key(namespace=list(op.namespace), key=op.key)
            return None
        self._saver.recall_put(
            namespace=list(op.namespace),
            key=op.key,
            value=dict(op.value or {}),
        )
        return None

    def _op_search(self, op: SearchOp) -> List[SearchItem]:
        backend_query: Dict[str, Any]
        if isinstance(op.query, str) and op.query.strip():
            backend_query = {"text": op.query.strip()}
        else:
            backend_query = {}
        recall_items = self._saver.recall_search(
            query=RecallQuery(
                namespace=list(op.namespace_prefix),
                query=backend_query,
                limit=op.limit + op.offset,
            )
        )
        filtered = [
            item for item in recall_items if self._match_filter(dict(item.value or {}), op.filter)
        ]
        selected = filtered[op.offset : op.offset + op.limit]
        out: List[SearchItem] = []
        for rec in selected:
            ts = _parse_dt(rec.updated_at)
            out.append(
                SearchItem(
                    namespace=tuple(rec.namespace),
                    key=rec.key,
                    value=dict(rec.value or {}),
                    created_at=ts,
                    updated_at=ts,
                    score=None,
                )
            )
        return out

    def _op_list_namespaces(self, op: ListNamespacesOp) -> List[tuple[str, ...]]:
        # Best-effort approach:
        # 1) Derive a candidate namespace prefix from first "prefix" match condition.
        # 2) Pull recall items under that prefix and dedupe namespaces.
        prefix: tuple[str, ...] = tuple()
        if op.match_conditions:
            for cond in op.match_conditions:
                if cond.match_type == "prefix":
                    prefix = cast(tuple[str, ...], tuple(p for p in cond.path if p != "*"))
                    break
        recall_items = self._saver.recall_search(
            query=RecallQuery(namespace=list(prefix), query={}, limit=max(op.limit + op.offset, 100))
        )
        namespaces = [tuple(item.namespace) for item in recall_items]
        return self._apply_namespace_filters(namespaces, op)

    def batch(self, ops: Iterable[Op]) -> List[Result]:
        results: List[Result] = []
        for op in ops:
            if isinstance(op, GetOp):
                results.append(self._op_get(op))
            elif isinstance(op, PutOp):
                self._op_put(op)
                results.append(None)
            elif isinstance(op, SearchOp):
                results.append(self._op_search(op))
            elif isinstance(op, ListNamespacesOp):
                results.append(self._op_list_namespaces(op))
            else:
                raise TypeError(f"Unsupported store operation: {type(op)!r}")
        return results

    async def abatch(self, ops: Iterable[Op]) -> List[Result]:
        return self.batch(ops)

