"""Shared CalorieCLIP loading and inference."""

from __future__ import annotations

import base64
import io
import logging
import os
from typing import Any

import open_clip
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from PIL import Image

logger = logging.getLogger("calorie_clip")

MODEL_REPO = os.getenv("CALORIE_CLIP_MODEL_REPO", "jc-builds/CalorieCLIP")
WEIGHTS_FILE = os.getenv("CALORIE_CLIP_WEIGHTS_FILE", "calorie_clip.pt")
MAX_IMAGE_BYTES = int(os.getenv("MAX_IMAGE_BYTES", str(8 * 1024 * 1024)))

_state: dict[str, Any] = {}


class RegressionHead(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(512, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


def load_models() -> None:
    if _state.get("ready"):
        return
    logger.info("Loading CalorieCLIP from %s/%s", MODEL_REPO, WEIGHTS_FILE)
    clip_model, _, preprocess = open_clip.create_model_and_transforms(
        "ViT-B-32",
        pretrained="openai",
    )
    weights_path = hf_hub_download(repo_id=MODEL_REPO, filename=WEIGHTS_FILE)
    checkpoint = torch.load(weights_path, map_location="cpu", weights_only=False)
    clip_model.load_state_dict(checkpoint["clip_state"], strict=False)
    head = RegressionHead()
    head.load_state_dict(checkpoint["regressor_state"])
    clip_model.eval()
    head.eval()

    _state["clip"] = clip_model
    _state["head"] = head
    _state["preprocess"] = preprocess
    _state["ready"] = True
    logger.info("CalorieCLIP ready")


def decode_image_bytes(raw: bytes) -> Image.Image:
    if len(raw) > MAX_IMAGE_BYTES:
        raise ValueError("Image too large")
    return Image.open(io.BytesIO(raw)).convert("RGB")


def decode_base64_image(b64: str) -> Image.Image:
    cleaned = b64.strip()
    if cleaned.startswith("data:"):
        cleaned = cleaned.split(",", 1)[1]
    try:
        raw = base64.b64decode(cleaned, validate=True)
    except Exception as exc:
        raise ValueError("Invalid base64 image") from exc
    return decode_image_bytes(raw)


def predict_calories(image: Image.Image) -> int:
    load_models()
    clip_model = _state["clip"]
    head = _state["head"]
    preprocess = _state["preprocess"]
    tensor = preprocess(image).unsqueeze(0)
    with torch.no_grad():
        features = clip_model.encode_image(tensor)
        calories = float(head(features).item())
    return round(max(0.0, calories))