"""CPU inference. TinyStories keeps backend/local.lua and grammar.lua intact."""
from functools import lru_cache
import json
import os
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
VENDOR = ROOT / 'space' / 'vendor' / 'lua_llama'
TINY_REPO = 'karpathy/tinyllamas'
QWEN_REPO = 'Qwen/Qwen2.5-1.5B-Instruct-GGUF'


def model_file(repo, filename):
    from huggingface_hub import hf_hub_download
    if repo == TINY_REPO and filename == 'tokenizer.bin':
        return str(VENDOR / 'tokenizer.bin')
    if repo == TINY_REPO:
        local = Path(os.environ.get('LUA_LLAMA_PATH', ROOT.parent / 'lua_llama')) / filename
        if local.is_file():
            return str(local)
    from space.revisions import REVISIONS
    return hf_hub_download(repo, filename, revision=REVISIONS[repo])


@lru_cache(maxsize=2)
def tiny_weights(device='cpu'):
    import torch
    from tools.llama2c import load
    torch.set_num_threads(min(2, os.cpu_count() or 1))
    cfg, weights = load(model_file(TINY_REPO, 'stories15M.bin'))
    return cfg, {name: tensor.to(device) for name, tensor in weights.items()}


def tiny_backend(lua, registry, config):
    import torch
    from tools.llama2c import forward
    cfg, weights = tiny_weights(config.get('device', 'cpu'))
    kcache = torch.zeros(cfg.n_layers, cfg.seq_len, cfg.kv_dim, device=weights['tok_emb'].device)
    vcache = torch.zeros_like(kcache)

    def step(_model, token, pos, *_unused):
        return lua.table_from(forward(cfg, weights, token, pos, kcache, vcache).tolist())

    lua.globals().lab_forward = step
    lua.globals().lab_cfg = lua.table_from(vars(cfg))
    # Inject only numerical kernels; the original local backend still owns
    # tokenization, cache positions, scaffolding, sampling and grammar masks.
    lua.execute('''
        package.loaded.model = {new=function()
            return {config=lab_cfg, create_run_state=function() return {} end}
        end}
        package.loaded.generate = {forward=function(...) return lab_forward(...) end}
    ''')
    options = lua.table_from({
        'llama_path': VENDOR.as_posix(),
        'tokenizer': model_file(TINY_REPO, 'tokenizer.bin'),
        'tools': registry, 'constrain': config['grammar'], 'temperature': 0.0,
        'max_thought': 24, 'max_call': 96,
    })
    backend = lua.eval("require('backend.local').new")(options)
    return backend, backend.tok


@lru_cache(maxsize=1)
def qwen_model():
    from llama_cpp import Llama
    return Llama(
        model_path=model_file(QWEN_REPO, 'qwen2.5-1.5b-instruct-q4_k_m.gguf'),
        n_ctx=4096, n_threads=min(2, os.cpu_count() or 1), n_batch=256,
        verbose=False, seed=0,
    )


def qwen_backend(lua, registry, config):
    model = qwen_model()
    encode_json = lua.eval("require('trace').encode")
    catalogue = json.loads(encode_json(lua.eval('function(reg) local out={} for _, t in ipairs(reg:list()) do out[#out+1]={name=t.name,args=t.args} end return out end')(registry)))
    schemas = []
    for tool in catalogue:
        props = {arg['name']: {'type': 'string'} for arg in tool['args']}
        required = [arg['name'] for arg in tool['args'] if arg.get('required')]
        schemas.append({'type': 'object', 'properties': {
            'tool': {'const': tool['name']},
            'args': {'type': 'object', 'properties': props, 'required': required, 'additionalProperties': False},
        }, 'required': ['tool', 'args'], 'additionalProperties': False})
    schema = {'oneOf': schemas}

    def complete(_self, _prompt, opts):
        pieces = json.loads(encode_json(lua.table_from({'system': opts.system, 'steps': opts.steps, 'note': opts.note})))
        messages = [{'role': 'system', 'content':
                     'Use the tools to solve the task. Return exactly one JSON object with tool and args. '
                     'All argument values must be strings. Use answer only after observing the required tool results.'},
                    {'role': 'user', 'content': pieces['system']}]
        for record in pieces.get('steps') or []:
            messages.append({'role': 'assistant', 'content': json.dumps(record['call']) if record.get('call') else record.get('raw', '')})
            if record.get('observation'):
                messages.append({'role': 'user', 'content': 'RESULT: ' + record['observation']})
        if pieces.get('note'):
            messages.append({'role': 'user', 'content': pieces['note']})
        kwargs = {'response_format': {'type': 'json_object', 'schema': schema}} if config['grammar'] else {}
        result = model.create_chat_completion(messages=messages, temperature=0, max_tokens=192, **kwargs)
        raw = result['choices'][0]['message']['content'] or ''
        # Both arms receive identical scaffolding. Free outputs are never repaired.
        return 'CALL: ' + raw.replace('\r', ''), lua.table_from({
            'tokens': result['usage']['completion_tokens'], 'constrained': config['grammar'],
            'finish_reason': result['choices'][0]['finish_reason'],
        })

    def encode(_self, text):
        return lua.table_from(model.tokenize(text.encode('utf-8'), add_bos=False))

    return lua.table_from({'complete': complete}), lua.table_from({'encode': encode})