import os
import json
import torch
import numpy as np
import faiss
import pickle
import logging
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
from config import (
    INDEX_FILE,
    MAPPING_FILE,
    CACHE_FILE,
    DEFAULT_GALLERY_DIR,
    MODEL_NAME,
    EMBEDDING_DIM,
    JEWELRY_VALIDATION_THRESHOLD,
    UPLOAD_FOLDER
)

logger = logging.getLogger("centralized_clip_service")
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")

os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# Device configuration
device = "cuda" if torch.cuda.is_available() else "cpu"
logger.info(f"Loading Centralized CLIP Model ({MODEL_NAME}) on device: {device}...")

try:
    model = CLIPModel.from_pretrained(MODEL_NAME).to(device)
    processor = CLIPProcessor.from_pretrained(MODEL_NAME)
    logger.info("✅ Centralized CLIP Model loaded successfully.")
except Exception as e:
    logger.error(f"❌ Failed to load CLIP model: {e}")
    raise e

# Global FAISS index & filename registry
index = None
filenames = []

def load_search_index():
    global index, filenames
    if os.path.exists(INDEX_FILE) and os.path.exists(MAPPING_FILE):
        try:
            logger.info(f"Loading FAISS index from {INDEX_FILE}...")
            index = faiss.read_index(INDEX_FILE)
            logger.info(f"Loading filename mapping from {MAPPING_FILE}...")
            with open(MAPPING_FILE, "r") as f:
                filenames = json.load(f)
            logger.info(f"✅ FAISS index ready. Total indexed items: {len(filenames)}.")
        except Exception as e:
            logger.error(f"❌ Error loading index files: {e}")
            index = None
            filenames = []
    else:
        logger.info("ℹ️ FAISS index not found. Starting with empty index.")
        index = None
        filenames = []

# Load index on startup
load_search_index()

def extract_embedding(image_source):
    """
    Generate L2-normalized 512-dim embedding vector from an image file path or PIL Image object.
    """
    try:
        if isinstance(image_source, str):
            image = Image.open(image_source).convert("RGB")
        elif isinstance(image_source, Image.Image):
            image = image_source.convert("RGB")
        else:
            image = Image.open(image_source).convert("RGB")
    except Exception as e:
        logger.error(f"❌ Error reading image: {e}")
        return None

    try:
        inputs = processor(images=image, return_tensors="pt")
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            features = model.get_image_features(**inputs)

        embedding = features.cpu().numpy().astype(np.float32)
        norm = np.linalg.norm(embedding, axis=1, keepdims=True)
        embedding = embedding / (norm + 1e-10)
        return embedding[0]
    except Exception as e:
        logger.error(f"❌ Embedding generation failed: {e}")
        return None

def is_jewelry_image(image_source, threshold=JEWELRY_VALIDATION_THRESHOLD):
    """
    Zero-shot CLIP classification to validate whether query image is jewelry.
    """
    try:
        if isinstance(image_source, str):
            image = Image.open(image_source).convert("RGB")
        elif isinstance(image_source, Image.Image):
            image = image_source.convert("RGB")
        else:
            image = Image.open(image_source).convert("RGB")
    except Exception as e:
        logger.error(f"❌ Image parsing failed during jewelry validation: {e}")
        return False

    labels = [
        "a photo of jewelry, ring, necklace, earring, bracelet, gold, gem, or diamond",
        "a photo of an animal, human, car, room, food, screen, or everyday object that is not jewelry"
    ]

    try:
        inputs = processor(text=labels, images=image, return_tensors="pt", padding=True)
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = model(**inputs)

        probs = outputs.logits_per_image.softmax(dim=1).cpu().numpy()[0]
        jewelry_prob = probs[0]
        logger.info(f"🔍 Validation: Jewelry={jewelry_prob:.2%}, Non-Jewelry={probs[1]:.2%}")
        return jewelry_prob >= threshold
    except Exception as e:
        logger.error(f"❌ CLIP zero-shot classification error: {e}")
        return False

def match_image(query_path, top_k=10, validate_jewelry=True):
    """
    Search FAISS index for top_k similar images.
    Returns list of [filename, score_integer_0_to_100].
    """
    global index, filenames

    if index is None or len(filenames) == 0:
        logger.warning("FAISS index is empty. Returning 0 matches.")
        return []

    if validate_jewelry:
        if not is_jewelry_image(query_path):
            raise ValueError("Uploaded image does not appear to be a valid jewelry item.")

    query_emb = extract_embedding(query_path)
    if query_emb is None:
        return []

    query_emb = np.expand_dims(query_emb, axis=0)
    distances, indices = index.search(query_emb, top_k)

    results = []
    for dist, idx in zip(distances[0], indices[0]):
        if idx < 0 or idx >= len(filenames):
            continue
        filename = filenames[idx]
        score = int(dist * 100)
        score = max(0, min(100, score))
        results.append([filename, score])

    return results

def add_image_to_index(image_path, original_filename=None):
    """
    Incremental addition of single product image vector to FAISS index.
    """
    global index, filenames

    filename = original_filename if original_filename else os.path.basename(image_path)
    if filename in filenames:
        logger.info(f"Image {filename} is already indexed. Skipping.")
        return False

    emb = extract_embedding(image_path)
    if emb is None:
        return False

    emb = np.expand_dims(emb, axis=0)

    if index is None:
        logger.info(f"Initializing new FAISS IndexFlatIP ({EMBEDDING_DIM}-dim)...")
        index = faiss.IndexFlatIP(EMBEDDING_DIM)
        filenames = []

    index.add(emb)
    filenames.append(filename)

    os.makedirs(os.path.dirname(INDEX_FILE), exist_ok=True)
    faiss.write_index(index, INDEX_FILE)
    with open(MAPPING_FILE, "w") as f:
        json.dump(filenames, f)

    logger.info(f"✅ Incrementally added '{filename}'. Total items: {len(filenames)}.")
    return True

def rebuild_index(gallery_dir=DEFAULT_GALLERY_DIR):
    """
    Batch index build scanning gallery directory and caching embeddings.
    """
    global index, filenames

    if not os.path.exists(gallery_dir):
        logger.error(f"Gallery directory does not exist: {gallery_dir}")
        return False, f"Directory not found: {gallery_dir}"

    valid_exts = (".jpg", ".jpeg", ".png", ".webp")
    image_files = [f for f in os.listdir(gallery_dir) if f.lower().endswith(valid_exts)]
    total_files = len(image_files)

    if total_files == 0:
        return False, "No valid image files found to index."

    embedding_cache = {}
    if os.path.exists(CACHE_FILE):
        try:
            with open(CACHE_FILE, "rb") as f:
                embedding_cache = pickle.load(f)
            logger.info(f"Loaded {len(embedding_cache)} cached embeddings.")
        except Exception as e:
            logger.warning(f"Failed to load cache: {e}. Starting fresh.")

    embeddings = []
    indexed_filenames = []
    cache_updated = False

    for idx, filename in enumerate(image_files):
        img_path = os.path.join(gallery_dir, filename)

        if filename in embedding_cache:
            emb = embedding_cache[filename]
            embeddings.append(emb)
            indexed_filenames.append(filename)
        else:
            emb = extract_embedding(img_path)
            if emb is not None:
                embeddings.append(emb)
                indexed_filenames.append(filename)
                embedding_cache[filename] = emb
                cache_updated = True

    if len(embeddings) == 0:
        return False, "No valid embeddings were generated."

    if cache_updated:
        try:
            os.makedirs(os.path.dirname(CACHE_FILE), exist_ok=True)
            with open(CACHE_FILE, "wb") as f:
                pickle.dump(embedding_cache, f)
        except Exception as e:
            logger.error(f"Failed to write cache: {e}")

    embeddings_matrix = np.vstack(embeddings).astype(np.float32)

    new_index = faiss.IndexFlatIP(EMBEDDING_DIM)
    new_index.add(embeddings_matrix)

    os.makedirs(os.path.dirname(INDEX_FILE), exist_ok=True)
    faiss.write_index(new_index, INDEX_FILE)
    with open(MAPPING_FILE, "w") as f:
        json.dump(indexed_filenames, f)

    index = new_index
    filenames = indexed_filenames

    logger.info(f"✅ Rebuilt index with {len(filenames)} items.")
    return True, f"Successfully indexed {len(filenames)} images."

def get_status():
    return {
        "model": MODEL_NAME,
        "device": device,
        "index_loaded": index is not None,
        "total_items": len(filenames) if filenames else 0,
        "index_file": INDEX_FILE,
        "mapping_file": MAPPING_FILE
    }