#!/usr/bin/env python3
"""Privacy-preserving telemetry for board pipeline jobs.

The database contains operational metadata only.  In particular, model prompts,
tool arguments, model replies and error messages are deliberately not accepted
by the public write API.  One row represents one role's work on one job and is
updated idempotently when a wrapper learns additional completion metadata.
"""

import json
import os
import sqlite3
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path


SCHEMA = """
CREATE TABLE IF NOT EXISTS job_telemetry (
  id INTEGER PRIMARY KEY,
  pipeline_run_id TEXT NOT NULL,
  slug TEXT NOT NULL,
  job_id TEXT NOT NULL,
  role TEXT NOT NULL,
  configured_runtime TEXT,
  configured_provider TEXT,
  configured_model TEXT,
  configured_reasoning_effort TEXT,
  actual_runtime TEXT,
  actual_provider TEXT,
  actual_model TEXT,
  actual_reasoning_effort TEXT,
  session_id TEXT,
  input_tokens INTEGER,
  cached_input_tokens INTEGER,
  cache_creation_input_tokens INTEGER,
  output_tokens INTEGER,
  reasoning_output_tokens INTEGER,
  total_tokens INTEGER,
  cost_usd REAL,
  ttft_seconds REAL,
  model_duration_seconds REAL,
  duration_seconds REAL,
  iterations INTEGER,
  status TEXT,
  success INTEGER,
  commit_hash TEXT,
  review_decision TEXT,
  deployment_status TEXT,
  error_kind TEXT,
  recorded_at TEXT NOT NULL,
  updated_at TEXT NOT NULL,
  UNIQUE(job_id, role)
);
CREATE INDEX IF NOT EXISTS job_telemetry_run_idx
  ON job_telemetry(pipeline_run_id, recorded_at);
CREATE INDEX IF NOT EXISTS job_telemetry_slug_idx
  ON job_telemetry(slug, recorded_at DESC);
"""

PROFILE_KEYS = ("runtime", "provider", "model", "reasoning_effort")
USAGE_KEYS = (
    "input_tokens",
    "cached_input_tokens",
    "cache_creation_input_tokens",
    "output_tokens",
    "reasoning_output_tokens",
    "total_tokens",
    "cost_usd",
)
TIMING_KEYS = ("ttft_seconds", "model_duration_seconds", "duration_seconds")
OUTCOME_KEYS = (
    "status",
    "success",
    "commit_hash",
    "review_decision",
    "deployment_status",
    "error_kind",
)


def _now_iso():
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def _required_text(name, value):
    text = str(value or "").strip()
    if not text:
        raise ValueError(f"{name} must not be empty")
    if len(text) > 512:
        raise ValueError(f"{name} is too long")
    return text


def _optional_text(value, maximum=512):
    if value is None:
        return None
    text = str(value).strip()
    if not text:
        return None
    return text[:maximum]


def _nonnegative_int(value):
    if value is None or isinstance(value, bool):
        return None
    try:
        number = int(value)
    except (TypeError, ValueError):
        return None
    return number if number >= 0 else None


def _nonnegative_float(value):
    if value is None or isinstance(value, bool):
        return None
    try:
        number = float(value)
    except (TypeError, ValueError):
        return None
    return number if number >= 0 else None


def _normalise_profile(profile):
    profile = profile if isinstance(profile, dict) else {}
    return {key: _optional_text(profile.get(key), 128) for key in PROFILE_KEYS}


def _normalise_usage(usage):
    usage = usage if isinstance(usage, dict) else {}
    aliases = {
        "cached_input_tokens": ("cache_read_input_tokens",),
        "reasoning_output_tokens": ("reasoning_tokens",),
    }
    result = {}
    for key in USAGE_KEYS[:-1]:
        value = usage.get(key)
        if value is None:
            value = next(
                (usage[name] for name in aliases.get(key, ()) if usage.get(name) is not None),
                None,
            )
        result[key] = _nonnegative_int(value)
    result["cost_usd"] = _nonnegative_float(usage.get("cost_usd"))
    if result["total_tokens"] is None and (
        result["input_tokens"] is not None or result["output_tokens"] is not None
    ):
        result["total_tokens"] = (result["input_tokens"] or 0) + (
            result["output_tokens"] or 0
        )
    return result


def _normalise_timing(timing):
    timing = timing if isinstance(timing, dict) else {}
    return {key: _nonnegative_float(timing.get(key)) for key in TIMING_KEYS}


def _normalise_outcome(outcome):
    outcome = outcome if isinstance(outcome, dict) else {}
    result = {
        key: _optional_text(outcome.get(key), 256)
        for key in OUTCOME_KEYS
        if key != "success"
    }
    success = outcome.get("success")
    result["success"] = int(success) if isinstance(success, bool) else None
    return result


def _connect(db_path):
    path = Path(db_path)
    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    connection = sqlite3.connect(path, timeout=10)
    connection.row_factory = sqlite3.Row
    connection.execute("PRAGMA journal_mode=WAL")
    connection.execute("PRAGMA synchronous=NORMAL")
    connection.execute("PRAGMA busy_timeout=10000")
    connection.executescript(SCHEMA)
    try:
        os.chmod(path, 0o600)
    except OSError:
        pass
    return connection


@contextmanager
def _connection(db_path):
    connection = _connect(db_path)
    try:
        yield connection
        connection.commit()
    finally:
        connection.close()


def record_job(
    db_path,
    *,
    pipeline_run_id,
    slug,
    job_id,
    role,
    configured_profile=None,
    actual_profile=None,
    usage=None,
    timing=None,
    iterations=None,
    outcome=None,
    session_id=None,
    recorded_at=None,
):
    """Insert or enrich one ``(job_id, role)`` observation.

    ``None`` fields never erase previously recorded data, which lets a wrapper
    write configured values at launch and usage/outcome values at completion.
    """
    configured = _normalise_profile(configured_profile)
    actual = _normalise_profile(actual_profile)
    usage_values = _normalise_usage(usage)
    timing_values = _normalise_timing(timing)
    outcome_values = _normalise_outcome(outcome)
    now = _optional_text(recorded_at, 64) or _now_iso()
    values = {
        "pipeline_run_id": _required_text("pipeline_run_id", pipeline_run_id),
        "slug": _required_text("slug", slug),
        "job_id": _required_text("job_id", job_id),
        "role": _required_text("role", role),
        **{f"configured_{key}": value for key, value in configured.items()},
        **{f"actual_{key}": value for key, value in actual.items()},
        "session_id": _optional_text(session_id, 512),
        **usage_values,
        **timing_values,
        "iterations": _nonnegative_int(iterations),
        **outcome_values,
        "recorded_at": now,
        "updated_at": _now_iso(),
    }
    columns = list(values)
    mutable = [
        column
        for column in columns
        if column not in ("job_id", "role", "recorded_at")
    ]
    updates = ", ".join(
        f"{column}=COALESCE(excluded.{column}, job_telemetry.{column})"
        for column in mutable
    )
    placeholders = ", ".join(f":{column}" for column in columns)
    with _connection(db_path) as connection:
        connection.execute(
            f"""
            INSERT INTO job_telemetry ({', '.join(columns)})
            VALUES ({placeholders})
            ON CONFLICT(job_id, role) DO UPDATE SET {updates}
            """,
            values,
        )
        row = connection.execute(
            "SELECT * FROM job_telemetry WHERE job_id = ? AND role = ?",
            (values["job_id"], values["role"]),
        ).fetchone()
    return _row_to_job(row)


def delete_slugs(db_path, slugs):
    """Delete every job_telemetry row for the given slugs.

    Used by the board's bulk smoke-test cleanup so throwaway pipeline-test
    runs don't linger in telemetry once their cards are gone.
    """
    slugs = sorted({slug for slug in slugs if slug})
    if not slugs:
        return 0
    with _connection(db_path) as connection:
        placeholders = ", ".join("?" for _ in slugs)
        cursor = connection.execute(
            f"DELETE FROM job_telemetry WHERE slug IN ({placeholders})", slugs
        )
        return cursor.rowcount


def _compact(mapping):
    return {key: value for key, value in mapping.items() if value is not None}


def _row_to_job(row):
    data = dict(row)
    return {
        "pipeline_run_id": data["pipeline_run_id"],
        "slug": data["slug"],
        "job_id": data["job_id"],
        "role": data["role"],
        "configured_profile": _compact(
            {key: data[f"configured_{key}"] for key in PROFILE_KEYS}
        ),
        "actual_profile": _compact(
            {key: data[f"actual_{key}"] for key in PROFILE_KEYS}
        ),
        "session_id": data["session_id"],
        "usage": _compact({key: data[key] for key in USAGE_KEYS}),
        "timing": _compact({key: data[key] for key in TIMING_KEYS}),
        "iterations": data["iterations"],
        "outcome": _compact(
            {
                key: bool(data[key]) if key == "success" and data[key] is not None else data[key]
                for key in OUTCOME_KEYS
            }
        ),
        "recorded_at": data["recorded_at"],
        "updated_at": data["updated_at"],
    }


def _totals(jobs):
    integer_keys = (
        "input_tokens",
        "cached_input_tokens",
        "cache_creation_input_tokens",
        "output_tokens",
        "reasoning_output_tokens",
        "total_tokens",
    )
    totals = {"job_count": len(jobs)}
    for key in integer_keys:
        totals[key] = sum(job["usage"].get(key, 0) for job in jobs)
    totals["non_cached_input_tokens"] = max(
        0, totals["input_tokens"] - totals["cached_input_tokens"]
    )
    totals["cost_usd"] = round(
        sum(job["usage"].get("cost_usd", 0.0) for job in jobs), 12
    )
    totals["cost_jobs"] = sum("cost_usd" in job["usage"] for job in jobs)
    for key in TIMING_KEYS:
        totals[key] = sum(job["timing"].get(key, 0.0) for job in jobs)
    totals["iterations"] = sum(job.get("iterations") or 0 for job in jobs)
    return totals


def aggregate(db_path, *, slug=None, pipeline_run_id=None):
    """Return JSON-safe matching jobs and their aggregate counters."""
    clauses = []
    parameters = []
    if slug is not None:
        clauses.append("slug = ?")
        parameters.append(_required_text("slug", slug))
    if pipeline_run_id is not None:
        clauses.append("pipeline_run_id = ?")
        parameters.append(_required_text("pipeline_run_id", pipeline_run_id))
    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    with _connection(db_path) as connection:
        rows = connection.execute(
            f"SELECT * FROM job_telemetry {where} ORDER BY recorded_at, id",
            parameters,
        ).fetchall()
    jobs = [_row_to_job(row) for row in rows]
    return {
        "pipeline_run_id": pipeline_run_id,
        "slug": slug,
        "slugs": sorted({job["slug"] for job in jobs}),
        "jobs": jobs,
        "totals": _totals(jobs),
    }


def get_pipeline_run(db_path, pipeline_run_id):
    return aggregate(db_path, pipeline_run_id=pipeline_run_id)


def get_slug_runs(db_path, slug, *, limit=20):
    """Return the newest distinct pipeline runs containing ``slug``."""
    limit = max(1, min(_nonnegative_int(limit) or 20, 200))
    slug = _required_text("slug", slug)
    with _connection(db_path) as connection:
        rows = connection.execute(
            """
            SELECT pipeline_run_id, MAX(recorded_at) AS newest
              FROM job_telemetry
             WHERE slug = ?
             GROUP BY pipeline_run_id
             ORDER BY newest DESC
             LIMIT ?
            """,
            (slug, limit),
        ).fetchall()
    return [get_pipeline_run(db_path, row["pipeline_run_id"]) for row in rows]


def _events(output):
    if isinstance(output, dict):
        return [output]
    try:
        document = json.loads(str(output or ""))
    except (TypeError, ValueError):
        document = None
    if isinstance(document, dict):
        return [document]
    events = []
    for line in str(output or "").splitlines():
        try:
            event = json.loads(line)
        except (TypeError, ValueError):
            continue
        if isinstance(event, dict):
            events.append(event)
    return events


def _timestamp(value):
    if not isinstance(value, str):
        return None
    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp()
    except ValueError:
        return None


def _usage_from_mapping(source):
    source = source if isinstance(source, dict) else {}

    def pick(*names):
        for name in names:
            if name in source:
                return _nonnegative_int(source[name])
        return None

    return {
        "input_tokens": pick("input_tokens", "inputTokens"),
        "cached_input_tokens": pick(
            "cached_input_tokens", "cache_read_input_tokens", "cacheReadInputTokens"
        ),
        "cache_creation_input_tokens": pick(
            "cache_creation_input_tokens", "cacheCreationInputTokens"
        ),
        "output_tokens": pick("output_tokens", "outputTokens"),
        "reasoning_output_tokens": pick(
            "reasoning_output_tokens", "reasoning_tokens", "reasoningOutputTokens"
        ),
        "total_tokens": pick("total_tokens", "totalTokens"),
    }


def _finalise_parsed(session_id, profile, usage, timing, iterations, cost_usd=None):
    usage = _normalise_usage({**usage, "cost_usd": cost_usd})
    return {
        "session_id": _optional_text(session_id, 512),
        "actual_profile": _compact(_normalise_profile(profile)),
        "usage": _compact(usage),
        "timing": _compact(_normalise_timing(timing)),
        "iterations": _nonnegative_int(iterations),
    }


def parse_codex_jsonl(output):
    """Extract only operational fields from Codex CLI or rollout JSONL."""
    events = _events(output)
    session_id = None
    profile = {"runtime": "codex"}
    iterations = 0
    usage_parts = []
    cumulative_usage = None
    timestamps = []
    first_output_at = None
    for event in events:
        payload = event.get("payload") if isinstance(event.get("payload"), dict) else {}
        event_type = event.get("type")
        payload_type = payload.get("type")
        session_id = session_id or event.get("thread_id")
        if event_type == "session_meta":
            session_id = session_id or payload.get("id")
        for source in (event, payload):
            if source.get("model"):
                profile["model"] = source["model"]
            if source.get("model_provider") or source.get("provider"):
                profile["provider"] = source.get("model_provider") or source.get("provider")
            if source.get("reasoning_effort") or source.get("effort"):
                profile["reasoning_effort"] = source.get("reasoning_effort") or source.get("effort")
        stamp = _timestamp(event.get("timestamp"))
        if stamp is not None:
            timestamps.append(stamp)
        if event_type == "item.completed":
            iterations += 1
            item = event.get("item") if isinstance(event.get("item"), dict) else {}
            if item.get("type") == "agent_message" and stamp is not None and first_output_at is None:
                first_output_at = stamp
        usage = event.get("usage")
        if event_type == "turn.completed" and isinstance(usage, dict):
            usage_parts.append(_usage_from_mapping(usage))
        if payload_type == "token_count":
            info = payload.get("info") if isinstance(payload.get("info"), dict) else {}
            total = info.get("total_token_usage")
            if isinstance(total, dict):
                cumulative_usage = _usage_from_mapping(total)

    if cumulative_usage is not None:
        usage = cumulative_usage
    else:
        usage = {
            key: (
                sum(part.get(key) or 0 for part in usage_parts)
                if any(part.get(key) is not None for part in usage_parts)
                else None
            )
            for key in USAGE_KEYS
            if key != "cost_usd"
        }
    start = min(timestamps) if timestamps else None
    end = max(timestamps) if timestamps else None
    timing = {
        "ttft_seconds": first_output_at - start
        if start is not None and first_output_at is not None
        else None,
        "model_duration_seconds": end - start
        if start is not None and end is not None
        else None,
    }
    return _finalise_parsed(session_id, profile, usage, timing, iterations)


def parse_claude_output(output):
    """Extract operational metadata from Claude Code stream/result JSON."""
    events = _events(output)
    result = next((event for event in reversed(events) if event.get("type") == "result"), {})
    raw_usage = _usage_from_mapping(result.get("usage"))
    raw_input = raw_usage.get("input_tokens") or 0
    cache_read = raw_usage.get("cached_input_tokens") or 0
    cache_creation = raw_usage.get("cache_creation_input_tokens") or 0
    # Anthropic reports uncached, cache-write and cache-read input separately;
    # expose a comparable total input count while retaining both cache fields.
    raw_usage["input_tokens"] = raw_input + cache_read + cache_creation
    models = result.get("modelUsage")
    model = None
    if isinstance(models, dict) and models:
        model = next(iter(models)) if len(models) == 1 else "multiple"
    duration_ms = _nonnegative_float(result.get("duration_ms"))
    api_ms = _nonnegative_float(result.get("duration_api_ms"))
    profile = {"runtime": "claude", "provider": "anthropic", "model": model}
    timing = {
        "duration_seconds": duration_ms / 1000 if duration_ms is not None else None,
        "model_duration_seconds": api_ms / 1000 if api_ms is not None else None,
    }
    return _finalise_parsed(
        result.get("session_id"),
        profile,
        raw_usage,
        timing,
        result.get("num_turns"),
        result.get("total_cost_usd"),
    )