"""CLIP-based indoor room classifier. Runs ViT-B/32 on CPU, ~100ms per image."""

import io
import logging
import os

import open_clip
import torch
from PIL import Image
from flask import Flask, jsonify, request

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("clip-server")
app = Flask(__name__)
app.config["MAX_CONTENT_LENGTH"] = 50 * 1024 * 1024  # 50 MB max upload

# Load once at startup
MODEL_NAME = os.getenv("CLIP_MODEL", "ViT-B-32")
PRETRAINED = os.getenv("CLIP_PRETRAINED", "laion2b_s34b_b79k")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

log.info(f"Loading CLIP model: {MODEL_NAME} pretrained={PRETRAINED} device={DEVICE}")
model, _, preprocess = open_clip.create_model_and_transforms(MODEL_NAME, pretrained=PRETRAINED)
model = model.to(DEVICE).eval()
tokenizer = open_clip.get_tokenizer(MODEL_NAME)

# Text prompts for indoor/outdoor classification
PROMPTS = [
    "a photo of an indoor room interior with floor visible",
    "a photo of an outdoor scene, landscape, street, or non-room image",
]


@app.route("/health", methods=["GET"])
def health():
    return jsonify({"status": "ok", "model": MODEL_NAME, "device": DEVICE})


@app.route("/check", methods=["POST"])
def check():
    """Classify whether an image is an indoor room photo suitable for floor replacement.
    Expects raw image bytes as the request body. Returns {score: 0-10, passed: bool}."""
    if not request.data:
        return jsonify({"error": "no image data"}), 400

    try:
        img = Image.open(io.BytesIO(request.data)).convert("RGB")
    except Exception as e:
        return jsonify({"error": f"bad image: {e}"}), 400

    # Multiple prompt pairs for robust ensemble (indoor vs outdoor)
    prompt_pairs = [
        ("a photo of an indoor room interior with floor", "a photo of an outdoor scene"),
        ("interior room photograph, real estate photo", "exterior photograph, outdoor photo"),
        ("an indoor space with visible flooring", "an outdoor space, landscape, street view"),
    ]

    image_input = preprocess(img).unsqueeze(0).to(DEVICE)
    wins = 0
    gap_sum = 0.0
    for indoor_text, outdoor_text in prompt_pairs:
        text_input = tokenizer([indoor_text, outdoor_text]).to(DEVICE)
        with torch.no_grad():
            image_features = model.encode_image(image_input)
            text_features = model.encode_text(text_input)
            image_features /= image_features.norm(dim=-1, keepdim=True)
            text_features /= text_features.norm(dim=-1, keepdim=True)
            raw = (image_features @ text_features.T)[0]  # cosine similarities
        indoor_sim = float(raw[0].item())
        outdoor_sim = float(raw[1].item())
        if indoor_sim > outdoor_sim:
            wins += 1
        gap_sum += indoor_sim - outdoor_sim  # positive = indoor wins

    indoor_wins = wins >= 2  # majority vote (2 out of 3)
    avg_gap = gap_sum / len(prompt_pairs)
    # Map avg_gap to 1-10: gap ~0.05 = score 6, gap ~0.15+ = score 10
    score_1_10 = max(1, min(10, int(5 + avg_gap * 40)))
    passed = indoor_wins

    log.info(
        f"wins={wins}/3 avg_gap={avg_gap:.3f} score_1_10={score_1_10} passed={passed}"
    )
    return jsonify({"score": score_1_10, "passed": passed})


if __name__ == "__main__":
    port = int(os.getenv("CLIP_PORT", "5100"))
    log.info(f"CLIP server listening on :{port}")
    app.run(host="0.0.0.0", port=port, debug=False)