"""FastAPI wrapper for RxDecide-1.7B / meddecide 0.6.2 (research prototype, not for clinical use).
Run from the repository folder: MEDDECIDE_DIR=. uvicorn server:app   (CATALOG and LABELS default to data/)
State may be a note string, a dict of sections, or a list of lines."""
import json, os
from pathlib import Path
from typing import Any, Dict, List, Optional, Union
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer
import meddecide as md

D = Path(os.environ['MEDDECIDE_DIR'])
CONFIG = json.loads((D / 'meddecide_config.json').read_text())
CATALOG = json.loads(Path(os.environ.get('CATALOG', D / 'data/catalog.json')).read_text())
LABELS_JSON = json.loads(Path(os.environ.get('LABELS', D / 'data/labels.json')).read_text())['labels']
REF_BY_SET = {v.get('set_id'): v['contraindications'] for v in LABELS_JSON.values()}
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
MD = D / 'model' if (D / 'model').exists() else D
TOK = AutoTokenizer.from_pretrained(MD)
MODEL = AutoModelForCausalLM.from_pretrained(MD, dtype=torch.bfloat16).to(DEVICE).eval()
READER = md.FastReader(MODEL, TOK, DEVICE)   # 0.4.0: adopted after speed_v2 (section 180)
import threading
_LOCK, _ASK = threading.Lock(), READER.ask
def _locked_ask(*args, **kwargs):            # FastReader keeps a key/value cache: one request on the GPU at a time
    with _LOCK:
        return _ASK(*args, **kwargs)
READER.ask = _locked_ask
DISCLAIMER = 'Research prototype; not validated for clinical use.'
app = FastAPI(title='RxDecide-1.7B (research prototype, not for clinical use)')
State = Union[str, Dict[str, Any], List[str]]


class DecideReq(BaseModel):
    state: State
    set_id: str                       # product-level label (see GET /products)


class FactReq(BaseModel):
    state: State
    statement: str


class ValueReq(BaseModel):
    state: State
    question: str


class ChoiceReq(BaseModel):
    state: State
    question: str
    options: List[str]


class EvalReq(BaseModel):
    state: State
    structure: Dict[str, Any]                      # {"facts": {...}, "logic": {...}}; logic optional when outcomes are given
    reference_text: Optional[str] = ''
    outcomes: Optional[List[Dict[str, Any]]] = None  # caller-defined result labels: [{"label", "when"}, ..., {"label", "otherwise": true}]
    title: Optional[str] = None


@app.get('/products')
def products():
    return md.products(CATALOG)


@app.post('/decide')
def decide(req: DecideReq):
    if req.set_id not in REF_BY_SET:
        raise HTTPException(404, 'unknown set_id; see GET /products')
    res = md.decide_product(READER, md.state_to_note(req.state), CATALOG, req.set_id, REF_BY_SET[req.set_id], CONFIG)
    return {'set_id': req.set_id, 'results': res, 'any_applies': any(r['decision'] == 'applies' for r in res),
            'needs_review': [r['criterion_id'] for r in res if r['route_to_review']], 'disclaimer': DISCLAIMER}


@app.post('/ask/fact')
def ask_fact(req: FactReq):
    return {**md.ask_fact(READER, req.state, req.statement, CONFIG), 'disclaimer': DISCLAIMER}


@app.post('/ask/value')
def ask_value(req: ValueReq):
    return {**md.ask_value(READER, req.state, req.question, CONFIG), 'disclaimer': DISCLAIMER}


@app.post('/ask/choice')
def ask_choice(req: ChoiceReq):
    return {**md.ask_choice(READER, req.state, req.question, req.options, CONFIG), 'disclaimer': DISCLAIMER}


@app.post('/evaluate')
def evaluate(req: EvalReq):
    try:
        res = md.evaluate(READER, req.state, req.structure, CONFIG, req.reference_text or '', title=req.title,
                          outcomes=req.outcomes)
    except ValueError as e:
        raise HTTPException(422, str(e))
    return {**res, 'disclaimer': DISCLAIMER}