"""CLIP-based semantic search for Muleline. Natural language photo search."""
from __future__ import annotations

import logging
import struct
import time
from pathlib import Path

from muleline.db import get_db

log = logging.getLogger(__name__)

MAX_IMAGE_PIXELS = 25_000_000
EMBEDDING_DIM = 512  # ViT-B-32 output dimension
MODEL_NAME = "ViT-B-32"

_clip_model = None
_clip_preprocess = None
_clip_tokenizer = None


def _load_model():
    global _clip_model, _clip_preprocess, _clip_tokenizer
    if _clip_model is not None:
        return

    try:
        import open_clip
        import torch
    except ImportError:
        raise ImportError("open-clip-torch and torch are required for CLIP search. Install with: pip install open-clip-torch torch")

    log.info("Loading CLIP model %s...", MODEL_NAME)
    device = "cpu"
    model, _, preprocess = open_clip.create_model_and_transforms(MODEL_NAME, pretrained="laion2b_s34b_b79k", device=device)
    tokenizer = open_clip.get_tokenizer(MODEL_NAME)
    model.eval()

    _clip_model = model
    _clip_preprocess = preprocess
    _clip_tokenizer = tokenizer
    log.info("CLIP model loaded")


def _embedding_to_blob(embedding) -> bytes:
    import numpy as np
    arr = np.array(embedding, dtype=np.float32).flatten()
    return arr.tobytes()


def _blob_to_embedding(blob: bytes):
    import numpy as np
    return np.frombuffer(blob, dtype=np.float32)


def encode_image(image_path: Path) -> bytes | None:
    try:
        import torch
        from PIL import Image
    except ImportError:
        log.error("torch and Pillow required for CLIP")
        return None

    _load_model()

    try:
        with Image.open(image_path) as img:
            w, h = img.size
            if w * h > MAX_IMAGE_PIXELS:
                ratio = (MAX_IMAGE_PIXELS / (w * h)) ** 0.5
                img = img.resize((int(w * ratio), int(h * ratio)))
            img = img.convert("RGB")
            image_tensor = _clip_preprocess(img).unsqueeze(0)
    except Exception as e:
        log.warning("Failed to load image %s: %s", image_path, e)
        return None

    with torch.no_grad():
        features = _clip_model.encode_image(image_tensor)
        features /= features.norm(dim=-1, keepdim=True)

    return _embedding_to_blob(features[0].cpu().numpy())


def encode_text(query: str) -> bytes | None:
    try:
        import torch
    except ImportError:
        return None

    _load_model()

    tokens = _clip_tokenizer([query])
    with torch.no_grad():
        features = _clip_model.encode_text(tokens)
        features /= features.norm(dim=-1, keepdim=True)

    return _embedding_to_blob(features[0].cpu().numpy())


def index_file(item_id: int, file_path: Path) -> bool:
    embedding_blob = encode_image(file_path)
    if embedding_blob is None:
        return False

    conn = get_db()
    try:
        conn.execute(
            "INSERT OR REPLACE INTO clip_embeddings (item_id, embedding, model_name, created_at) VALUES (?, ?, ?, ?)",
            (item_id, embedding_blob, MODEL_NAME, time.time()),
        )
        conn.commit()
        return True
    finally:
        conn.close()


def search(owner_id: str, query: str, limit: int = 50) -> list[dict]:
    import numpy as np

    query_embedding = encode_text(query)
    if query_embedding is None:
        return []

    query_vec = _blob_to_embedding(query_embedding)

    conn = get_db()
    try:
        rows = conn.execute(
            """SELECT c.item_id, c.embedding,
                      f.id, f.filename, f.original_name, f.mime_type, f.file_size,
                      f.file_hash, f.storage_key, f.date_taken, f.gps_lat, f.gps_lon,
                      f.width, f.height, f.created_at
               FROM clip_embeddings c
               JOIN files f ON f.id = c.item_id
               WHERE f.owner_id = ?""",
            (owner_id,),
        ).fetchall()

        if not rows:
            return []

        results = []
        for row in rows:
            item_vec = _blob_to_embedding(row["embedding"])
            similarity = float(np.dot(query_vec, item_vec))
            results.append({
                "id": row["id"],
                "filename": row["filename"],
                "original_name": row["original_name"],
                "mime_type": row["mime_type"],
                "file_size": row["file_size"],
                "file_hash": row["file_hash"],
                "storage_key": row["storage_key"],
                "date_taken": row["date_taken"],
                "gps_lat": row["gps_lat"],
                "gps_lon": row["gps_lon"],
                "width": row["width"],
                "height": row["height"],
                "created_at": row["created_at"],
                "score": round(similarity, 4),
            })

        results.sort(key=lambda x: x["score"], reverse=True)
        return results[:limit]
    finally:
        conn.close()


def index_all(owner_id: str, upload_dir: Path) -> dict:
    conn = get_db()
    try:
        files = conn.execute(
            """SELECT f.id, f.filename, f.mime_type FROM files f
               LEFT JOIN clip_embeddings c ON c.item_id = f.id
               WHERE f.owner_id = ? AND f.mime_type LIKE 'image/%' AND c.id IS NULL""",
            (owner_id,),
        ).fetchall()
    finally:
        conn.close()

    indexed = 0
    failed = 0
    for f in files:
        path = upload_dir / f["filename"]
        if path.exists():
            if index_file(f["id"], path):
                indexed += 1
            else:
                failed += 1
        else:
            failed += 1

    return {"indexed": indexed, "failed": failed, "total": len(files)}


def stats(owner_id: str) -> dict:
    conn = get_db()
    try:
        total_images = conn.execute(
            "SELECT COUNT(*) FROM files WHERE owner_id = ? AND mime_type LIKE 'image/%'",
            (owner_id,),
        ).fetchone()[0]
        indexed = conn.execute(
            """SELECT COUNT(*) FROM clip_embeddings c
               JOIN files f ON f.id = c.item_id
               WHERE f.owner_id = ?""",
            (owner_id,),
        ).fetchone()[0]
        return {"total_images": total_images, "indexed": indexed, "pending": total_images - indexed}
    finally:
        conn.close()
