import json
import os
import chromadb

class VectorDBManager:
    def __init__(self, collection_name: str = "shl_catalog"):
        # Persistent storage locally, fits stateless server startups perfectly
        self.client = chromadb.Client()
        self.collection = self.client.get_or_create_collection(name=collection_name)

    def initialize_catalog(self, catalog_path: str):
        """Loads and indexes the catalog JSON file if the collection is empty."""
        if self.collection.count() > 0:
            return  # Already populated

        if not os.path.exists(catalog_path):
            raise FileNotFoundError(f"Catalog file not found at: {catalog_path}")

        # Added utf-8 encoding and strict=False to bypass dirty JSON characters
        with open(catalog_path, "r", encoding="utf-8") as f:
            catalog = json.load(f, strict=False)

        docs = []
        metadatas = []
        ids = []

        for item in catalog:
            # Combine content for meaningful semantic retrieval
            text_chunk = (
                f"Name: {item['name']}. "
                f"Description: {item['description']}. "
                f"Job Levels: {', '.join(item.get('job_levels', []))}"
            )
            docs.append(text_chunk)

            # Metadata aligns precisely with the output Recommendation model
            metadatas.append({
                "name": item["name"],
                "url": item.get("link", ""),
                "test_type": item.get("keys", ["General"])[0] if item.get("keys") else "General"
            })
            ids.append(str(item["entity_id"]))

        self.collection.add(documents=docs, metadatas=metadatas, ids=ids)

    def query_catalog(self, query_text: str, n_results: int = 8) -> str:
        """Queries the collection and returns a structured string for prompt context.
        Kept for backward compatibility — prefer query_catalog_structured for new code,
        since the structured form is what enables post-generation validation."""
        results = self.collection.query(query_texts=[query_text], n_results=n_results)

        context_str = ""
        if results['documents'] and results['documents'][0]:
            for i, doc in enumerate(results['documents'][0]):
                meta = results['metadatas'][0][i]
                context_str += f"- {doc} (URL: {meta['url']}, Type: {meta['test_type']})\n"
        return context_str

    def query_catalog_structured(self, query_text: str, n_results: int = 8) -> list:
        """Same query, returned as a list of dicts (name/url/test_type/description)
        instead of a flattened string. This is the ground truth used both to build
        the responder's candidate list AND to validate its output afterward —
        anything the model returns that isn't in this list gets dropped."""
        results = self.collection.query(query_texts=[query_text], n_results=n_results)

        candidates = []
        if results['documents'] and results['documents'][0]:
            for i, doc in enumerate(results['documents'][0]):
                meta = results['metadatas'][0][i]
                candidates.append({
                    "name": meta["name"],
                    "url": meta["url"],
                    "test_type": meta["test_type"],
                    "description": doc,
                })
        return candidates

# Global singleton instance
db_manager = VectorDBManager()