"""Tests for StoreBase derived methods using a mock store."""

from __future__ import annotations

import json
import os
from datetime import datetime, timedelta, timezone
from typing import Any, Iterator

from bandito.models import Trace
from bandito.storage import LocalStore
from bandito.storage.base import StoreBase, TraceStats


def _make_trace(trace_id: str, project: str = "bot", started_at: str | None = None) -> dict:
    now = datetime.now(timezone.utc)
    return {
        "trace_id": trace_id,
        "project": project,
        "input": f"input-{trace_id}",
        "output": f"output-{trace_id}",
        "started_at": started_at or now.isoformat().replace("+00:00", "Z"),
        "ended_at": (now + timedelta(seconds=1)).isoformat().replace("+00:00", "Z"),
        "status": "ok",
        "spans": [{
            "span_id": f"sp-{trace_id}",
            "name": "llm",
            "kind": "llm",
            "started_at": now.isoformat().replace("+00:00", "Z"),
            "ended_at": (now + timedelta(seconds=1)).isoformat().replace("+00:00", "Z"),
            "status": "ok",
            "model": "gpt-4o",
            "cost": 0.01,
            "input_tokens": 100,
            "output_tokens": 50,
        }],
    }


def _setup_store(tmp_path, traces=None) -> LocalStore:
    """Helper to create a store with test data in project subdirectories."""
    (tmp_path / "traces").mkdir(parents=True, exist_ok=True)

    for t in (traces or []):
        project = t.get("project", "default")
        project_dir = tmp_path / "traces" / project
        project_dir.mkdir(parents=True, exist_ok=True)
        with open(project_dir / f"{t['trace_id']}.json", "w") as f:
            json.dump(t, f)

    return LocalStore(str(tmp_path))


class TestListProjects:
    def test_extracts_unique_projects(self, tmp_path):
        traces = [
            _make_trace("t-1", project="bot-a"),
            _make_trace("t-2", project="bot-b"),
            _make_trace("t-3", project="bot-a"),
        ]
        store = _setup_store(tmp_path, traces=traces)
        assert store.list_projects() == ["bot-a", "bot-b"]

    def test_empty(self, tmp_path):
        store = _setup_store(tmp_path)
        assert store.list_projects() == []


class TestListTraces:
    def test_filters_by_project(self, tmp_path):
        traces = [
            _make_trace("t-1", project="a"),
            _make_trace("t-2", project="b"),
            _make_trace("t-3", project="a"),
        ]
        store = _setup_store(tmp_path, traces=traces)
        result = store.list_traces("a")
        assert len(result) == 2
        assert all(t["project"] == "a" for t in result)

    def test_returns_all_when_no_project(self, tmp_path):
        traces = [_make_trace(f"t-{i}") for i in range(5)]
        store = _setup_store(tmp_path, traces=traces)
        assert len(store.list_traces()) == 5

    def test_sorted_by_started_at_desc(self, tmp_path):
        traces = [
            _make_trace("t-old", started_at="2026-01-01T00:00:00Z"),
            _make_trace("t-new", started_at="2026-03-01T00:00:00Z"),
            _make_trace("t-mid", started_at="2026-02-01T00:00:00Z"),
        ]
        store = _setup_store(tmp_path, traces=traces)
        result = store.list_traces()
        ids = [t["trace_id"] for t in result]
        assert ids == ["t-new", "t-mid", "t-old"]


class TestCountTraces:
    def test_counts_all_traces(self, tmp_path):
        traces = [_make_trace(f"t-{i}") for i in range(10)]
        store = _setup_store(tmp_path, traces=traces)
        assert store.count_traces() == 10

    def test_counts_by_project(self, tmp_path):
        traces = [
            _make_trace("t-1", project="a"),
            _make_trace("t-2", project="b"),
            _make_trace("t-3", project="a"),
        ]
        store = _setup_store(tmp_path, traces=traces)
        assert store.count_traces("a") == 2
        assert store.count_traces("b") == 1

    def test_empty_project(self, tmp_path):
        store = _setup_store(tmp_path)
        assert store.count_traces("nonexistent") == 0


class TestListTracesFiltered:
    def test_limit(self, tmp_path):
        traces = [_make_trace(f"t-{i}") for i in range(10)]
        store = _setup_store(tmp_path, traces=traces)
        assert len(store.list_traces_filtered(limit=3)) == 3

    def test_since(self, tmp_path):
        traces = [
            _make_trace("t-old", started_at="2026-01-01T00:00:00Z"),
            _make_trace("t-new", started_at="2026-03-01T00:00:00Z"),
        ]
        store = _setup_store(tmp_path, traces=traces)
        since = datetime(2026, 2, 1, tzinfo=timezone.utc)
        result = store.list_traces_filtered(since=since)
        assert len(result) == 1
        assert result[0]["trace_id"] == "t-new"


class TestGetStats:
    def test_basic_stats(self, tmp_path):
        traces = [_make_trace(f"t-{i}") for i in range(5)]
        store = _setup_store(tmp_path, traces=traces)
        stats = store.get_stats("bot")

        assert stats.total_traces == 5
        assert stats.total_cost > 0
        assert "llm" in stats.span_kinds

    def test_empty_project(self, tmp_path):
        store = _setup_store(tmp_path)
        stats = store.get_stats("nonexistent")
        assert stats.total_traces == 0
