"""Langfuse storage backend — read traces from Langfuse API.

Read-only. Traces are created via Langfuse's own instrumentation (SDK, OpenAI
integration, etc.). Bandito reads them for analysis, grading, and judging.

Requires: pip install bandito[langfuse]
Auth: LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY, LANGFUSE_BASE_URL

Langfuse trace names map to Bandito project names.
"""

from __future__ import annotations

import datetime as dt
import json
import logging
import os
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any, Iterator

logger = logging.getLogger(__name__)

from .base import TraceReader, sanitize_project_name

# Max pages to scan for list_projects (avoids unbounded pagination)
_MAX_PROJECT_SCAN_PAGES = 10

# Max pages for iter_traces (200 pages * 50/page = 10,000 traces)
_MAX_ITER_PAGES = 200

# Cache traces for 10 minutes by default
_DEFAULT_CACHE_TTL_SECONDS = 600

# Note: LangfuseStore no longer uses internal disk caching.
# Caching is handled by RetrievalLayer's DiskCache instead.
_CACHE_DIR = str(Path.home() / ".bandito" / "cache" / "langfuse")


class LangfuseStore(TraceReader):
    """Read-only Langfuse storage backend with local disk cache."""

    def __init__(
        self,
        public_key: str | None = None,
        secret_key: str | None = None,
        base_url: str | None = None,
        project: str | None = None,
        cache_dir: str | None = None,
        cache_ttl: int = _DEFAULT_CACHE_TTL_SECONDS,
        timeout: float = 30.0,
    ) -> None:
        from langfuse import Langfuse

        kwargs: dict[str, Any] = {}
        if public_key:
            kwargs["public_key"] = public_key
        if secret_key:
            kwargs["secret_key"] = secret_key
        if base_url:
            kwargs["base_url"] = base_url
        kwargs["timeout"] = timeout

        self._client = Langfuse(**kwargs)
        self._project = project
        self._cache_dir = cache_dir or _CACHE_DIR
        self._cache_ttl = cache_ttl

    def validate(self) -> str | None:
        try:
            self._client.api.trace.list(limit=1)
            return None
        except Exception as e:
            return f"Cannot connect to Langfuse: {e}"

    # --- Cache ---

    def _cache_path(self, project: str, trace_id: str) -> str:
        safe_project = sanitize_project_name(project)
        return os.path.join(self._cache_dir, safe_project, f"{trace_id}.json")

    def _cache_read(self, project: str, trace_id: str) -> dict[str, Any] | None:
        path = self._cache_path(project, trace_id)
        try:
            stat = os.stat(path)
            age = dt.datetime.now().timestamp() - stat.st_mtime
            if age > self._cache_ttl:
                return None
            with open(path) as f:
                return json.load(f)
        except (OSError, json.JSONDecodeError):
            return None

    def _cache_write(self, trace: dict[str, Any]) -> None:
        project = sanitize_project_name(trace.get("project", "unknown"))
        trace_id = trace.get("trace_id", "")
        if not trace_id:
            return
        path = self._cache_path(project, trace_id)
        try:
            os.makedirs(os.path.dirname(path), exist_ok=True)
            with open(path, "w") as f:
                json.dump(trace, f)
        except OSError as e:
            logger.debug("Cache write failed for %s/%s: %s", project, trace_id, e)

    def _cache_list_traces(self, project: str) -> list[dict[str, Any]]:
        """Read all cached traces for a project that haven't expired."""
        safe_project = sanitize_project_name(project)
        project_dir = os.path.join(self._cache_dir, safe_project)
        if not os.path.isdir(project_dir):
            return []
        traces = []
        now = dt.datetime.now().timestamp()
        for fname in os.listdir(project_dir):
            if not fname.endswith(".json"):
                continue
            path = os.path.join(project_dir, fname)
            try:
                age = now - os.stat(path).st_mtime
                if age > self._cache_ttl:
                    continue
                with open(path) as f:
                    traces.append(json.load(f))
            except (OSError, json.JSONDecodeError):
                continue
        return traces

    def _cache_has_recent(self, project: str) -> bool:
        """Check if cache has any non-expired traces for this project."""
        safe_project = sanitize_project_name(project)
        project_dir = os.path.join(self._cache_dir, safe_project)
        if not os.path.isdir(project_dir):
            return False
        now = dt.datetime.now().timestamp()
        for fname in os.listdir(project_dir):
            if fname.endswith(".json"):
                try:
                    age = now - os.stat(os.path.join(project_dir, fname)).st_mtime
                    if age <= self._cache_ttl:
                        return True
                except OSError:
                    continue
        return False

    # --- Read ---

    def _read_trace(self, trace_id: str) -> dict[str, Any] | None:
        """Fetch a single trace. Checks cache first, then Langfuse API."""
        # Check cache across all projects
        for entry in os.listdir(self._cache_dir) if os.path.isdir(self._cache_dir) else []:
            cached = self._cache_read(entry, trace_id)
            if cached:
                return cached

        try:
            lf_trace = self._client.api.trace.get(trace_id)
            observations = getattr(lf_trace, "observations", []) or []
            mapped = _map_trace(lf_trace, observations, self._project)
            if mapped:
                self._cache_write(mapped)
            return mapped
        except Exception as e:
            print(f"  Warning: failed to fetch trace {trace_id}: {e}", file=sys.stderr)
            return None

    def iter_traces(self, project: str | None = None) -> Iterator[dict[str, Any]]:
        """Iterate traces from Langfuse.

        Serves from disk cache if fresh. Otherwise fetches from API,
        caches each trace, and yields.
        """
        target_project = project or self._project

        # Serve from cache if available and fresh
        if target_project and self._cache_has_recent(target_project):
            yield from self._cache_list_traces(target_project)
            return

        page = 1
        error_count = 0

        while True:
            kwargs: dict[str, Any] = {"page": page, "limit": 50}
            if target_project:
                kwargs["name"] = target_project

            try:
                response = self._client.api.trace.list(**kwargs)
            except Exception as e:
                if error_count == 0:
                    print(f"  Warning: Langfuse API error: {e}", file=sys.stderr)
                return

            trace_summaries = response.data if hasattr(response, "data") else []
            if not trace_summaries:
                return

            # Collect trace IDs for this page
            page_trace_ids: set[str] = set()
            for ts in trace_summaries:
                tid = getattr(ts, "id", None)
                if tid:
                    page_trace_ids.add(tid)

            # Batch-fetch observations for this page (1-2 API calls, not N)
            obs_by_trace = _fetch_observations_batch(
                self._client, trace_summaries, page_trace_ids
            )

            for ts in trace_summaries:
                tid = getattr(ts, "id", None)
                if not tid:
                    continue
                observations = obs_by_trace.get(tid, [])
                mapped = _map_trace(ts, observations, self._project)
                if mapped:
                    self._cache_write(mapped)
                    yield mapped

            meta = response.meta if hasattr(response, "meta") else None
            if meta:
                total_pages = (
                    getattr(meta, "total_pages", None)
                    or getattr(meta, "totalPages", None)
                )
                if total_pages and page >= total_pages:
                    return
            page += 1
            if page > _MAX_ITER_PAGES:
                print(
                    f"  Warning: reached max page limit ({_MAX_ITER_PAGES}), "
                    "results may be truncated",
                    file=sys.stderr,
                )
                return

    def list_projects(self) -> list[str]:
        """List unique trace names (Bandito projects).

        Scans up to _MAX_PROJECT_SCAN_PAGES pages of trace summaries.
        Caching is handled by RetrievalLayer's DiskCache, not here.
        """
        projects: set[str] = set()

        for page in range(1, _MAX_PROJECT_SCAN_PAGES + 1):
            try:
                response = self._client.api.trace.list(page=page, limit=100)
            except Exception as e:
                print(f"  Warning: Langfuse API error listing projects: {e}", file=sys.stderr)
                break

            traces = response.data if hasattr(response, "data") else []
            if not traces:
                break

            for trace in traces:
                name = getattr(trace, "name", None)
                if name:
                    projects.add(name)

            meta = response.meta if hasattr(response, "meta") else None
            if meta:
                total_pages = (
                    getattr(meta, "total_pages", None)
                    or getattr(meta, "totalPoints", None)
                )
                if total_pages and page >= total_pages:
                    break

        return sorted(projects)

    def count_traces(self, project: str | None = None) -> int:
        """Count traces using Langfuse API metadata - efficient."""
        target_project = project or self._project
        kwargs: dict[str, Any] = {"page": 1, "limit": 1}
        if target_project:
            kwargs["name"] = target_project

        try:
            response = self._client.api.trace.list(**kwargs)
        except Exception as e:
            print(f"  Warning: Langfuse API error counting traces: {e}", file=sys.stderr)
            return 0

        meta = response.meta if hasattr(response, "meta") else None
        if meta:
            total = (
                getattr(meta, "total", None)
                or getattr(meta, "totalCount", None)
                or getattr(meta, "total_items", None)
            )
            if total is not None:
                return int(total)

        return 0

    def list_trace_summaries(
        self,
        project: str | None = None,
        tag: str | None = None,
        since: "dt.datetime | None" = None,
        limit: int | None = None,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        """List trace summaries from Langfuse.

        Uses trace.list() which returns pre-computed cost and latency.
        Caching is handled by RetrievalLayer's DiskCache, not here.
        """
        from .summary import content_preview  # noqa: F401

        target_project = project or self._project
        summaries: list[dict[str, Any]] = []
        page = 1
        # Need offset + limit items to sort correctly before slicing
        collect_target = (offset + limit) if limit else None

        while True:
            kwargs: dict[str, Any] = {"page": page, "limit": 100}
            if target_project:
                kwargs["name"] = target_project
            if since:
                kwargs["from_timestamp"] = since

            try:
                response = self._client.api.trace.list(**kwargs)
            except Exception as e:
                print(f"  Warning: Langfuse API error: {e}", file=sys.stderr)
                break

            traces = response.data if hasattr(response, "data") else []
            if not traces:
                break

            for lf_trace in traces:
                summary = _map_langfuse_summary(lf_trace, self._project, tag)
                if summary is None:
                    continue
                summaries.append(summary)
                if collect_target and len(summaries) >= collect_target:
                    summaries.sort(key=lambda s: s.get("started_at", ""), reverse=True)
                    if offset:
                        summaries = summaries[offset:]
                    return summaries[:limit]

            meta = response.meta if hasattr(response, "meta") else None
            if meta:
                total_pages = (
                    getattr(meta, "total_pages", None)
                    or getattr(meta, "totalPages", None)
                )
                if total_pages and page >= total_pages:
                    break
            page += 1

        summaries.sort(key=lambda s: s.get("started_at", ""), reverse=True)
        if offset:
            summaries = summaries[offset:]
        if limit is not None and limit > 0:
            summaries = summaries[:limit]
        return summaries

    def list_traces_filtered(
        self,
        project: str | None = None,
        tag: str | None = None,
        since: "dt.datetime | None" = None,
        limit: int | None = None,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        """Override to pass time filters to Langfuse API directly."""
        target_project = project or self._project

        # Check cache first
        if target_project and self._cache_has_recent(target_project):
            cached = self._cache_list_traces(target_project)
            # Apply filters
            if since:
                cached = [
                    t for t in cached
                    if _parse_ts_for_filter(t.get("started_at")) is None
                    or _parse_ts_for_filter(t.get("started_at")) >= since
                ]
            if tag:
                cached = [t for t in cached if tag in t.get("tags", [])]
            cached.sort(key=lambda t: t.get("started_at", ""), reverse=True)
            if offset:
                cached = cached[offset:]
            if limit:
                cached = cached[:limit]
            return cached

        traces: list[dict[str, Any]] = []
        page = 1
        # Need offset + limit items to sort correctly before slicing
        collect_target = (offset + limit) if limit else None

        while True:
            kwargs: dict[str, Any] = {"page": page, "limit": 50}
            if target_project:
                kwargs["name"] = target_project
            if since:
                kwargs["from_timestamp"] = since

            try:
                response = self._client.api.trace.list(**kwargs)
            except Exception as e:
                print(f"  Warning: Langfuse API error: {e}", file=sys.stderr)
                break

            trace_summaries = response.data if hasattr(response, "data") else []
            if not trace_summaries:
                break

            page_trace_ids: set[str] = set()
            for ts in trace_summaries:
                tid = getattr(ts, "id", None)
                if tid:
                    page_trace_ids.add(tid)

            obs_by_trace = _fetch_observations_batch(
                self._client, trace_summaries, page_trace_ids
            )

            for ts in trace_summaries:
                tid = getattr(ts, "id", None)
                if not tid:
                    continue
                observations = obs_by_trace.get(tid, [])
                mapped = _map_trace(ts, observations, self._project)
                if not mapped:
                    continue
                self._cache_write(mapped)
                if tag and tag not in mapped.get("tags", []):
                    continue
                traces.append(mapped)
                if collect_target and len(traces) >= collect_target:
                    traces.sort(key=lambda t: t.get("started_at", ""), reverse=True)
                    if offset:
                        traces = traces[offset:]
                    return traces[:limit]

            meta = response.meta if hasattr(response, "meta") else None
            if meta:
                total_pages = (
                    getattr(meta, "total_pages", None)
                    or getattr(meta, "totalPages", None)
                )
                if total_pages and page >= total_pages:
                    break
            page += 1

        traces.sort(key=lambda t: t.get("started_at", ""), reverse=True)
        if offset:
            traces = traces[offset:]
        if limit is not None and limit > 0:
            traces = traces[:limit]
        return traces


# --- Observation fetching ---


def _fetch_observations_batch(
    client: Any,
    trace_summaries: list[Any],
    trace_ids: set[str],
) -> dict[str, list[Any]]:
    """Batch-fetch observations for a page of traces.

    Uses time-window filtering on the observations endpoint to get all
    observations in one paginated call, then groups by trace_id.
    This is 1-2 API calls per page instead of N per-trace calls.
    """
    obs_by_trace: dict[str, list[Any]] = defaultdict(list)
    if not trace_ids:
        return obs_by_trace

    # Compute time bounds from trace summaries
    min_time = None
    max_time = None
    for ts in trace_summaries:
        timestamp = getattr(ts, "timestamp", None)
        if timestamp:
            if min_time is None or timestamp < min_time:
                min_time = timestamp
            if max_time is None or timestamp > max_time:
                max_time = timestamp

    cursor = None
    while True:
        kwargs: dict[str, Any] = {
            "limit": 1000,
            "fields": "core,basic,model,usage,io",
        }
        if min_time:
            kwargs["from_start_time"] = min_time
        if max_time:
            # Buffer for observations that start after trace timestamp
            if isinstance(max_time, dt.datetime):
                kwargs["to_start_time"] = max_time + dt.timedelta(hours=1)
        if cursor:
            kwargs["cursor"] = cursor

        try:
            response = client.api.observations.get_many(**kwargs)
        except Exception as e:
            print(f"  Warning: Langfuse API error fetching observations: {e}", file=sys.stderr)
            break

        data = response.data if hasattr(response, "data") else []
        if not data:
            break

        for obs in data:
            trace_id = (
                getattr(obs, "trace_id", None)
                or getattr(obs, "traceId", None)
            )
            if trace_id and trace_id in trace_ids:
                obs_by_trace[trace_id].append(obs)

        meta = response.meta if hasattr(response, "meta") else None
        if meta:
            next_cursor = (
                getattr(meta, "next_cursor", None)
                or getattr(meta, "nextCursor", None)
            )
            if next_cursor and next_cursor != cursor:
                cursor = next_cursor
            else:
                break
        else:
            break

    return obs_by_trace


# --- Mapping: Langfuse → Bandito ---


def _map_langfuse_summary(
    lf_trace: Any,
    default_project: str | None = None,
    tag: str | None = None,
) -> dict[str, Any] | None:
    """Map a Langfuse TraceWithDetails to a Bandito summary dict.

    Returns None if the trace should be filtered out (e.g., tag mismatch).
    """
    from .summary import content_preview

    trace_name = getattr(lf_trace, "name", None) or default_project or "langfuse"
    if tag:
        trace_tags = getattr(lf_trace, "tags", []) or []
        if tag not in trace_tags:
            return None

    inp = getattr(lf_trace, "input", None)
    out = getattr(lf_trace, "output", None)

    summary: dict[str, Any] = {
        "trace_id": getattr(lf_trace, "id", ""),
        "project": trace_name,
        "started_at": _format_ts(getattr(lf_trace, "timestamp", None)) or "",
        "ended_at": None,
        "status": "ok",
        "tags": getattr(lf_trace, "tags", []) or [],
        "user_id": getattr(lf_trace, "user_id", None) or getattr(lf_trace, "userId", None),
        "session_id": getattr(lf_trace, "session_id", None) or getattr(lf_trace, "sessionId", None),
        "input_preview": content_preview(_unwrap_message(inp)) if inp else "",
        "output_preview": content_preview(_unwrap_message(out)) if out else "",
        "total_cost": getattr(lf_trace, "total_cost", None) or getattr(lf_trace, "totalCost", None),
        "latency": getattr(lf_trace, "latency", None),
        "span_count": len(getattr(lf_trace, "observations", []) or []),
        "span_kinds": [],
        "primary_model": None,
        "total_tokens": None,
        "error_count": 0,
    }

    if summary["latency"] and summary["started_at"]:
        try:
            start_dt = dt.datetime.fromisoformat(summary["started_at"].replace("Z", "+00:00"))
            summary["ended_at"] = (
                start_dt + dt.timedelta(seconds=summary["latency"])
            ).isoformat().replace("+00:00", "Z")
        except (ValueError, TypeError):
            pass

    return summary


def _map_trace(
    lf_trace: Any,
    observations: list[Any],
    default_project: str | None = None,
) -> dict[str, Any] | None:
    """Map a Langfuse trace + observations to Bandito trace dict."""
    trace_id = getattr(lf_trace, "id", None)
    if not trace_id:
        return None

    spans = [s for s in (_map_observation(obs) for obs in observations) if s is not None]

    started_at = _format_ts(getattr(lf_trace, "timestamp", None))

    # End time: prefer latency, fall back to latest span
    latency = getattr(lf_trace, "latency", None)
    ended_at = None
    if latency is not None and started_at:
        try:
            start_dt = dt.datetime.fromisoformat(started_at.replace("Z", "+00:00"))
            ended_at = (start_dt + dt.timedelta(seconds=latency)).isoformat().replace("+00:00", "Z")
        except (ValueError, TypeError):
            pass
    if not ended_at and spans:
        end_times = [s.get("ended_at") for s in spans if s.get("ended_at")]
        if end_times:
            ended_at = max(end_times)

    trace: dict[str, Any] = {
        "trace_id": trace_id,
        "project": getattr(lf_trace, "name", None) or default_project or "langfuse",
        "started_at": started_at,
        "status": "ok",
        "spans": spans,
    }

    if ended_at:
        trace["ended_at"] = ended_at

    inp = getattr(lf_trace, "input", None)
    if inp is not None:
        trace["input"] = _unwrap_message(inp)

    out = getattr(lf_trace, "output", None)
    if out is not None:
        trace["output"] = _unwrap_message(out)

    user_id = getattr(lf_trace, "user_id", None) or getattr(lf_trace, "userId", None)
    if user_id:
        trace["user_id"] = user_id

    session_id = getattr(lf_trace, "session_id", None) or getattr(lf_trace, "sessionId", None)
    if session_id:
        trace["session_id"] = session_id

    metadata = getattr(lf_trace, "metadata", None)
    if metadata:
        trace["metadata"] = metadata

    tags = getattr(lf_trace, "tags", None)
    if tags:
        trace["tags"] = tags

    if any(s.get("status") == "error" for s in spans):
        trace["status"] = "error"

    return trace


def _map_observation(obs: Any) -> dict[str, Any] | None:
    """Map a Langfuse observation to a Bandito span dict."""
    span_id = getattr(obs, "id", None)
    if not span_id:
        return None

    obs_type = getattr(obs, "type", None) or "SPAN"
    kind = _map_observation_type(obs_type)

    span: dict[str, Any] = {
        "span_id": span_id,
        "parent_id": (
            getattr(obs, "parent_observation_id", None)
            or getattr(obs, "parentObservationId", None)
        ),
        "name": getattr(obs, "name", None) or kind,
        "kind": kind,
        "started_at": _format_ts(
            getattr(obs, "start_time", None) or getattr(obs, "startTime", None)
        ),
        "status": _map_level(getattr(obs, "level", None)),
    }

    ended_at = _format_ts(
        getattr(obs, "end_time", None) or getattr(obs, "endTime", None)
    )
    if ended_at:
        span["ended_at"] = ended_at

    inp = getattr(obs, "input", None)
    if inp is not None:
        span["input"] = _unwrap_message(inp)

    out = getattr(obs, "output", None)
    if out is not None:
        span["output"] = _unwrap_message(out)

    # Model — if present, this is an LLM span
    model = (
        getattr(obs, "model", None)
        or getattr(obs, "provided_model_name", None)
        or getattr(obs, "providedModelName", None)
    )
    if model:
        span["model"] = model
        span["kind"] = "llm"

    # Token usage
    usage = (
        getattr(obs, "usage_details", None)
        or getattr(obs, "usageDetails", None)
    )
    if isinstance(usage, dict):
        input_tokens = usage.get("input") or usage.get("input_tokens")
        output_tokens = usage.get("output") or usage.get("output_tokens")
        if input_tokens:
            span["input_tokens"] = int(input_tokens)
        if output_tokens:
            span["output_tokens"] = int(output_tokens)

    # Cost
    cost = (
        getattr(obs, "calculated_total_cost", None)
        or getattr(obs, "calculatedTotalCost", None)
        or getattr(obs, "total_cost", None)
        or getattr(obs, "totalCost", None)
    )
    if cost is not None:
        span["cost"] = float(cost)

    # Error details
    level = getattr(obs, "level", None)
    if level and str(level).upper() == "ERROR":
        error = (
            getattr(obs, "status_message", None)
            or getattr(obs, "statusMessage", None)
        )
        if error:
            span["error"] = str(error)

    metadata = getattr(obs, "metadata", None)
    if metadata:
        span["metadata"] = metadata

    return span


def _map_observation_type(obs_type: str | None) -> str:
    if obs_type is None:
        return "custom"
    t = str(obs_type).upper()
    if t == "GENERATION":
        return "llm"
    if t == "SPAN":
        return "chain"
    if t == "EVENT":
        return "custom"
    return "custom"


def _map_level(level: Any) -> str:
    if level is None:
        return "ok"
    if str(level).upper() == "ERROR":
        return "error"
    return "ok"


def _format_ts(ts: Any) -> str | None:
    if ts is None:
        return None
    if isinstance(ts, dt.datetime):
        s = ts.isoformat()
        return s.replace("+00:00", "Z") if "+00:00" in s else s
    return str(ts)


def _unwrap_message(value: Any) -> Any:
    """Unwrap LLM message dicts to their content.

    Single message: {"role": "assistant", "content": "..."} → content string
    Message list: [{"role": "system", ...}, {"role": "user", ...}] → formatted string
    Other dicts/lists: returned as pretty-printed JSON
    """
    if isinstance(value, dict):
        if "content" in value and ("role" in value or "refusal" in value):
            return value["content"]
        # Other dicts — pretty print
        return json.dumps(value, indent=2, ensure_ascii=False)
    if isinstance(value, list):
        # Check if it's a message list
        if value and all(
            isinstance(item, dict) and "content" in item and "role" in item
            for item in value
        ):
            if len(value) == 1:
                return value[0]["content"]
            # Multi-message: format with role labels
            parts = []
            for msg in value:
                role = msg.get("role", "unknown")
                content = msg.get("content", "")
                if content:
                    parts.append(f"[{role}]\n{content}")
            return "\n\n".join(parts)
        # Other lists — pretty print
        return json.dumps(value, indent=2, ensure_ascii=False)
    return value


def _parse_ts_for_filter(ts_str: str | None) -> dt.datetime | None:
    if not ts_str:
        return None
    try:
        return dt.datetime.fromisoformat(str(ts_str).replace("Z", "+00:00"))
    except (ValueError, TypeError):
        return None
