from __future__ import annotations

import os
import random

MODEL_ID = os.environ.get("PILLPAL_MODEL", "openbmb/MiniCPM-V-4_5")
USE_MOCK = os.environ.get("PILLPAL_MOCK", "0") == "1"
GPU_DURATION = int(os.environ.get("PILLPAL_GPU_DURATION", "60"))

BACKEND = os.environ.get("PILLPAL_BACKEND", "api").lower()
LLAMA_URL = os.environ.get("PILLPAL_LLAMA_URL", "http://localhost:8080/v1/chat/completions")

API_URL = os.environ.get("PILLPAL_API_URL", "https://api.modelbest.cn/v1/chat/completions")
API_KEY = os.environ.get("PILLPAL_API_KEY", "XXXX")
API_MODEL = os.environ.get("PILLPAL_API_MODEL", "MiniCPM-V-4.6-Instruct")

ON_SPACE = bool(os.environ.get("SPACE_ID"))


try:
    import spaces

    def gpu(fn):
        return spaces.GPU(duration=GPU_DURATION)(fn)
except Exception:
    def gpu(fn):
        return fn


PROMPT = (
    "You are reading a medication bottle label. Return ONLY a JSON object with "
    "exactly these keys: drug_name, dose, frequency_text, quantity, "
    "refill_or_expiry_date. "
    "- frequency_text: copy the dosing instruction verbatim (e.g. 'take 1 tablet twice daily'). "
    "- quantity: the number of pills/units in the bottle, as a number. "
    "- refill_or_expiry_date: any refill-by or expiry date, as YYYY-MM-DD if possible. "
    "Use null for any field not visible. Do NOT invent or infer dosing. "
    "Return the JSON and nothing else."
)

_model = None
_tokenizer = None


def _load_model():
    global _model, _tokenizer
    if _model is not None:
        return
    import torch
    from transformers import AutoModel, AutoTokenizer

    _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    _model = AutoModel.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
        attn_implementation="sdpa",
        torch_dtype=torch.bfloat16,
    ).eval()
    if ON_SPACE or torch.cuda.is_available():
        _model = _model.to("cuda")


@gpu
def _run_model(image) -> str:
    _load_model()

    msgs = [{"role": "user", "content": [image.convert("RGB"), PROMPT]}]
    answer = _model.chat(
        msgs=msgs,
        tokenizer=_tokenizer,
        sampling=False,
        max_new_tokens=256,
    )
    return answer if isinstance(answer, str) else str(answer)


_MOCK_SAMPLES = [
    '{"drug_name": "Metformin", "dose": "500 mg", "frequency_text": "take 1 tablet twice daily", "quantity": 60, "refill_or_expiry_date": "2026-07-01"}',
    '{"drug_name": "Lisinopril", "dose": "10 mg", "frequency_text": "take 1 tablet once daily", "quantity": 30, "refill_or_expiry_date": null}',
    '{"drug_name": "Atorvastatin", "dose": "20 mg", "frequency_text": "take 1 tablet at bedtime", "quantity": 90, "refill_or_expiry_date": "2026-06-09"}',
    '{"drug_name": "Amoxicillin", "dose": "500 mg", "frequency_text": "take 1 capsule every 8 hours", "quantity": 21, "refill_or_expiry_date": null}',
    '{"drug_name": "Ibuprofen", "dose": "200 mg", "frequency_text": "take as needed for pain", "quantity": 50, "refill_or_expiry_date": null}',
]


def _image_to_data_url(image) -> str:
    import base64
    import io

    buf = io.BytesIO()
    image.convert("RGB").save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
    return f"data:image/png;base64,{b64}"


def _run_llama_server(image) -> str:
    import json
    import urllib.request

    payload = {
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": PROMPT},
                    {"type": "image_url", "image_url": {"url": _image_to_data_url(image)}},
                ],
            }
        ],
        "temperature": 0,
        "max_tokens": 256,
    }
    req = urllib.request.Request(
        LLAMA_URL,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=120) as resp:
        body = json.loads(resp.read().decode("utf-8"))
    return body["choices"][0]["message"]["content"]


def _run_modelbest_api(image) -> str:
    import json
    import urllib.request

    payload = {
        "model": API_MODEL,
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": PROMPT},
                    {"type": "image_url", "image_url": {"url": _image_to_data_url(image)}},
                ],
            }
        ],
        "temperature": 0,
        "max_tokens": 256,
    }
    req = urllib.request.Request(
        API_URL,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=120) as resp:
        body = json.loads(resp.read().decode("utf-8"))
    return body["choices"][0]["message"]["content"]


def extract_label(image) -> str:
    if USE_MOCK or image is None:
        return random.choice(_MOCK_SAMPLES)
    if BACKEND == "api":
        return _run_modelbest_api(image)
    if BACKEND == "llama":
        return _run_llama_server(image)
    return _run_model(image)


if ON_SPACE and not USE_MOCK and BACKEND == "transformers":
    try:
        _load_model()
    except Exception as exc:
        print(f"[pillpal] deferred model load (will retry on first request): {exc}")