"""Shared helpers for the module 2 experiments (English edition).

All calls go to AI Builder Space, the same API readers use. The key is read from AI_BUILDER_TOKEN; run scripts via
  op run --env-file experiments/aib.env -- .venv/bin/python experiments_en/<script>.py
Differences from experiments/common.py: transcribe() defaults to language='en', and there is a small cost logger
(log_cost) that appends one line per script run to experiments_en/COSTS.md.
"""
import datetime
import json
import os
import pathlib
import time

import httpx
from openai import OpenAI

ROOT = pathlib.Path(__file__).resolve().parent
BASE = os.getenv('AI_BUILDER_SPACE_BASE_URL', 'https://space.ai-builders.com/backend/v1')

# USD per 1M tokens (input, output). Approximate list prices, used only for the cost log.
PRICE = {'deepseek-v4.1-flash': (0.21, 0.42), 'deepseek-v4-flash': (0.21, 0.42),
         'gemini-3.8-flash': (0.50, 3.00), 'grok-4-fast': (1.25, 2.5)}
# Speech-to-text, USD per audio hour (grok from the gateway billing field; compat ~ whisper-1 list price).
ASR_PRICE_PER_HOUR = {'grok': 0.10, 'compat': 0.36}

_usage = {'calls': 0, 'prompt_tokens': 0, 'completion_tokens': 0, 'usd': 0.0}


def token():
    t = os.getenv('AI_BUILDER_TOKEN')
    if not t:
        raise SystemExit('AI_BUILDER_TOKEN is not set; run through op run --env-file experiments/aib.env')
    return t


def client():
    return OpenAI(base_url=BASE, api_key=token(), timeout=300)


def transcribe(path, engine='grok', terms=None, language='en'):
    """Transcribe one audio file. engine='grok' uses /audio/grok-transcription (supports a terms list);
    engine='compat' uses the OpenAI-compatible /audio/transcriptions route."""
    path = pathlib.Path(path)
    headers = {'Authorization': f'Bearer {token()}'}
    with httpx.Client(timeout=300) as h:
        if engine == 'grok':
            data = {'simple': 'true', 'language': language}
            if terms:
                data['terms'] = ','.join(terms)
            r = h.post(f'{BASE}/audio/grok-transcription', headers=headers, data=data,
                       files={'audio_file': (path.name, path.read_bytes())})
        else:
            r = h.post(f'{BASE}/audio/transcriptions', headers=headers, data={'model': 'whisper-1', 'language': language},
                       files={'file': (path.name, path.read_bytes())})
    r.raise_for_status()
    return r.json()['text']


def chat(messages, model, **kw):
    """One chat call; returns (text, usage dict, seconds). Usage is accumulated for log_cost()."""
    api_model = 'deepseek-v4-flash' if model == 'deepseek-v4.1-flash' else model
    if model.startswith('gemini'):
        # Gateway quirk (2026-10-14): some plain-string Gemini prompts (e.g. any containing "Found a gotcha today.")
        # deterministically fail with 400 "Requests ending with a model turn are not supported". Sending the same text
        # as a list of content parts avoids it. Same text, same model; only the wire format differs.
        messages = [dict(m, content=[{'type': 'text', 'text': m['content']}]) if isinstance(m.get('content'), str) else m
                    for m in messages]
    t0 = time.time()
    for attempt in range(4):
        try:
            r = client().chat.completions.create(model=api_model, messages=messages, **kw)
            break
        except Exception as e:  # transient gateway errors (e.g. Gemini 400 "ending with a model turn") are retried
            if attempt == 3:
                raise
            time.sleep(2 + 3 * attempt)
    u = r.usage.model_dump() if r.usage else {}
    pin, pout = PRICE.get(model, (1.0, 5.0))
    pt, ct = u.get('prompt_tokens') or 0, u.get('completion_tokens') or 0
    _usage['calls'] += 1; _usage['prompt_tokens'] += pt; _usage['completion_tokens'] += ct
    _usage['usd'] += pt * pin / 1e6 + ct * pout / 1e6
    return r.choices[0].message.content or '', u, time.time() - t0


def log_cost(script, extra_usd=0.0, note=''):
    """Append one line to COSTS.md for this script run."""
    f = ROOT / 'COSTS.md'
    if not f.exists():
        f.write_text('# Cost log (English experiments)\n\nOne line per script run. Token prices are approximate list prices '
                     '(see PRICE in common.py); ASR is priced per audio hour.\n\n'
                     '| date | script | calls | prompt tok | completion tok | est. $ | note |\n|---|---|---|---|---|---|---|\n',
                     encoding='utf-8')
    usd = _usage['usd'] + extra_usd
    with open(f, 'a', encoding='utf-8') as fh:
        fh.write(f"| {datetime.date.today()} | {script} | {_usage['calls']} | {_usage['prompt_tokens']} | "
                 f"{_usage['completion_tokens']} | {usd:.4f} | {note} |\n")
    print(f'cost logged: {usd:.4f} USD, {_usage["calls"]} chat calls')


def dump(obj, path):
    pathlib.Path(path).write_text(json.dumps(obj, ensure_ascii=False, indent=1), encoding='utf-8')