# MIT. The same forced-tool request is used for both scored arms.
import json
import os

CATEGORIES = ['billing', 'access', 'bug', 'feature', 'unknown']
PRIORITIES = ['P1', 'P2', 'P3']
TEAMS = ['finance', 'identity', 'engineering', 'product', 'support']


def request(item, config):
    return {
        'messages': [{'role': 'system', 'content': config['prompt']}, {'role': 'user', 'content': item['text']}],
        'tools': [{'type': 'function', 'function': {
            'name': 'triage_ticket', 'description': 'Classify a ticket and assign priority and team.',
            'parameters': {'type': 'object', 'additionalProperties': False, 'properties': {
                'category': {'type': 'string', 'enum': CATEGORIES},
                'priority': {'type': 'string', 'enum': PRIORITIES},
                'team': {'type': 'string', 'enum': TEAMS},
                'confidence': {'type': 'string', 'enum': ['low', 'medium', 'high']},
                'reason': {'type': 'string', 'description': 'A short explanation grounded in the ticket.'},
            }, 'required': ['category', 'priority', 'team', 'confidence', 'reason']},
        }}],
        'tool_choice': config['tool_choice'], 'parallel_tool_calls': False,
        'max_tokens': config['max_tokens'], 'reasoning_effort': config['reasoning_effort'],
    }


def triage(text, config):
    """Optional SDK adapter. Scored run_eval.py sends request() through stdlib HTTP."""
    from openai import OpenAI
    model = os.environ.get('MODEL', 'pareto')
    base_url = os.environ.get('BASE_URL', 'https://api.unbiased.ai/v1')
    key_name = 'UNBIASED_API_KEY' if base_url == 'https://api.unbiased.ai/v1' else 'OPENAI_API_KEY'
    client = OpenAI(api_key=os.environ[key_name], base_url=base_url, max_retries=0)
    response = client.chat.completions.create(model=model, **request({'text': text}, config))
    calls = response.choices[0].message.tool_calls or []
    # The server may return extra calls even with parallel_tool_calls=False.
    if len(calls) != 1 or calls[0].function.name != 'triage_ticket':
        raise ValueError('Expected exactly one triage_ticket call')
    if response.choices[0].finish_reason == 'length':
        raise ValueError('Truncated output needs human review')
    return json.loads(calls[0].function.arguments)