"""SQLite storage layer for the plyra-trace collector."""

import json
import logging
import os
import sqlite3
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

logger = logging.getLogger("plyra_trace.collector.storage")

# ── Schema ─────────────────────────────────────────────────────────────────────

_SCHEMA_SQL = """
CREATE TABLE IF NOT EXISTS projects (
    id TEXT PRIMARY KEY,
    created_at TEXT NOT NULL,
    last_seen_at TEXT NOT NULL,
    span_count INTEGER DEFAULT 0
);

CREATE TABLE IF NOT EXISTS traces (
    trace_id TEXT PRIMARY KEY,
    project_id TEXT NOT NULL,
    root_name TEXT NOT NULL,
    root_span_id TEXT NOT NULL,
    started_at TEXT NOT NULL,
    ended_at TEXT,
    duration_ms INTEGER,
    span_count INTEGER DEFAULT 0,
    llm_cost_usd REAL DEFAULT 0.0,
    llm_token_total INTEGER DEFAULT 0,
    llm_cache_hits INTEGER DEFAULT 0,
    llm_cache_total INTEGER DEFAULT 0,
    has_guard_block INTEGER DEFAULT 0,
    has_error INTEGER DEFAULT 0,
    status TEXT DEFAULT 'ok',
    session_id TEXT,
    user_id TEXT
);

CREATE TABLE IF NOT EXISTS spans (
    span_id TEXT PRIMARY KEY,
    trace_id TEXT NOT NULL,
    parent_id TEXT,
    project_id TEXT NOT NULL,
    name TEXT NOT NULL,
    kind TEXT NOT NULL,
    started_at TEXT NOT NULL,
    ended_at TEXT NOT NULL,
    duration_ms INTEGER NOT NULL,
    attributes TEXT NOT NULL,
    events TEXT DEFAULT '[]',
    status TEXT DEFAULT 'ok'
);

CREATE INDEX IF NOT EXISTS idx_spans_trace_id ON spans(trace_id);
CREATE INDEX IF NOT EXISTS idx_traces_project_started ON traces(project_id, started_at DESC);
CREATE INDEX IF NOT EXISTS idx_spans_kind ON spans(kind);
CREATE INDEX IF NOT EXISTS idx_spans_project ON spans(project_id);
"""


def open_db(db_path: str) -> sqlite3.Connection:
    """
    Open (or create) the SQLite database, applying schema migrations.

    Args:
        db_path: Path to the database file. '~' is expanded.

    Returns:
        sqlite3.Connection with row_factory set to sqlite3.Row
    """
    expanded = os.path.expanduser(db_path)
    Path(expanded).parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(expanded, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.executescript(_SCHEMA_SQL)
    conn.commit()
    return conn


def _now_iso() -> str:
    return datetime.now(UTC).isoformat()


def upsert_project(db: sqlite3.Connection, project_id: str) -> None:
    """
    Insert or update a project record.

    Args:
        db: Open database connection
        project_id: Project identifier (service.name)
    """
    now = _now_iso()
    db.execute(
        """
        INSERT INTO projects (id, created_at, last_seen_at, span_count)
        VALUES (?, ?, ?, 0)
        ON CONFLICT(id) DO UPDATE SET last_seen_at = excluded.last_seen_at
        """,
        (project_id, now, now),
    )
    db.commit()


def insert_span(db: sqlite3.Connection, span: dict[str, Any]) -> None:
    """
    Insert a single span record. Upserts to handle duplicates.

    Args:
        db: Open database connection
        span: Span dict with keys:
            span_id, trace_id, parent_id, project_id, name, kind,
            started_at, ended_at, duration_ms, attributes (dict), events (list),
            status
    """
    db.execute(
        """
        INSERT OR REPLACE INTO spans
            (span_id, trace_id, parent_id, project_id, name, kind,
             started_at, ended_at, duration_ms, attributes, events, status)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        (
            span["span_id"],
            span["trace_id"],
            span.get("parent_id"),
            span["project_id"],
            span["name"],
            span["kind"],
            span["started_at"],
            span["ended_at"],
            span["duration_ms"],
            json.dumps(span.get("attributes", {})),
            json.dumps(span.get("events", [])),
            span.get("status", "ok"),
        ),
    )
    # update project span count
    db.execute(
        "UPDATE projects SET span_count = span_count + 1, last_seen_at = ? WHERE id = ?",
        (_now_iso(), span["project_id"]),
    )
    db.commit()


def update_trace_summary(db: sqlite3.Connection, trace_id: str) -> None:
    """
    Recalculate and upsert the traces table entry for a given trace_id.
    Should be called after inserting all spans for a trace.

    Args:
        db: Open database connection
        trace_id: Trace ID to summarise
    """
    rows = db.execute(
        "SELECT * FROM spans WHERE trace_id = ? ORDER BY started_at ASC",
        (trace_id,),
    ).fetchall()

    if not rows:
        return

    # Find root span (no parent_id, or parent_id not in this trace)
    span_ids = {r["span_id"] for r in rows}
    root = None
    for row in rows:
        if not row["parent_id"] or row["parent_id"] not in span_ids:
            root = row
            break
    if root is None:
        root = rows[0]

    total_cost = 0.0
    total_tokens = 0
    cache_hits = 0
    cache_total = 0
    has_guard_block = 0
    has_error = 0
    status = "ok"
    session_id = None
    user_id = None

    for row in rows:
        attrs: dict[str, Any] = json.loads(row["attributes"])
        # LLM cost
        cost = attrs.get("llm.cost_usd") or attrs.get("plyra.llm.cost_usd")
        if cost:
            try:
                total_cost += float(cost)
            except (TypeError, ValueError):
                pass
        # Tokens
        tokens = attrs.get("llm.token_count.total")
        if tokens:
            try:
                total_tokens += int(tokens)
            except (TypeError, ValueError):
                pass
        # Cache
        kind = row["kind"]
        if kind == "LLM":
            cache_total += 1
            if attrs.get("llm.cache_hit") in (True, "true", "True", 1, "1"):
                cache_hits += 1
        # Guard blocks
        guard_action = attrs.get("guard.action") or attrs.get("plyra.guard.action")
        if guard_action == "block":
            has_guard_block = 1
        # Errors
        if row["status"] == "error":
            has_error = 1
            status = "error"
        # Session/user from root or first occurrence
        if not session_id:
            session_id = attrs.get("session.id")
        if not user_id:
            user_id = attrs.get("user.id")

    root_attrs: dict[str, Any] = json.loads(root["attributes"])
    if not session_id:
        session_id = root_attrs.get("session.id")
    if not user_id:
        user_id = root_attrs.get("user.id")

    project_id = root["project_id"]
    started_at = root["started_at"]

    # Compute trace end/duration from the last-ending span
    ended_at = max((r["ended_at"] for r in rows), default=None)
    duration_ms = None
    if ended_at:
        try:
            start = datetime.fromisoformat(started_at.replace("Z", "+00:00"))
            end = datetime.fromisoformat(ended_at.replace("Z", "+00:00"))
            duration_ms = int((end - start).total_seconds() * 1000)
        except Exception:
            pass

    db.execute(
        """
        INSERT OR REPLACE INTO traces
            (trace_id, project_id, root_name, root_span_id, started_at, ended_at,
             duration_ms, span_count, llm_cost_usd, llm_token_total,
             llm_cache_hits, llm_cache_total, has_guard_block, has_error,
             status, session_id, user_id)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        (
            trace_id,
            project_id,
            root["name"],
            root["span_id"],
            started_at,
            ended_at,
            duration_ms,
            len(rows),
            round(total_cost, 6),
            total_tokens,
            cache_hits,
            cache_total,
            has_guard_block,
            has_error,
            status,
            session_id,
            user_id,
        ),
    )
    db.commit()


def get_traces(
    db: sqlite3.Connection,
    project_id: str | None = None,
    limit: int = 50,
    offset: int = 0,
    since_hours: int | None = None,
) -> tuple[list[dict[str, Any]], int]:
    """
    Retrieve trace summaries.

    Args:
        db: Open database connection
        project_id: Filter by project (None = all projects)
        limit: Max results
        offset: Pagination offset
        since_hours: Only return traces from the last N hours

    Returns:
        (list of trace dicts, total count)
    """
    conditions: list[str] = []
    params: list[Any] = []

    if project_id:
        conditions.append("project_id = ?")
        params.append(project_id)

    if since_hours is not None:
        cutoff = datetime.now(UTC)
        from datetime import timedelta

        cutoff = cutoff - timedelta(hours=since_hours)
        conditions.append("started_at >= ?")
        params.append(cutoff.isoformat())

    where = ("WHERE " + " AND ".join(conditions)) if conditions else ""

    total_row = db.execute(f"SELECT COUNT(*) as cnt FROM traces {where}", params).fetchone()
    total = total_row["cnt"] if total_row else 0

    rows = db.execute(
        f"SELECT * FROM traces {where} ORDER BY started_at DESC LIMIT ? OFFSET ?",
        params + [limit, offset],
    ).fetchall()

    return [dict(r) for r in rows], total


def get_spans(
    db: sqlite3.Connection,
    trace_id: str,
) -> list[dict[str, Any]]:
    """
    Retrieve all spans for a trace.

    Args:
        db: Open database connection
        trace_id: Trace to fetch spans for

    Returns:
        List of span dicts with attributes parsed from JSON
    """
    rows = db.execute(
        "SELECT * FROM spans WHERE trace_id = ? ORDER BY started_at ASC",
        (trace_id,),
    ).fetchall()
    result = []
    for r in rows:
        d = dict(r)
        d["attributes"] = json.loads(d["attributes"])
        d["events"] = json.loads(d["events"])
        result.append(d)
    return result


def get_stats_llm(
    db: sqlite3.Connection,
    project_id: str | None = None,
    since_hours: int = 24,
) -> dict[str, Any]:
    """
    Compute LLM cost, token, and cache statistics.

    Args:
        db: Open database connection
        project_id: Filter by project (None = all)
        since_hours: Time window in hours

    Returns:
        Dict with cost/token/cache statistics
    """
    from datetime import timedelta

    cutoff = (datetime.now(UTC) - timedelta(hours=since_hours)).isoformat()

    cond = "WHERE s.started_at >= ?"
    params: list[Any] = [cutoff]
    if project_id:
        cond += " AND s.project_id = ?"
        params.append(project_id)

    rows = db.execute(
        f"SELECT s.attributes, s.kind FROM spans s {cond}",
        params,
    ).fetchall()

    total_cost = 0.0
    total_calls = 0
    cache_hits = 0
    cost_by_agent: dict[str, float] = {}
    cost_by_model: dict[str, dict[str, Any]] = {}
    token_by_hour: dict[str, dict[str, int]] = {f"{h:02d}": {"prompt": 0, "completion": 0} for h in range(24)}

    for row in rows:
        attrs: dict[str, Any] = json.loads(row["attributes"])
        kind = row["kind"]
        if kind != "LLM":
            continue

        total_calls += 1
        cost = attrs.get("llm.cost_usd") or attrs.get("plyra.llm.cost_usd") or 0.0
        try:
            cost = float(cost)
        except (TypeError, ValueError):
            cost = 0.0
        total_cost += cost

        if attrs.get("llm.cache_hit") in (True, "true", "True", 1, "1"):
            cache_hits += 1

        agent = attrs.get("agent.name", "unknown")
        cost_by_agent[agent] = cost_by_agent.get(agent, 0.0) + cost

        model = attrs.get("llm.model_name", "unknown")
        if model not in cost_by_model:
            cost_by_model[model] = {
                "model_name": model,
                "cost_usd": 0.0,
                "call_count": 0,
            }
        cost_by_model[model]["cost_usd"] += cost
        cost_by_model[model]["call_count"] += 1

    # Token by hour requires started_at
    rows2 = db.execute(
        f"SELECT s.attributes, s.kind, s.started_at FROM spans s {cond} AND s.kind = 'LLM'",
        params,
    ).fetchall()
    for row in rows2:
        attrs = json.loads(row["attributes"])
        try:
            hour = row["started_at"][11:13]  # e.g. "10" from ISO string
        except Exception:
            continue
        if hour in token_by_hour:
            token_by_hour[hour]["prompt"] += int(attrs.get("llm.token_count.prompt") or 0)
            token_by_hour[hour]["completion"] += int(attrs.get("llm.token_count.completion") or 0)

    cache_hit_rate = cache_hits / total_calls if total_calls else 0.0

    return {
        "total_cost_usd": round(total_cost, 6),
        "total_calls": total_calls,
        "cache_hit_rate": round(cache_hit_rate, 4),
        "cache_hits": cache_hits,
        "cost_by_agent": [
            {"agent_name": k, "cost_usd": round(v, 6)} for k, v in sorted(cost_by_agent.items(), key=lambda x: -x[1])
        ],
        "cost_by_model": sorted(cost_by_model.values(), key=lambda x: -x["cost_usd"]),
        "token_by_hour": [{"hour": h, **token_by_hour[h]} for h in sorted(token_by_hour)],
    }


def get_stats_agents(
    db: sqlite3.Connection,
    project_id: str | None = None,
    since_hours: int = 24,
) -> list[dict[str, Any]]:
    """
    Compute per-agent performance statistics.

    Args:
        db: Open database connection
        project_id: Filter by project (None = all)
        since_hours: Time window in hours

    Returns:
        List of agent stat dicts
    """
    from datetime import timedelta

    cutoff = (datetime.now(UTC) - timedelta(hours=since_hours)).isoformat()

    cond = "WHERE started_at >= ? AND kind = 'AGENT'"
    params: list[Any] = [cutoff]
    if project_id:
        cond += " AND project_id = ?"
        params.append(project_id)

    rows = db.execute(
        f"SELECT name, duration_ms, trace_id, attributes FROM spans {cond}",
        params,
    ).fetchall()

    agents: dict[str, dict[str, Any]] = {}
    for row in rows:
        name = row["name"]
        attrs: dict[str, Any] = json.loads(row["attributes"])
        agent_name = attrs.get("agent.name", name)
        if agent_name not in agents:
            agents[agent_name] = {
                "name": agent_name,
                "durations": [],
                "llm_calls": 0,
                "cache_hits": 0,
                "tools_used": set(),
            }
        agents[agent_name]["durations"].append(row["duration_ms"] or 0)

    # Gather tool and LLM info per trace for each agent
    for agent_name, data in agents.items():
        durs = sorted(data["durations"])
        n = len(durs)
        avg = sum(durs) / n if n else 0
        p95 = durs[int(n * 0.95)] if n > 1 else (durs[0] if durs else 0)
        data["invocations"] = n
        data["avg_duration_ms"] = round(avg, 1)
        data["p95_duration_ms"] = float(p95)
        data["tools_used"] = list(data["tools_used"])
        llm_calls = data.get("llm_calls", 0)
        cache_hits = data.get("cache_hits", 0)
        data["cache_hit_rate"] = round(cache_hits / llm_calls, 4) if llm_calls else 0.0
        del data["durations"]

    return sorted(agents.values(), key=lambda x: -x["avg_duration_ms"])


def get_stats_guards(
    db: sqlite3.Connection,
    project_id: str | None = None,
    since_hours: int = 24,
) -> dict[str, Any]:
    """
    Retrieve GUARD span events for the audit log.

    Args:
        db: Open database connection
        project_id: Filter by project (None = all)
        since_hours: Time window in hours

    Returns:
        Dict with total/allowed/blocked counts and events list
    """
    from datetime import timedelta

    cutoff = (datetime.now(UTC) - timedelta(hours=since_hours)).isoformat()

    cond = "WHERE started_at >= ? AND kind = 'GUARD'"
    params: list[Any] = [cutoff]
    if project_id:
        cond += " AND project_id = ?"
        params.append(project_id)

    rows = db.execute(
        f"SELECT span_id, trace_id, name, started_at, attributes FROM spans {cond} ORDER BY started_at DESC",
        params,
    ).fetchall()

    total = len(rows)
    allowed = 0
    blocked = 0
    events = []

    for row in rows:
        attrs: dict[str, Any] = json.loads(row["attributes"])
        action = attrs.get("guard.action") or attrs.get("plyra.guard.action") or "allow"
        if action == "block":
            blocked += 1
        else:
            allowed += 1

        input_val = attrs.get("input.value", "")
        if isinstance(input_val, str) and len(input_val) > 60:
            input_preview = input_val[:60] + "…"
        else:
            input_preview = str(input_val)[:60]

        # Determine tool category from tool name
        tool_name = attrs.get("tool.name") or row["name"]
        tool_cat = _categorize_tool(tool_name)

        events.append(
            {
                "time": row["started_at"],
                "tool_name": tool_name,
                "tool_category": tool_cat,
                "action": action,
                "agent_name": attrs.get("agent.name", "unknown"),
                "input_preview": input_preview,
                "confidence": attrs.get("guard.confidence") or attrs.get("plyra.guard.confidence"),
                "trace_id": row["trace_id"],
                "policy": attrs.get("guard.policy") or attrs.get("plyra.guard.policy") or "default",
            }
        )

    return {
        "total": total,
        "allowed": allowed,
        "blocked": blocked,
        "events": events,
    }


def _categorize_tool(tool_name: str) -> str:
    """Heuristic tool category from name."""
    name_lower = tool_name.lower()
    write_keywords = (
        "create",
        "write",
        "send",
        "post",
        "update",
        "delete",
        "draft",
        "reply",
    )
    communicate_keywords = ("email", "message", "notify", "slack", "chat")
    if any(k in name_lower for k in communicate_keywords):
        return "communicate"
    if any(k in name_lower for k in write_keywords):
        return "write"
    return "read"


def get_projects(db: sqlite3.Connection) -> list[dict[str, Any]]:
    """
    Retrieve all project summaries.

    Args:
        db: Open database connection

    Returns:
        List of project dicts
    """
    rows = db.execute("SELECT id, span_count, last_seen_at FROM projects ORDER BY last_seen_at DESC").fetchall()
    return [dict(r) for r in rows]


def get_projects_with_stats(db: sqlite3.Connection) -> list[dict[str, Any]]:
    """
    Return each project with aggregated trace statistics.

    Returns:
        List of dicts with: id, span_count, last_seen_at, trace_count, total_cost, last_seen
    """
    rows = db.execute(
        """
        SELECT
            p.id,
            p.span_count,
            p.last_seen_at,
            COUNT(DISTINCT t.trace_id) AS trace_count,
            COALESCE(SUM(t.llm_cost_usd), 0) AS total_cost,
            MAX(t.started_at) AS last_seen
        FROM projects p
        LEFT JOIN traces t ON t.project_id = p.id
        GROUP BY p.id
        ORDER BY last_seen DESC
        """
    ).fetchall()
    return [dict(r) for r in rows]


def get_db_meta(db: sqlite3.Connection, db_path: str) -> dict[str, Any]:
    """
    Return database metadata: size, trace count, span count, oldest trace.

    Args:
        db: Open database connection
        db_path: Path to the db file (for size calculation)

    Returns:
        Dict with size_mb, trace_count, span_count, oldest_trace_at
    """
    expanded = os.path.expanduser(db_path)
    try:
        size_bytes = os.path.getsize(expanded)
        size_mb = round(size_bytes / (1024 * 1024), 2)
    except OSError:
        size_mb = 0.0

    trace_count = db.execute("SELECT COUNT(*) as c FROM traces").fetchone()["c"]
    span_count = db.execute("SELECT COUNT(*) as c FROM spans").fetchone()["c"]
    oldest_row = db.execute("SELECT MIN(started_at) as t FROM traces").fetchone()
    oldest = oldest_row["t"] if oldest_row else None

    return {
        "size_mb": size_mb,
        "trace_count": trace_count,
        "span_count": span_count,
        "oldest_trace_at": oldest,
    }
