import json
import time
import urllib.error
import urllib.request
from pathlib import Path

URL = "https://api.typesafe.ai/v1/systemone"
MODEL = "jev-latest"
KEY_FILE = Path.home() / ".config/typesafe/key"
PRICE_PER_INPUT_TOKEN = 0.042 / 1_000_000
REQUEST_BUDGET, STATE_BUDGET = 64_000, 32_000
CALLS = []


def has_key():
    return KEY_FILE.is_file()


def estimate_tokens(value):
    return len(json.dumps(value)) // 3


def fits(state, questions):
    longest = max(estimate_tokens(q) for q in questions.values())
    total = estimate_tokens(state) + sum(estimate_tokens(q) for q in questions.values())
    return estimate_tokens(state) + longest < STATE_BUDGET and total < REQUEST_BUDGET


def ask(state, questions, stage="", retries=6):
    body = json.dumps({"state": state, "model": MODEL, "questions": questions}).encode()
    headers = {
        "Authorization": f"Bearer {KEY_FILE.read_text().strip()}",
        "Content-Type": "application/json",
    }
    started = time.perf_counter()
    for attempt in range(retries):
        try:
            with urllib.request.urlopen(urllib.request.Request(URL, body, headers), timeout=60) as resp:
                data = json.load(resp)
            break
        except urllib.error.HTTPError as err:
            if err.code not in (429, 529) or attempt == retries - 1:
                raise RuntimeError(f"jev {err.code}: {err.read().decode()[:300]}") from None
            time.sleep(float(err.headers.get("retry-after") or 2 ** attempt))
    seconds = time.perf_counter() - started
    tokens = data["usage"]["input_tokens"]
    CALLS.append({"stage": stage, "model": data["model"], "input_tokens": tokens,
                  "questions": len(questions), "seconds": seconds,
                  "dollars": tokens * PRICE_PER_INPUT_TOKEN})
    return data["answers"]