"""Output schemas are application-owned, never inferred from user prompt text."""
from functools import lru_cache


def obj(properties, required=None):
    return {'type':'object', 'properties':properties,
            'required':list(properties) if required is None else required,
            'additionalProperties':False}


def array(items, minimum=0):
    return {'type':'array','items':items,'minItems':minimum,'maxItems':6}


INDEX = {'type':'integer','enum':[0,1,2,3,4,5]}
BOOL = {'type':'boolean'}
SCHEMAS = {
    'claims':obj({
        'status':{'type':'string','enum':['answer','insufficient_evidence','needs_clarification']},
        'claims':array(obj({'text':{'type':'string','minLength':1,'maxLength':450},
                            'source_id':{'type':'integer'},'excerpt_id':{'type':'integer'}}))}),
    'support':obj({'verdicts':array(obj({'claim_index':INDEX, 'supported':BOOL,
        'reason_code':{'type':'string','enum':['supported','extra_facts','reversed_meaning',
        'wrong_region','not_answering_question','wrong_language','uncertain']}},
        required=['claim_index','supported']), minimum=1)}),
    'translation':obj({'steps':array(obj({'index':INDEX,
        'text':{'type':'string','minLength':1,'maxLength':900}}), minimum=1)}),
    'translation_review':obj({'verdicts':array(obj({'index':INDEX,
        'correct_language':BOOL,'same_meaning':BOOL}), minimum=1)}),
}


class StructuredPrompt(str):
    def __new__(cls, prompt, kind):
        if kind not in SCHEMAS:
            raise ValueError('Unknown output schema')
        instance = super().__new__(cls, prompt)
        instance.kind = kind
        return instance

    def __getnewargs__(self):
        # ZeroGPU transfers function arguments into a worker process.
        return str(self), self.kind


@lru_cache(maxsize=2)
def _tokenizer_data(tokenizer):
    from lmformatenforcer.integrations.transformers import build_token_enforcer_tokenizer_data
    return build_token_enforcer_tokenizer_data(tokenizer)


def prefix_constraint(tokenizer, prompt):
    """Fail explicitly if formatting support is missing; do not silently run unconstrained."""
    from lmformatenforcer import JsonSchemaParser
    from lmformatenforcer.integrations.transformers import build_transformers_prefix_allowed_tokens_fn
    if not isinstance(prompt, StructuredPrompt):
        raise ValueError('Generation requires an application-owned output schema.')
    return build_transformers_prefix_allowed_tokens_fn(_tokenizer_data(tokenizer), JsonSchemaParser(SCHEMAS[prompt.kind]))