"""Shared helpers for the module 2 experiments.

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/<script>.py
"""
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')


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='zh'):
    """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'},
                       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)."""
    api_model = 'deepseek-v4-flash' if model == 'deepseek-v4.1-flash' else model
    t0 = time.time()
    r = client().chat.completions.create(model=api_model, messages=messages, **kw)
    u = r.usage.model_dump() if r.usage else {}
    return r.choices[0].message.content, u, time.time() - t0


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