"""MIT. Forced-tool extraction; arithmetic is checked locally, not by a model."""
import json
import os
from pathlib import Path

FIELDS = ['invoice_id', 'document_type', 'vendor', 'email', 'invoice_date', 'due_date', 'currency', 'subtotal', 'tax', 'total', 'bank_code', 'purchase_order']
LINE_FIELDS = ['description', 'quantity', 'unit_price', 'amount']
TOOL = {'type': 'function', 'function': {
    'name': 'record', 'description': 'Extract the printed invoice, without correcting totals.',
    'parameters': {'type': 'object', 'properties': {
        **{name: {'type': ['string', 'null']} for name in FIELDS},
        'line_items': {'type': 'array', 'items': {'type': 'object', 'properties': {name: {'type': 'string'} for name in LINE_FIELDS}, 'required': LINE_FIELDS, 'additionalProperties': False}},
        'note': {'type': 'string', 'description': 'Briefly describe any absent field or unusual sign. Not scored.'},
    }, 'required': FIELDS + ['line_items', 'note'], 'additionalProperties': False},
}}


def request(item, config):
    return {'messages': [{'role': 'system', 'content': config['prompt']}, {'role': 'user', 'content': item['text']}],
            'tools': [TOOL], 'tool_choice': config['tool_choice'], 'max_tokens': config['max_tokens'],
            'reasoning_effort': config['reasoning_effort'], 'parallel_tool_calls': False}


def extract(item, config, client):
    response = client.chat.completions.create(model=os.environ.get('MODEL', 'pareto'), **request(item, config))
    calls = response.choices[0].message.tool_calls or []
    if len(calls) != 1 or calls[0].function.name != 'record':
        raise ValueError('Expected exactly one record call; extra calls are rejected')
    return json.loads(calls[0].function.arguments)


if __name__ == '__main__':
    from openai import OpenAI
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--item', default='dev-00')
    args = parser.parse_args()
    here = Path(__file__).resolve().parent
    config = json.loads((here / 'eval.json').read_text())
    item = next(json.loads(line) for line in (here / 'fixture.jsonl').read_text().splitlines() if json.loads(line)['item_id'] == args.item)
    model = os.environ.get('MODEL', 'pareto')
    client = OpenAI(base_url=os.environ.get('BASE_URL', 'https://api.unbiased.ai/v1'), api_key=os.environ['UNBIASED_API_KEY' if model == 'pareto' else 'OPENAI_API_KEY'], max_retries=0)
    print(json.dumps(extract(item, config, client), indent=2))