"""Open typed-decision server with a Jev-compatible API.

POST /v1/systemone  {"state": ..., "questions": {id: {"type": "choice|score|noul", "instructions": ..., "criteria": ...}},
                     "model": "kivok" | "migom" | "auto"}
  -> {"model": ..., "answers": {id: {...TypeSafe answer shape...}}, "usage": {...}}
GET  /v1/models

Models:
  kivok - small (0.3B) and fast, runs on CPU;
  migom - larger (2B), better at reasoning; GPU recommended;
  auto  - cascade: Kivok answers, questions where its top probability is below AUTO_THRESHOLD go to Migom.

Env: KIVOK_PATH, MIGOM_PATH (either may be empty), DEVICE (cpu|cuda|xpu), DTYPE (fp16|bf16|fp32; use fp16 on Intel XPU),
     GRAPHS (0 disables graphs), AUTO_THRESHOLD (default 0.8),
     API_KEY (if set, require "Authorization: Bearer <key>"), HOST (default 127.0.0.1), PORT (default 8010).
"""
import hmac, os, threading, time
from typing import Any, Dict, Optional
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel

KIVOK_PATH = os.environ.get('KIVOK_PATH', '')
MIGOM_PATH = os.environ.get('MIGOM_PATH', '')
DEVICE = os.environ.get('DEVICE') or None
THRESHOLD = float(os.environ.get('AUTO_THRESHOLD', '0.8'))
API_KEY = os.environ.get('API_KEY', '')
DTYPE = os.environ.get('DTYPE', '')  # fp16 on Intel XPU: bf16 gives wrong answers there
GRAPHS = os.environ.get('GRAPHS', '')  # '0' disables CUDA/XPU graphs

app = FastAPI(title='typed-decision server', version='0.1')
_lock = threading.Lock()
_models: Dict[str, Any] = {}


def kivok():
    if 'kivok' not in _models:
        if not KIVOK_PATH:
            raise HTTPException(503, 'kivok is not configured (KIVOK_PATH)')
        import laya
        _models['kivok'] = laya.load(KIVOK_PATH, device=DEVICE) if DEVICE else laya.load(KIVOK_PATH)
    return _models['kivok']


def migom():
    if 'migom' not in _models:
        if not MIGOM_PATH:
            raise HTTPException(503, 'migom is not configured (MIGOM_PATH)')
        from decider.infer import Decider
        import torch
        dt = {'fp16': torch.float16, 'bf16': torch.bfloat16, 'fp32': torch.float32}.get(DTYPE)
        _models['migom'] = Decider(MIGOM_PATH, device=DEVICE, dtype=dt, use_graphs=False if GRAPHS == '0' else None)
    return _models['migom']


def top_prob(ans: dict) -> float:
    if 'probabilities' in ans and ans['probabilities']:
        return max(ans['probabilities'].values())
    if 'noul' in ans:
        p = float(ans['noul']); return max(p, 1 - p)
    return 1.0


def ask_kivok(state, questions):
    return kivok().predict(state, questions)


def ask_migom(state, questions):
    return migom().system_one(state, questions)


class Request_(BaseModel):
    state: Any
    questions: Dict[str, Dict[str, Any]]
    model: Optional[str] = 'auto'


@app.middleware('http')
async def auth(request: Request, call_next):
    if API_KEY and request.url.path.startswith('/v1') and not hmac.compare_digest(
            request.headers.get('authorization', ''), f'Bearer {API_KEY}'):
        return JSONResponse({'detail': 'missing or invalid API key'}, 401)
    return await call_next(request)


@app.get('/v1/models')
def models():
    return {'data': [{'id': m, 'configured': bool(p)} for m, p in
                     (('kivok', KIVOK_PATH), ('migom', MIGOM_PATH), ('auto', KIVOK_PATH or MIGOM_PATH))]}


@app.post('/v1/systemone')
def systemone(req: Request_):
    if not req.questions:
        raise HTTPException(422, 'questions must not be empty')
    name = (req.model or 'auto').lower()
    t0 = time.time()
    with _lock:
        try:
            if name == 'kivok' or (name == 'auto' and not MIGOM_PATH):
                out = ask_kivok(req.state, req.questions); used = {k: 'kivok' for k in req.questions}
            elif name == 'migom' or (name == 'auto' and not KIVOK_PATH):
                out = ask_migom(req.state, req.questions); used = {k: 'migom' for k in req.questions}
            elif name == 'auto':
                out = ask_kivok(req.state, req.questions)
                used = {k: 'kivok' for k in req.questions}
                hard = {k: q for k, q in req.questions.items() if top_prob(out['answers'][k]) < THRESHOLD}
                if hard:
                    second = ask_migom(req.state, hard)
                    for k in hard:
                        out['answers'][k] = second['answers'][k]; used[k] = 'migom'
            else:
                raise HTTPException(404, f'unknown model {name!r}')
        except ValueError as e:
            raise HTTPException(422, str(e))
    return {'model': name, 'answers': out['answers'], 'routed': used,
            'usage': {'latency_ms': round(1000 * (time.time() - t0), 1)}}


if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host=os.environ.get('HOST', '127.0.0.1'), port=int(os.environ.get('PORT', '8010')))