import json
from pathlib import Path
from datetime import datetime

import chromadb
from chromadb.config import Settings

CHROMA_DIR = Path(".vectormind/chroma")
COLLECTION_NAME = "direction"
META_FILE = Path(".vectormind/index_meta.json")


# ---------------------------------------------------------------------------
# Client + Collection
# ---------------------------------------------------------------------------

def get_client() -> chromadb.Client:
    """Get persistent ChromaDB client pointed at .vectormind/chroma."""
    return chromadb.PersistentClient(
        path=str(CHROMA_DIR),
        settings=Settings(anonymized_telemetry=False)
    )


def get_collection():
    """Get or create the Direction collection."""
    client = get_client()
    return client.get_or_create_collection(
        name=COLLECTION_NAME,
        metadata={"hnsw:space": "cosine"}  # cosine similarity for code embeddings
    )


# ---------------------------------------------------------------------------
# Write operations
# ---------------------------------------------------------------------------

def upsert_chunks(
    collection,
    ids: list[str],
    embeddings: list[list[float]],
    documents: list[str],
    metadatas: list[dict]
):
    """Upsert chunks into ChromaDB collection."""
    if not ids:
        return

    collection.upsert(
        ids=ids,
        embeddings=embeddings,
        documents=documents,
        metadatas=metadatas
    )


def delete_file_chunks(collection, file_path: str):
    """Delete all chunks belonging to a specific file."""
    try:
        results = collection.get(
            where={"file": file_path}
        )
        if results and results["ids"]:
            collection.delete(ids=results["ids"])
    except Exception:
        pass  # File not yet indexed, nothing to delete


def delete_all():
    """Wipe entire collection. Used for full re-index."""
    client = get_client()
    try:
        client.delete_collection(COLLECTION_NAME)
    except Exception:
        pass
    return client.get_or_create_collection(
        name=COLLECTION_NAME,
        metadata={"hnsw:space": "cosine"}
    )


# ---------------------------------------------------------------------------
# Read operations
# ---------------------------------------------------------------------------

def query_chunks(
    collection,
    query_embedding: list[float],
    top_k: int = 5,
    file_filter: str = None
) -> dict:
    """
    Query ChromaDB for most relevant chunks.
    Returns dict with documents, metadatas, distances.
    """
    where = {"file": file_filter} if file_filter else None

    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k,
        where=where,
        include=["documents", "metadatas", "distances"]
    )

    return {
        "documents": results["documents"][0] if results["documents"] else [],
        "metadatas": results["metadatas"][0] if results["metadatas"] else [],
        "distances": results["distances"][0] if results["distances"] else []
    }


def get_file_chunks(collection, file_path: str) -> dict:
    """Retrieve all chunks for a specific file in order."""
    results = collection.get(
        where={"file": file_path},
        include=["documents", "metadatas"]
    )

    if not results or not results["ids"]:
        return {"documents": [], "metadatas": []}

    # Sort by chunk_index
    paired = sorted(
        zip(results["documents"], results["metadatas"]),
        key=lambda x: x[1].get("chunk_index", 0)
    )

    documents, metadatas = zip(*paired) if paired else ([], [])

    return {
        "documents": list(documents),
        "metadatas": list(metadatas)
    }


def list_indexed_files(collection) -> list[str]:
    """Return list of all unique files currently indexed."""
    try:
        results = collection.get(include=["metadatas"])
        if not results or not results["metadatas"]:
            return []
        files = {m["file"] for m in results["metadatas"] if "file" in m}
        return sorted(files)
    except Exception:
        return []


# ---------------------------------------------------------------------------
# Stats
# ---------------------------------------------------------------------------

def get_stats() -> dict:
    """Return DB stats for status command."""
    collection = get_collection()

    total_chunks = collection.count()

    files = list_indexed_files(collection)
    total_files = len(files)

    last_indexed = "never"
    if META_FILE.exists():
        with open(META_FILE) as f:
            meta = json.load(f)
            last_indexed = meta.get("last_indexed", "never")

    return {
        "total_chunks": total_chunks,
        "total_files": total_files,
        "last_indexed": last_indexed,
        "files": files
    }


# ---------------------------------------------------------------------------
# Similarity score helper
# ---------------------------------------------------------------------------

def distance_to_score(distance: float) -> float:
    """
    Convert cosine distance to a 0-1 relevance score.
    ChromaDB cosine distance: 0 = identical, 2 = opposite.
    Score: 1 = perfect match, 0 = no relevance.
    """
    return round(1 - (distance / 2), 4)