"""Cost estimation and cost recording.

Zero third-party dependencies; standard library only.

Actual token counts come from the model/provider response when
available: host.py collects the app-server's thread/tokenUsage/updated
events during each turn (one per upstream completion, summed over the
turn window) and the worker records them with ``source="actual"``.
When no usage events arrive (old server, turn produced no completion),
the worker falls back to ``estimate_tokens`` with
``source="estimated"``; the two are stored separately so estimates are
never mixed into actual totals.

USD is tokens times a price table, not the relay's bill: the defaults
are OpenAI list prices (see ``_DEFAULT_MODEL_PRICES``), and the relay
may charge differently. Override per model with
``CLOUD_MODEL_PRICES_JSON`` or ``CLOUD_MODEL_PRICES_FILE``; models
with no price record ``usd_micros=None`` and reports count their
tokens in ``unpriced_actual_tokens`` instead of pricing them at zero.

Dead end, documented so nobody retries it: the app-server's
account/usage/read method requires ChatGPT account authentication and
fails on a relay API key (2026-10-26 probe: ``codex account
authentication required to read token usage``). Account-level usage
is therefore not a source here; per-task actuals come from turn
events only.
"""

import json
import os
import sqlite3
from cloud.db import connect as _db_connect
import time
from datetime import datetime, timezone


# OpenAI list prices, USD per 1M tokens, as (input, cached_input,
# output). gpt-6.1-sol $2.00/$0.10/$10.00 and gpt-6-sol
# $2.00/$0.20/$10.00, as reported at the 2026-09-30 launch.
# Again: list prices, not the relay's bill; override below.
_DEFAULT_MODEL_PRICES = {
    "gpt-6.1-sol": (2.00, 0.10, 10.00),
    "gpt-6-sol": (2.00, 0.20, 10.00),
}


def _coerce_rates(rates):
    """Accept [input, cached, output] or {input, cached_input, output}."""
    try:
        if isinstance(rates, dict):
            return (float(rates["input"]), float(rates["cached_input"]),
                    float(rates["output"]))
        return (float(rates[0]), float(rates[1]), float(rates[2]))
    except (TypeError, ValueError, KeyError, IndexError):
        return None


def model_prices():
    """Effective price table: defaults merged with env overrides.

    ``CLOUD_MODEL_PRICES_JSON`` (a JSON object) and
    ``CLOUD_MODEL_PRICES_FILE`` (path to a JSON file) both merge over
    the defaults per model; unparsable content is ignored, never fatal.
    """
    prices = dict(_DEFAULT_MODEL_PRICES)
    raw = os.environ.get("CLOUD_MODEL_PRICES_JSON") or ""
    if raw:
        try:
            override = json.loads(raw)
        except ValueError:
            override = None
        if isinstance(override, dict):
            for model, rates in override.items():
                coerced = _coerce_rates(rates)
                if coerced is not None:
                    prices[model] = coerced
    path = os.environ.get("CLOUD_MODEL_PRICES_FILE") or ""
    if path:
        try:
            with open(path) as f:
                override = json.load(f)
        except (OSError, ValueError):
            override = None
        if isinstance(override, dict):
            for model, rates in override.items():
                coerced = _coerce_rates(rates)
                if coerced is not None:
                    prices[model] = coerced
    return prices


def price_usd_micros(model, input_tokens, cached_input_tokens,
                     output_tokens):
    """USD micros for provider-measured tokens at the price table.

    cached is a subset of input (prompt-cache reads): non-cached input
    pays the input rate, cached pays the cached rate, output pays the
    output rate. Reasoning tokens are already inside output_tokens
    (provider total = input + output), never added twice. Returns None
    when the model has no price entry (caller records usd as unknown).
    """
    rates = model_prices().get(model)
    if rates is None:
        return None
    input_rate, cached_rate, output_rate = rates
    total_in = max(0, int(input_tokens or 0))
    cached = max(0, min(int(cached_input_tokens or 0), total_in))
    non_cached = total_in - cached
    usd = (non_cached * input_rate + cached * cached_rate
           + max(0, int(output_tokens or 0)) * output_rate) / 1000000
    return int(round(usd * 1000000))


_USAGE_SUM_KEYS = ("input_tokens", "cached_input_tokens",
                   "cache_write_input_tokens", "output_tokens",
                   "reasoning_output_tokens", "total_tokens", "completions")


def sum_usage(a, b):
    """Add two usage dicts (host.py shape); either side may be None."""
    if not a:
        return dict(b) if b else None
    if not b:
        return dict(a)
    out = dict(a)
    for k in _USAGE_SUM_KEYS:
        if isinstance(b.get(k), int):
            out[k] = out.get(k, 0) + b[k]
    return out


def estimate_tokens(prompt_text, output_text) -> int:
    """Estimate total tokens as (chars of prompt + output) // 4."""
    prompt_text = prompt_text or ""
    output_text = output_text or ""
    return (len(prompt_text) + len(output_text)) // 4


def fetch_account_usage(usage_client) -> dict:
    """Fetch account usage through an injected ``usage_client`` callable.

    This is a call wrapper around the app-server account/usage/read
    concept: the real ``usage_client`` is injected by the control plane
    once the billing route is confirmed available. This function itself
    never runs a model turn and never touches the network.

    On success it returns ``{"available": True, "data": <raw data>}``.
    If ``usage_client`` raises, or returns None / an empty dict,
    it returns ``{"available": False, "data": None}`` and never raises
    to the caller. On a relay, usage being unreadable is a normal
    outcome (available=False), not a bug.
    """
    try:
        data = usage_client()
    except Exception:
        return {"available": False, "data": None}
    if data is None:
        return {"available": False, "data": None}
    if isinstance(data, dict) and not data:
        return {"available": False, "data": None}
    return {"available": True, "data": data}


class CostStore:
    """SQLite store for per-task / per-project / per-day cost records."""

    _SOURCES = ("actual", "estimated")

    _COST_DDL = """CREATE TABLE IF NOT EXISTS cost_records (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id TEXT,
                project_id TEXT,
                model TEXT,
                input_tokens INTEGER,
                output_tokens INTEGER,
                cached_input_tokens INTEGER DEFAULT 0,
                total_tokens INTEGER,
                usd_micros INTEGER,
                source TEXT,
                fallback_count INTEGER DEFAULT 0,
                day TEXT,
                created_at REAL
            )"""

    def _ensure_cached_column(self, conn):
        """老库补 cached_input_tokens 列（SQLite 与 Postgres 两支：

        PRAGMA table_info 经 db.py 适配层翻译，缺列时 ALTER 补列；
        并发下两个进程同时补列时第二个 ALTER 报错，直接吞掉）。
        """
        try:
            cols = {row[1] for row in conn.execute(
                "PRAGMA table_info(cost_records)").fetchall()}
        except Exception:
            return
        if "cached_input_tokens" not in cols:
            try:
                conn.execute(
                    "ALTER TABLE cost_records"
                    " ADD COLUMN cached_input_tokens INTEGER DEFAULT 0")
                conn.commit()
            except Exception:
                pass

    def __init__(self, db_path):
        self.db_path = db_path
        if db_path != ":memory:":
            parent = os.path.dirname(os.path.abspath(db_path))
            os.makedirs(parent, exist_ok=True)
        self._conn = _db_connect(db_path)
        self._conn.execute(self._COST_DDL)
        self._conn.commit()
        self._ensure_cached_column(self._conn)

    def _new_conn(self):
        """每次写/读开一个新连接，用完即关。

        旧实现在 __init__ 里建一条连接并终身共用，单线程没问题；
        worker 并发后多个线程共用同一条 sqlite 连接会直接报
        ProgrammingError（连接是线程绑定的）。其它 store（Worker
        的 _set_status、_audit 等）本来就是每次开新连接，这里
        向它们看齐。_conn 仍保留（建表 + 旧测试直接用它读），
        但 record/_report 不再走它。
        例外：db_path 为 :memory: 时新连接是另一个空库，必须
        复用 _conn（:memory: 只用于单线程单测，并发 worker
        从不用它）。
        """
        if self.db_path == ":memory:":
            return self._conn
        conn = _db_connect(self.db_path)
        conn.execute(self._COST_DDL)
        conn.commit()
        self._ensure_cached_column(conn)
        return conn

    def record(self, task_id, project_id, model, input_tokens=0,
               output_tokens=0, cached_input_tokens=0, usd_micros=None,
               source="actual", fallback_count=0):
        """Insert one cost record; returns its row id.

        ``total_tokens`` is computed as input + output, and ``day`` is
        the UTC date (YYYY-MM-DD) of ``created_at``. cached is a
        subset of input (prompt-cache reads), stored separately so
        reports can show why a task was cheap; it is never added to
        total on top of input.
        """
        if source not in self._SOURCES:
            raise ValueError(
                "source must be one of %r, got %r" % (self._SOURCES, source)
            )
        input_tokens = int(input_tokens or 0)
        output_tokens = int(output_tokens or 0)
        cached_input_tokens = int(cached_input_tokens or 0)
        total_tokens = input_tokens + output_tokens
        created_at = time.time()
        day = datetime.fromtimestamp(created_at, tz=timezone.utc).strftime("%Y-%m-%d")
        conn = self._new_conn()
        try:
            cursor = conn.execute(
                """INSERT INTO cost_records
                   (task_id, project_id, model, input_tokens, output_tokens,
                    cached_input_tokens, total_tokens, usd_micros, source,
                    fallback_count, day, created_at)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (task_id, project_id, model, input_tokens, output_tokens,
                 cached_input_tokens, total_tokens, usd_micros, source,
                 int(fallback_count or 0), day, created_at),
            )
            conn.commit()
            return cursor.lastrowid
        finally:
            # :memory: 时 conn 就是共享的 _conn，绝不能关掉它。
            if conn is not self._conn:
                conn.close()

    def record_estimated(self, task_id, project_id, model, prompt_text,
                         output_text):
        """Record an estimated cost row; returns its row id.

        The token split cannot be recovered from a character estimate,
        so ``input_tokens`` holds the whole estimate and
        ``output_tokens`` is 0; ``source`` is ``"estimated"`` and
        ``usd_micros`` is None.
        """
        total = estimate_tokens(prompt_text, output_text)
        return self.record(
            task_id, project_id, model,
            input_tokens=total, output_tokens=0,
            usd_micros=None, source="estimated", fallback_count=0,
        )

    def _report(self, where, params):
        conn = self._new_conn()
        try:
            row = conn.execute(
                """SELECT
                     COALESCE(SUM(total_tokens), 0),
                     COALESCE(SUM(usd_micros), 0),
                     COUNT(*),
                     COALESCE(SUM(fallback_count), 0),
                     COALESCE(SUM(CASE WHEN source = 'actual'
                                       THEN total_tokens ELSE 0 END), 0),
                     COALESCE(SUM(CASE WHEN source = 'estimated'
                                       THEN total_tokens ELSE 0 END), 0),
                     COALESCE(SUM(cached_input_tokens), 0),
                     COALESCE(SUM(CASE WHEN source = 'actual'
                                       AND usd_micros IS NULL
                                       THEN total_tokens ELSE 0 END), 0)
                   FROM cost_records WHERE %s""" % where,
                params,
            ).fetchone()
        finally:
            if conn is not self._conn:
                conn.close()
        (total_tokens, usd_micros, turns, fallback_count, actual,
         estimated, cached, unpriced) = row
        return {
            "total_tokens": int(total_tokens),
            "usd_micros": int(usd_micros),
            "turns": int(turns),
            "fallback_count": int(fallback_count),
            "actual_tokens": int(actual),
            "estimated_tokens": int(estimated),
            "cached_input_tokens": int(cached),
            # usd 求和只覆盖有价的行；无价 actual 行的 token 在此，
            # 不会被静默算成 0 美元。
            "unpriced_actual_tokens": int(unpriced),
        }

    def report_by_task(self, task_id) -> dict:
        """Aggregate cost records for one task."""
        return self._report("task_id = ?", (task_id,))

    def report_by_project(self, project_id) -> dict:
        """Aggregate cost records for one project."""
        return self._report("project_id = ?", (project_id,))

    def report_by_day(self, day) -> dict:
        """Aggregate cost records for one UTC day (YYYY-MM-DD)."""
        return self._report("day = ?", (day,))