# -*- coding: utf-8 -*-
"""SQLite store of the AI usage history: every call, its cost and the summaries the "Ceny AI" dialog shows.

The only module that runs SQL of the usage history. One connection per call (the plugin, the service and the game all
write), WAL journal, every write one transaction.

`calls` holds one row per AI call (a failed call too, with its error). Rows older than RETENTION_MONTHS are rolled up
into `daily` (one row per day, provider, model and purpose, kept for ever), so the totals stay right while the detail
shrinks. Every summary reads the view `usage`, which is both tables together; every delete removes from both.
"""
import csv
import datetime
import json
import os
import sqlite3
from collections import defaultdict
from contextlib import closing, contextmanager

SCHEMA_VERSION = 1
DB_NAME = 'ai_usage.db'
LEGACY_JSON = 'ai_usage_stats.json'
RETENTION_MONTHS = 12
LEGACY_PURPOSE = 'legacy'
_override_path = None

SCHEMA = '''
CREATE TABLE calls(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  ts TEXT NOT NULL, day TEXT NOT NULL,
  provider TEXT NOT NULL, model TEXT NOT NULL, purpose TEXT NOT NULL,
  session_id TEXT NOT NULL DEFAULT '', title TEXT NOT NULL DEFAULT '',
  status TEXT NOT NULL DEFAULT 'ok', error TEXT NOT NULL DEFAULT '',
  latency_ms INTEGER NOT NULL DEFAULT 0, requests INTEGER NOT NULL DEFAULT 1,
  input_tokens INTEGER NOT NULL DEFAULT 0, cached_tokens INTEGER NOT NULL DEFAULT 0, output_tokens INTEGER NOT NULL DEFAULT 0,
  input_cost REAL NOT NULL DEFAULT 0, cached_cost REAL NOT NULL DEFAULT 0, output_cost REAL NOT NULL DEFAULT 0,
  cost REAL NOT NULL DEFAULT 0, savings REAL NOT NULL DEFAULT 0);
CREATE INDEX calls_day ON calls(day);
CREATE INDEX calls_session ON calls(session_id);
CREATE TABLE daily(
  day TEXT NOT NULL, provider TEXT NOT NULL, model TEXT NOT NULL, purpose TEXT NOT NULL,
  requests INTEGER NOT NULL DEFAULT 0, errors INTEGER NOT NULL DEFAULT 0,
  input_tokens INTEGER NOT NULL DEFAULT 0, cached_tokens INTEGER NOT NULL DEFAULT 0, output_tokens INTEGER NOT NULL DEFAULT 0,
  cost REAL NOT NULL DEFAULT 0, savings REAL NOT NULL DEFAULT 0,
  PRIMARY KEY(day, provider, model, purpose));
CREATE TABLE meta(key TEXT PRIMARY KEY, value TEXT);
CREATE VIEW usage AS
  SELECT day, provider, model, purpose, requests, CASE WHEN status = 'error' THEN requests ELSE 0 END AS errors,
         input_tokens, cached_tokens, output_tokens, cost, savings FROM calls
  UNION ALL
  SELECT day, provider, model, purpose, requests, errors, input_tokens, cached_tokens, output_tokens, cost, savings
  FROM daily;
'''

_SUMS = ('SUM(requests) AS requests, SUM(errors) AS errors, SUM(input_tokens) AS input_tokens, '
         'SUM(cached_tokens) AS cached_tokens, SUM(output_tokens) AS output_tokens, SUM(cost) AS cost, SUM(savings) AS savings')


def set_path(path):
    """Tests point the store at a temporary file."""
    global _override_path
    _override_path = path


def _profile_dir():
    import xbmcvfs
    return xbmcvfs.translatePath('special://profile/addon_data/plugin.video.milionar/')


def db_path():
    return _override_path or os.path.join(_profile_dir(), DB_NAME)


def legacy_json_path():
    return os.path.join(os.path.dirname(db_path()), LEGACY_JSON)


@contextmanager
def _db():
    """A connection with the schema ready; commits on success, rolls back on an error."""
    path = db_path()
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with closing(sqlite3.connect(path, timeout=5)) as con:
        con.row_factory = sqlite3.Row
        con.execute('PRAGMA journal_mode=WAL')
        _ensure_schema(con)
        try:
            yield con
            con.commit()
        except Exception:
            con.rollback()
            raise


def _ensure_schema(con):
    if con.execute('PRAGMA user_version').fetchone()[0] >= SCHEMA_VERSION:
        return
    con.executescript(SCHEMA)
    _import_legacy(con)
    con.execute('PRAGMA user_version = {}'.format(SCHEMA_VERSION))
    con.commit()


def _today():
    return datetime.date.today()


def _ranges(since, until):
    """WHERE clause and parameters for an inclusive range of days (YYYY-MM-DD); None leaves a side open."""
    clauses, params = [], []
    if since:
        clauses.append('day >= ?')
        params.append(since)
    if until:
        clauses.append('day <= ?')
        params.append(until)
    return (' WHERE ' + ' AND '.join(clauses)) if clauses else '', params


# ---- writing
def record_call(provider, model, purpose, session_id='', title='', status='ok', error='', latency_ms=0, requests=1,
                input_tokens=0, cached_tokens=0, output_tokens=0, input_cost=0.0, cached_cost=0.0, output_cost=0.0,
                savings=0.0, now=None):
    """Stores one AI call; returns its id."""
    now = now or datetime.datetime.now()
    cost = input_cost + cached_cost + output_cost
    with _db() as con:
        cursor = con.execute(
            'INSERT INTO calls(ts, day, provider, model, purpose, session_id, title, status, error, latency_ms, requests, '
            'input_tokens, cached_tokens, output_tokens, input_cost, cached_cost, output_cost, cost, savings) '
            'VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)',
            (now.isoformat(), now.date().isoformat(), provider, model, purpose, session_id or '', title or '', status,
             (error or '')[:300], int(latency_ms or 0), requests, input_tokens, cached_tokens, output_tokens,
             input_cost, cached_cost, output_cost, cost, savings))
        return cursor.lastrowid


def prune(today=None, months=RETENTION_MONTHS):
    """Rolls the calls older than `months` months up into `daily` and removes them; returns how many were rolled up."""
    today = today or _today()
    month_index = today.year * 12 + today.month - 1 - months
    cutoff = datetime.date(month_index // 12, month_index % 12 + 1, 1).isoformat()
    with _db() as con:
        con.execute(
            'INSERT INTO daily(day, provider, model, purpose, requests, errors, input_tokens, cached_tokens, output_tokens, cost, savings) '
            'SELECT day, provider, model, purpose, SUM(requests), SUM(CASE WHEN status = \'error\' THEN requests ELSE 0 END), '
            'SUM(input_tokens), SUM(cached_tokens), SUM(output_tokens), SUM(cost), SUM(savings) '
            'FROM calls WHERE day < ? GROUP BY day, provider, model, purpose '
            'ON CONFLICT(day, provider, model, purpose) DO UPDATE SET requests = requests + excluded.requests, '
            'errors = errors + excluded.errors, input_tokens = input_tokens + excluded.input_tokens, '
            'cached_tokens = cached_tokens + excluded.cached_tokens, output_tokens = output_tokens + excluded.output_tokens, '
            'cost = cost + excluded.cost, savings = savings + excluded.savings', (cutoff,))
        return con.execute('DELETE FROM calls WHERE day < ?', (cutoff,)).rowcount


def prune_if_due(today=None):
    """prune() at most once a day."""
    today = today or _today()
    with _db() as con:
        row = con.execute("SELECT value FROM meta WHERE key = 'pruned'").fetchone()
        if row and row['value'] == today.isoformat():
            return 0
        con.execute("INSERT OR REPLACE INTO meta(key, value) VALUES('pruned', ?)", (today.isoformat(),))
    return prune(today)


# ---- summaries
def _row(row):
    return {key: (row[key] or 0) if key not in ('provider', 'model', 'purpose', 'day') else row[key] for key in row.keys()}


def totals(since=None, until=None):
    where, params = _ranges(since, until)
    with _db() as con:
        return _row(con.execute('SELECT {} FROM usage{}'.format(_SUMS, where), params).fetchone())


def per_day(since, until):
    """One entry for every day of the range, empty days included (a chart wants all of them)."""
    where, params = _ranges(since, until)
    with _db() as con:
        found = {row['day']: _row(row) for row in con.execute(
            'SELECT day, {} FROM usage{} GROUP BY day'.format(_SUMS, where), params)}
    first, last = datetime.date.fromisoformat(since), datetime.date.fromisoformat(until)
    days = []
    for offset in range((last - first).days + 1):
        day = (first + datetime.timedelta(days=offset)).isoformat()
        days.append(found.get(day) or dict(day=day, requests=0, errors=0, input_tokens=0, cached_tokens=0,
                                           output_tokens=0, cost=0.0, savings=0.0))
    return days


def per_month(limit=12):
    with _db() as con:
        rows = con.execute('SELECT substr(day, 1, 7) AS month, {} FROM usage GROUP BY month ORDER BY month DESC LIMIT ?'.format(_SUMS),
                           (limit,)).fetchall()
    return [{key: row[key] for key in row.keys()} for row in rows]


def _grouped(columns, since, until, order='cost DESC'):
    where, params = _ranges(since, until)
    with _db() as con:
        return [_row(row) for row in con.execute(
            'SELECT {cols}, {sums} FROM usage{where} GROUP BY {cols} ORDER BY {order}'.format(
                cols=columns, sums=_SUMS, where=where, order=order), params)]


def per_provider(since=None, until=None):
    return _grouped('provider', since, until)


def per_model(since=None, until=None):
    return _grouped('provider, model', since, until)


def per_purpose(since=None, until=None):
    return _grouped('purpose', since, until)


def sessions(since=None, until=None, limit=300):
    """The history: a game (all calls of one session) or a single call, newest first."""
    where, params = _ranges(since, until)
    with _db() as con:
        rows = con.execute(
            "SELECT CASE WHEN session_id = '' THEN 'c:' || id ELSE 's:' || session_id END AS key, "
            "MAX(ts) AS last_ts, MIN(ts) AS first_ts, MAX(title) AS title, MAX(purpose) AS purpose, MAX(provider) AS provider, "
            "MAX(model) AS model, SUM(requests) AS requests, SUM(CASE WHEN status = 'error' THEN requests ELSE 0 END) AS errors, "
            "SUM(input_tokens) AS input_tokens, SUM(cached_tokens) AS cached_tokens, SUM(output_tokens) AS output_tokens, "
            "SUM(cost) AS cost, SUM(savings) AS savings, SUM(latency_ms) AS latency_ms "
            "FROM calls{} GROUP BY key ORDER BY last_ts DESC LIMIT ?".format(where), params + [limit]).fetchall()
    return [{key: row[key] for key in row.keys()} for row in rows]


def calls_of(key):
    """The calls of a history entry (`sessions()` key), oldest first."""
    with _db() as con:
        if key.startswith('s:'):
            rows = con.execute('SELECT * FROM calls WHERE session_id = ? ORDER BY ts', (key[2:],))
        else:
            rows = con.execute('SELECT * FROM calls WHERE id = ?', (int(key[2:]),))
        return [{k: row[k] for k in row.keys()} for row in rows]


def span():
    """(first day, last day) with any usage, or None."""
    with _db() as con:
        row = con.execute('SELECT MIN(day) AS first, MAX(day) AS last FROM usage').fetchone()
    return (row['first'], row['last']) if row['first'] else None


# ---- deleting; a scope is (kind, value): entry (a `sessions()` key), call (id), day, month (YYYY-MM), provider, all
def _scope(kind, value):
    """(WHERE for `calls`, WHERE for `daily` or None when only `calls` can hold it, parameters)."""
    if kind == 'entry':
        if value.startswith('s:'):
            return 'session_id = ?', None, (value[2:],)
        return 'id = ?', None, (int(value[2:]),)
    if kind == 'call':
        return 'id = ?', None, (int(value),)
    if kind == 'day':
        return 'day = ?', 'day = ?', (value,)
    if kind == 'month':
        return 'substr(day, 1, 7) = ?', 'substr(day, 1, 7) = ?', (value,)
    if kind == 'provider':
        return 'provider = ?', 'provider = ?', (value,)
    if kind == 'all':
        return '1 = 1', '1 = 1', ()
    raise ValueError('unknown delete scope: {}'.format(kind))


def preview_delete(kind, value=None):
    """What deleting a scope would remove: {'records': rows (calls + rolled-up days), 'requests': calls, 'cost': USD}."""
    calls_where, daily_where, params = _scope(kind, value)
    with _db() as con:
        calls = con.execute('SELECT COUNT(*) AS n, COALESCE(SUM(requests), 0) AS requests, COALESCE(SUM(cost), 0) AS cost '
                            'FROM calls WHERE {}'.format(calls_where), params).fetchone()
        days = con.execute('SELECT COUNT(*) AS n, COALESCE(SUM(requests), 0) AS requests, COALESCE(SUM(cost), 0) AS cost '
                           'FROM daily WHERE {}'.format(daily_where), params).fetchone() if daily_where else None
    return {'records': calls['n'] + (days['n'] if days else 0), 'requests': calls['requests'] + (days['requests'] if days else 0),
            'cost': calls['cost'] + (days['cost'] if days else 0.0)}


def delete(kind, value=None):
    """Removes a scope from both tables; returns how many rows were removed."""
    calls_where, daily_where, params = _scope(kind, value)
    with _db() as con:
        removed = con.execute('DELETE FROM calls WHERE {}'.format(calls_where), params).rowcount
        if daily_where:
            removed += con.execute('DELETE FROM daily WHERE {}'.format(daily_where), params).rowcount
        return removed


def delete_entry(key):
    return delete('entry', key)


def delete_call(call_id):
    return delete('call', call_id)


def delete_day(day):
    return delete('day', day)


def delete_month(month):
    return delete('month', month)


def delete_provider(provider):
    return delete('provider', provider)


def delete_all():
    return delete('all')


# ---- export
EXPORT_COLUMNS = ('ts', 'provider', 'model', 'purpose', 'title', 'status', 'error', 'latency_ms', 'requests',
                  'input_tokens', 'cached_tokens', 'output_tokens', 'cost', 'savings')


def export_csv(path):
    """Writes every call (and the rolled-up days) as CSV; returns the number of rows."""
    with _db() as con:
        calls = con.execute('SELECT {} FROM calls ORDER BY ts'.format(', '.join(EXPORT_COLUMNS))).fetchall()
        daily = con.execute(
            "SELECT day AS ts, provider, model, purpose, '' AS title, 'summary' AS status, '' AS error, 0 AS latency_ms, "
            "requests, input_tokens, cached_tokens, output_tokens, cost, savings FROM daily ORDER BY day").fetchall()
    with open(path, 'w', encoding='utf-8', newline='') as handle:
        writer = csv.writer(handle)
        writer.writerow(EXPORT_COLUMNS)
        for row in list(daily) + list(calls):
            writer.writerow([round(row[column], 8) if isinstance(row[column], float) else row[column] for column in EXPORT_COLUMNS])
    return len(calls) + len(daily)


# ---- the old JSON file (ai_usage_stats.json) becomes the first rows of the database, once
def _import_legacy(con):
    path = legacy_json_path()
    if not os.path.exists(path):
        return
    try:
        with open(path, encoding='utf-8') as handle:
            data = json.load(handle)
    except (OSError, ValueError):
        return
    imported = defaultdict(lambda: [0, 0.0, 0, 0, 0])               # (provider, model) -> requests, cost, tokens
    for record in data.get('call_history', []):
        stamp = record.get('timestamp') or ''
        provider, model = record.get('provider', ''), record.get('model', '')
        if not (stamp and provider and model):
            continue
        requests = max(1, int(record.get('inference_count') or 1))
        values = (stamp, stamp[:10], provider, model, record.get('request_type') or 'other', record.get('session_id') or '',
                  requests, record.get('input_tokens', 0), record.get('cached_tokens', 0), record.get('output_tokens', 0),
                  record.get('cost', 0.0))
        con.execute('INSERT INTO calls(ts, day, provider, model, purpose, session_id, requests, input_tokens, cached_tokens, '
                    'output_tokens, cost) VALUES(?,?,?,?,?,?,?,?,?,?,?)', values)
        total = imported[(provider, model)]
        for index, value in enumerate((requests, record.get('cost', 0.0), record.get('input_tokens', 0),
                                       record.get('cached_tokens', 0), record.get('output_tokens', 0))):
            total[index] += value
    day = (data.get('created') or '')[:10] or _today().isoformat()
    for provider, provider_data in data.get('providers', {}).items():           # the totals older than the kept calls
        for model, model_data in provider_data.get('models', {}).items():
            done = imported[(provider, model)]
            rest = [max(0, model_data.get('requests', 0) - done[0]), max(0.0, model_data.get('total_cost', 0.0) - done[1]),
                    max(0, model_data.get('input_tokens', 0) - done[2]), max(0, model_data.get('cached_tokens', 0) - done[3]),
                    max(0, model_data.get('output_tokens', 0) - done[4])]
            if rest[0] or rest[1] > 1e-9:
                con.execute('INSERT INTO daily(day, provider, model, purpose, requests, input_tokens, cached_tokens, '
                            'output_tokens, cost) VALUES(?,?,?,?,?,?,?,?,?)',
                            (day, provider, model, LEGACY_PURPOSE, rest[0], rest[2], rest[3], rest[4], rest[1]))
    try:
        os.replace(path, path + '.migrated')
    except OSError:
        pass