# rag_utils.py
import faiss
import pickle
import numpy as np
from typing import Dict, Optional

FAISS_INDEX_PATH = "faiss_travel_index"
PKL_PATH = "faiss_travel_index.pkl"

def get_relevant_plan(travel_data: Dict) -> Optional[str]:
    try:
        index = faiss.read_index(FAISS_INDEX_PATH)
        with open(PKL_PATH, "rb") as f:
            plan_data = pickle.load(f)

        if not isinstance(plan_data, dict):
            print("Warning: plan_data is not a dict, assuming list with index as key")
            plan_data = {i: plan for i, plan in enumerate(plan_data)}

        # --- IMPORTANT ---
        # This part needs a real embedding model.
        # You need to install 'sentence-transformers' and use it here.
        # For now, it uses a random vector, which WON'T give relevant results.
        # Example (uncomment and install 'sentence-transformers'):
        # from sentence_transformers import SentenceTransformer
        # model = SentenceTransformer('all-MiniLM-L6-v2')
        query_text = f"{travel_data.get('destination', '')} {travel_data.get('location_type', '')} {travel_data.get('province', '')} {' '.join(travel_data.get('preferences', []))}"
        # query_vector = model.encode([query_text]).astype('float32') # Use this line with SentenceTransformers
        query_vector = np.random.rand(1, index.d).astype('float32')  # Dummy vector - REPLACE THIS
        # --- END IMPORTANT ---

        k = 1
        distances, indices = index.search(query_vector, k)

        if indices.size > 0:
            relevant_index = indices[0][0]
            # Check if the index is valid for plan_data
            if relevant_index in plan_data:
                relevant_plan = plan_data[relevant_index]
                print(f"Found relevant plan at index {relevant_index} with distance {distances[0][0]}")
                return relevant_plan
            else:
                 # Try to map the index if plan_data is a list
                if isinstance(plan_data, dict) and relevant_index < len(plan_data):
                     plan_keys = list(plan_data.keys())
                     mapped_key = plan_keys[relevant_index]
                     relevant_plan = plan_data[mapped_key]
                     print(f"Found relevant plan at mapped index {mapped_key} with distance {distances[0][0]}")
                     return relevant_plan
                else:
                    print(f"No plan found at index {relevant_index}")

    except FileNotFoundError:
        print(f"Error: FAISS index or pickle file not found at {FAISS_INDEX_PATH} or {PKL_PATH}")
        return None
    except Exception as e:
        print(f"Error in RAG functionality: {e}")
        return None

    print("No relevant plan found.")
    return None