"""Shared helpers for 2.7: raw HTTP calls (so we can read the gateway's orchestrator trace), prompts, judge.

Gateway facts found while probing (see probe.json and README):
  * deepseek-v4-flash goes through the gateway's orchestrator, which gives the model its own `search` / `extract`
    web tools and ignores client-supplied tools. We ask for debug=true and record every web tool call so those
    answers can be flagged. The index condition for deepseek uses a plain-text READ_SKILL protocol instead of tools.
  * gpt-5 is forced to temperature=1.0 by the gateway and rejects the debug flag.
  * grok-4-fast is a passthrough to xAI grok-4.3 (non-reasoning); client tools work.
"""
import json, pathlib, sys, time, random
import httpx

ROOT = pathlib.Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from common import token, BASE  # noqa: E402

HERE = pathlib.Path(__file__).resolve().parent

# USD per 1M tokens (input, output). deepseek / grok are read from the gateway trace when available;
# these are only used as a fallback and for gpt-5 (OpenAI list price; cached-input discount not applied).
PRICE = {'gpt-5': (1.25, 10.0), 'grok-4-fast': (1.25, 2.5), 'deepseek-v4.1-flash': (0.21, 0.42), 'deepseek-v4-flash': (0.21, 0.42)}

PREAMBLE = ('你在一个离线环境里工作，不能联网，也不要调用任何搜索或网页抓取工具。'
            '请只根据下面提供的资料和你自己已有的知识回答用户的问题。回答要具体，直接给出可以照做的做法。')

_http = None


def http():
    global _http
    if _http is None:
        _http = httpx.Client(timeout=900, limits=httpx.Limits(max_connections=32))
    return _http


def call(model, messages, tools=None, tool_choice=None, max_tokens=None, retries=4):
    """One chat completion. Returns dict(message, usage, sec, web_tools, cost, error)."""
    api_model = 'deepseek-v4-flash' if model == 'deepseek-v4.1-flash' else model
    body = {'model': api_model, 'messages': messages, 'temperature': 0}
    if tools:
        body['tools'] = tools
    if tool_choice:
        body['tool_choice'] = tool_choice
    if max_tokens:
        body['max_tokens'] = max_tokens
    url = f'{BASE}/chat/completions'
    if model.startswith('deepseek'):
        body['debug'] = True
        url += '?debug=true'
    last = None
    for attempt in range(retries):
        t0 = time.time()
        try:
            r = http().post(url, headers={'Authorization': f'Bearer {token()}'}, json=body)
        except Exception as e:  # network / timeout
            last = f'{type(e).__name__}: {e}'
            time.sleep(5 * (attempt + 1))
            continue
        sec = time.time() - t0
        if r.status_code >= 500 or r.status_code == 429:
            last = f'HTTP {r.status_code}: {r.text[:300]}'
            time.sleep(10 * (attempt + 1))
            continue
        if r.status_code != 200:
            return {'error': f'HTTP {r.status_code}: {r.text[:600]}', 'sec': sec}
        j = r.json()
        msg = j['choices'][0]['message']
        usage = j.get('usage') or {}
        tr = j.get('orchestrator_trace') or {}
        web = []
        for rd in tr.get('rounds', []):
            for tc in rd.get('tool_calls') or []:
                web.append({'name': tc.get('name'), 'arguments': tc.get('arguments')})
        cost = (tr.get('totals') or {}).get('total_cost_usd')
        if cost is None:
            pin, pout = PRICE.get(model, (0, 0))
            cost = usage.get('prompt_tokens', 0) * pin / 1e6 + usage.get('completion_tokens', 0) * pout / 1e6
        return {'message': msg, 'usage': usage, 'sec': sec, 'web_tools': web, 'cost': cost,
                'finish_reason': j['choices'][0].get('finish_reason')}
    return {'error': last or 'unknown', 'sec': 0}


JUDGE_SYS = ('You grade answers for an experiment. A user asked a question whose correct answer depends on one specific '
             'rule or fact from a reference document. You get the question, the verbatim excerpt that states the rule, '
             'a grading rubric, and the answer. Mark correct=true only if the answer clearly states or applies the rule '
             'as the excerpt gives it (the specific value / flag / forbidden step / gotcha), and does not contradict it. '
             'Listing the right option only as one of several equally-weighted guesses is NOT correct. Paraphrase in any '
             'language is fine. Extra harmless detail is fine. Reply with JSON only: {"correct": true|false, "reason": "<one sentence>"}')


def judge(task, answer):
    if not answer or not answer.strip():
        return {'correct': False, 'reason': 'empty answer', 'judge_error': None}
    user = (f'QUESTION:\n{task["question"]}\n\nEXCERPT (ground truth):\n{task["excerpt"]}\n\n'
            f'RUBRIC:\n{json.dumps(task["rubric"], ensure_ascii=False)}\n\nANSWER TO GRADE:\n{answer}')
    for _ in range(3):
        r = call('gpt-5', [{'role': 'system', 'content': JUDGE_SYS}, {'role': 'user', 'content': user}], max_tokens=4000)
        if 'error' in r:
            continue
        txt = (r['message'].get('content') or '').strip()
        try:
            s, e = txt.find('{'), txt.rfind('}')
            d = json.loads(txt[s:e + 1])
            return {'correct': bool(d.get('correct')), 'reason': d.get('reason', ''), 'judge_cost': r['cost']}
        except Exception:
            continue
    return {'correct': None, 'reason': 'judge failed', 'judge_error': True}


def load_pool():
    pool = json.loads((HERE / 'pool.json').read_text())
    texts = json.loads((HERE / 'pool_text.json').read_text())
    return pool, texts


def one_line(desc, cap=240):
    d = ' '.join(desc.split())
    return d if len(d) <= cap else d[:cap - 1] + '…'


def index_text(pool, overrides=None):
    overrides = overrides or {}
    lines = [f'- {s["name"]}: {one_line(overrides.get(s["name"], s["description"]))}' for s in pool['skills']]
    return '\n'.join(lines)


def skills_block(names, texts):
    return '\n\n'.join(f'<skill name="{n}">\n{texts[n]}\n</skill>' for n in names)