"""Unified retrieval layer with caching.

Features should use RetrievalLayer instead of calling storage/backend directly.
This ensures consistent bounds, caching, and error handling.

Usage:
    r = RetrievalLayer("my-project")

    # Simple access
    traces = r.traces(limit=100)
    evals = r.evals()

    # Time-scoped access
    scoped = r.with_time_window(days=14)
    traces = scoped.traces(tag="prod")
    evals = scoped.evals()

    # Single trace
    trace = r.get_trace("trace-id-123")
"""

from __future__ import annotations

import logging
from datetime import datetime, timedelta, timezone
from typing import Any

from .cache import DiskCache, get_cache
from ..config import load_config, resolve_cache_enabled, resolve_traces_dir
from ..storage import get_store

logger = logging.getLogger(__name__)


class ProjectStats:
    """Aggregated project statistics."""

    def __init__(
        self,
        total_traces: int = 0,
        total_cost: float = 0.0,
        avg_cost: float = 0.0,
        span_kinds: dict[str, int | None] | None = None,
        daily_counts: list[int] | None = None,
    ):
        self.total_traces = total_traces
        self.total_cost = total_cost
        self.avg_cost = avg_cost
        self.span_kinds = span_kinds or {}
        self.daily_counts = daily_counts or []


class ScopedRetrieval:
    """Retrieval with a pre-set time window.

    Use with_time_window() to create.
    """

    def __init__(self, retrieval: RetrievalLayer, since: datetime | None):
        self._r = retrieval
        self._since = since

    def traces(self, **kwargs: Any) -> list[dict[str, Any]]:
        """Load traces with inherited time window."""
        if self._since is not None and "since" not in kwargs:
            kwargs["since"] = self._since
        return self._r.traces(**kwargs)

    def evals(self, **kwargs: Any) -> dict[str, dict[str, Any]]:
        """Load evaluations with inherited time window."""
        if self._since is not None and "since" not in kwargs:
            kwargs["since"] = self._since
        return self._r.evals(**kwargs)


class RetrievalLayer:
    """Unified data retrieval with caching.

    Features should use this instead of calling storage/backend directly.
    """

    def __init__(
        self,
        project: str,
        *,
        cache: DiskCache | None = None,
        use_cache: bool | None = None,
    ):
        self._project = project
        self._store = None  # Lazy init
        self._client = None  # Lazy init

        if use_cache is None:
            config = load_config()
            use_cache = resolve_cache_enabled(config)
        
        if use_cache:
            self._cache = cache or get_cache()
        else:
            self._cache = None

    @property
    def project(self) -> str:
        return self._project

    @property
    def store(self) -> Any:
        if self._store is None:
            self._store = get_store(project=self._project)
        return self._store

    def _get_client(self) -> Any:
        if self._client is None:
            try:
                from .api import BanditoClient

                self._client = BanditoClient()
            except RuntimeError:
                self._client = None
        return self._client

    def traces(
        self,
        *,
        since: datetime | None = None,
        limit: int | None = None,
        offset: int = 0,
        tag: str | None = None,
        tier: str = "summary",
    ) -> list[dict[str, Any]]:
        """Load traces with automatic caching.

        Args:
            since: Only traces after this time
            limit: Max traces to return
            offset: Pagination offset
            tag: Filter by tag
            tier: "summary" (lightweight) or "trace" (full with spans)

        Returns:
            List of trace dicts
        """
        # For tier="summary", try API first (layer 2 from Bandito backend)
        if tier == "summary":
            api_result = self._fetch_summaries_from_api(
                since=since, limit=limit, offset=offset, tag=tag
            )
            if api_result is not None:
                return api_result

        cache_key = {
            "op": "traces",
            "tier": tier,
            "since": since.isoformat() if since else None,
            "limit": limit,
            "offset": offset,
            "tag": tag,
        }

        if self._cache is not None:
            cached = self._cache.get(self._project, cache_key)
            if cached is not None:
                return cached

        if tier == "summary":
            result = self.store.list_trace_summaries(
                project=self._project,
                tag=tag,
                since=since,
                limit=limit,
                offset=offset,
            )
        else:
            result = self.store.list_traces_filtered(
                project=self._project,
                tag=tag,
                since=since,
                limit=limit,
                offset=offset,
            )

        if self._cache is not None:
            self._cache.set(self._project, cache_key, result)
        return result

    def _fetch_summaries_from_api(
        self,
        *,
        since: datetime | None = None,
        limit: int | None = None,
        offset: int = 0,
        tag: str | None = None,
    ) -> list[dict[str, Any]] | None:
        """Fetch summaries from Bandito API. Returns None if unavailable."""
        client = self._get_client()
        if client is None:
            return None

        try:
            resp = client.list_summaries(
                project=self._project,
                limit=limit or 100,
                offset=offset,
                since=since.isoformat() if since else None,
                tag=tag,
            )
            summaries = resp.get("summaries", [])
            return [self._summary_to_trace(s) for s in summaries]
        except RuntimeError:
            logger.debug("Failed to fetch summaries from API")
            return None

    def _summary_to_trace(self, summary: dict[str, Any]) -> dict[str, Any]:
        """Convert API summary to trace-shaped dict for compatibility."""
        return {
            "trace_id": summary.get("trace_id"),
            "project": summary.get("project"),
            "started_at": summary.get("started_at"),
            "ended_at": summary.get("ended_at"),
            "status": summary.get("status", "ok"),
            "span_count": summary.get("span_count", 0),
            "total_cost": summary.get("total_cost"),
            "latency": summary.get("latency"),
            "tags": summary.get("tags", []),
        }

    def evals(
        self,
        *,
        since: datetime | None = None,
        limit: int | None = None,
        offset: int = 0,
    ) -> dict[str, dict[str, Any]]:
        """Load evaluations with automatic caching.

        Args:
            since: Only evaluations after this time (ISO string)
            limit: Max evaluations to return
            offset: Pagination offset

        Returns:
            Dict keyed by trace_id
        """
        client = self._get_client()
        if client is None:
            return {}

        cache_key = {
            "op": "evals",
            "since": since.isoformat() if since else None,
            "limit": limit,
            "offset": offset,
        }

        if self._cache is not None:
            cached = self._cache.get(self._project, cache_key)
            if cached is not None:
                return cached

        try:
            result = client.list_evaluations(
                self._project,
                limit=limit,
                offset=offset,
                since=since.isoformat() if since else None,
            )
        except RuntimeError:
            logger.debug("Failed to fetch evaluations from backend")
            result = {}

        if self._cache is not None:
            self._cache.set(self._project, cache_key, result)
        return result

    def stats(self) -> ProjectStats:
        """Get project statistics efficiently.

        Uses cache to avoid recounting traces on every call.
        Tries API first (pre-computed), falls back to local.
        """
        cache_key = {"op": "stats"}

        if self._cache is not None:
            cached = self._cache.get(self._project, cache_key)
            if cached is not None:
                return self._make_stats(cached)

        # Try API first (pre-computed stats from Bandito backend)
        client = self._get_client()
        if client is not None:
            try:
                api_stats = client.get_project_stats(self._project)
                if self._cache is not None:
                    self._cache.set(self._project, cache_key, api_stats)
                return self._make_stats(api_stats)
            except RuntimeError:
                logger.debug("Failed to fetch stats from API")

        store_stats = self.store.get_stats(self._project)

        result = {
            "total_traces": store_stats.total_traces,
            "total_cost": store_stats.total_cost,
            "avg_cost": store_stats.avg_cost,
            "span_kinds": store_stats.span_kinds,
            "daily_counts": store_stats.daily_counts,
        }

        if self._cache is not None:
            self._cache.set(self._project, cache_key, result)
        return self._make_stats(result)

    def _make_stats(self, data: dict[str, Any]) -> ProjectStats:
        return ProjectStats(
            total_traces=data.get("total_traces", 0),
            total_cost=data.get("total_cost", 0.0),
            avg_cost=data.get("avg_cost", 0.0),
            span_kinds=data.get("span_kinds", {}),
            daily_counts=data.get("daily_counts", []),
        )

    def get_trace(self, trace_id: str) -> dict[str, Any] | None:
        """Get a single trace by ID."""
        return self.store.get_trace(trace_id)

    def with_time_window(self, days: int) -> ScopedRetrieval:
        """Create a scoped retrieval with a time window.

        All traces/evals will be filtered to the last N days.
        """
        since = datetime.now(timezone.utc) - timedelta(days=days)
        return ScopedRetrieval(self, since)

    def invalidate(self) -> None:
        """Clear all cached data for this project."""
        if self._cache is not None:
            self._cache.invalidate_project(self._project)

    def refresh(
        self,
        *,
        layers: list[int] | None = None,
        lite: bool = False,
    ) -> dict[str, int]:
        """Refresh cache for specified layers.
        
        Args:
            layers: List of layers to refresh [1, 2, 3]. None = all.
            lite: If True, only refresh layers 1 + 2 (faster).
        
        Returns:
            Dict with counts: {"projects": X, "traces": Y, "full_traces": Z}
        """
        counts = {"projects": 0, "traces": 0, "full_traces": 0}
        
        # Determine layers to refresh
        if lite:
            layers = [1, 2]
        elif layers is None:
            layers = [1, 2, 3]
        
        # Layer 1: Project list (global, not project-scoped)
        if 1 in layers:
            global _global_projects_cache
            _global_projects_cache = None
            counts["projects"] = 1
        
        # Layer 2: Trace summaries
        if 2 in layers:
            self._refresh_layer("summary", counts, "traces")
        
        # Layer 3: Full traces
        if 3 in layers:
            self._refresh_layer("trace", counts, "full_traces")
        
        return counts

    def _refresh_layer(self, tier: str, counts: dict[str, int], count_key: str) -> None:
        """Helper to refresh a specific layer."""
        import logging
        logger = logging.getLogger(__name__)
        
        if self._cache is None:
            return
        
        cache: DiskCache = self._cache  # type: ignore[assignment]
        cache_key = {"op": "traces", "tier": tier}
        cache.delete(self._project, cache_key)
        
        try:
            if tier == "summary":
                data = self.store.list_trace_summaries(project=self._project)
            else:
                data = self.store.list_traces_filtered(project=self._project)
            cache.set(self._project, cache_key, data)
            counts[count_key] = len(data)
        except Exception as e:
            logger.warning(f"Failed to refresh layer {tier} for {self._project}: {e}")


_retrieval_cache: dict[str, RetrievalLayer] = {}
_retrieval_lock: Any = None
_global_projects_cache: list[str] | None = None


def list_projects(use_cache: bool = True) -> list[str]:
    """Get list of projects, with optional caching.
    
    Tries API first (Bandito backend has aggregated project list),
    falls back to local storage scan.
    
    Args:
        use_cache: If True, serve from cache if available.
    
    Returns:
        List of project names.
    """
    global _global_projects_cache
    
    cache_key = {"op": "list_projects"}
    
    if use_cache:
        disk_cache = get_cache()
        cached = disk_cache.get("__global__", cache_key)
        if cached is not None:
            return cached
    
    # Try API first (layer 1 from Bandito backend)
    try:
        from .api import BanditoClient
        client = BanditoClient()
        projects = client.list_projects()
        _global_projects_cache = projects
        if use_cache:
            disk_cache.set("__global__", cache_key, projects)
        return projects
    except RuntimeError:
        pass
    
    # Fall back to local storage
    store = get_store()
    projects = store.list_projects()
    _global_projects_cache = projects
    if use_cache:
        disk_cache.set("__global__", cache_key, projects)
    return projects


def get_retrieval(project: str) -> RetrievalLayer:
    """Get or create a RetrievalLayer for a project.

    RetrievalLayers are cached per project within a session.
    """
    global _retrieval_cache, _retrieval_lock

    try:
        import threading

        if _retrieval_lock is None:
            _retrieval_lock = threading.Lock()
    except ImportError:
        _retrieval_lock = None

    if _retrieval_lock:
        with _retrieval_lock:
            if project not in _retrieval_cache:
                _retrieval_cache[project] = RetrievalLayer(project)
            return _retrieval_cache[project]
    else:
        if project not in _retrieval_cache:
            _retrieval_cache[project] = RetrievalLayer(project)
        return _retrieval_cache[project]


def reset_retrieval() -> None:
    """Reset retrieval cache (for testing)."""
    global _retrieval_cache
    _retrieval_cache = {}
