"""Chroma vector store wrapper for rules and actions."""

from __future__ import annotations

import hashlib
from pathlib import Path
from typing import Any

import chromadb
from chromadb.config import Settings
from chromadb.utils.embedding_functions import ONNXMiniLM_L6_V2

from theow._core._logging import get_logger
from theow._core._models import Rule

logger = get_logger(__name__)


def extract_query_text(context: dict[str, Any], max_chars: int = 800) -> str:
    """Extract query text from context by finding the longest string value.

    Uses the tail of long strings because errors appear at the end of command
    output, and the embedding model (MiniLM-L6-V2, 256 tokens) truncates
    inputs beyond its context window.
    """
    longest = ""
    for value in context.values():
        if isinstance(value, str) and len(value) > len(longest):
            longest = value
    if len(longest) > max_chars:
        longest = longest[-max_chars:]
    return longest


# TODO: explore other available providers or maybe not.
_embedding_function = ONNXMiniLM_L6_V2(preferred_providers=["CPUExecutionProvider"])


class ChromaStore:
    """Wrapper around ChromaDB for rule and action storage."""

    def __init__(
        self,
        path: Path,
        embedding_model: str = "all-MiniLM-L6-v2",
    ) -> None:
        self._path = path
        self._embedding_model = embedding_model

        path.mkdir(parents=True, exist_ok=True)

        self._client = chromadb.PersistentClient(
            path=str(path),
            settings=Settings(anonymized_telemetry=False),
        )

        self._actions_collection = self._client.get_or_create_collection(
            name="theow-actions",
            metadata={"hnsw:space": "cosine"},
            embedding_function=_embedding_function,  # type: ignore[arg-type]
        )

        # Cache of known metadata keys from indexed rules
        self._metadata_keys: set[str] = set()

    def _get_collection(self, name: str) -> chromadb.Collection:
        return self._client.get_or_create_collection(
            name=name,
            metadata={"hnsw:space": "cosine"},
            embedding_function=_embedding_function,  # type: ignore[arg-type]
        )

    def index_rule(self, rule: Rule) -> None:
        """Index a rule into its collection."""
        collection = self._get_collection(rule.collection)

        metadata = rule.get_metadata()
        self._metadata_keys.update(
            k
            for k, v in metadata.items()
            if isinstance(v, str) and k not in ("type", "content_hash")
        )

        collection.upsert(
            ids=[rule.name],
            documents=[rule.get_embedding_text()],
            metadatas=[metadata],
        )
        logger.debug("Rule indexed", rule=rule.name, collection=rule.collection)

    def index_action(self, name: str, docstring: str, signature: str) -> None:
        """Index an action for semantic search."""
        embedding_text = f"{name}: {docstring}\nSignature: {signature}"
        content_hash = hashlib.sha256(embedding_text.encode()).hexdigest()[:16]

        existing = self._actions_collection.get(ids=[name], include=["metadatas"])
        metadatas = existing.get("metadatas")

        if existing["ids"] and metadatas:
            stored_hash = metadatas[0].get("content_hash", "")
            if stored_hash == content_hash:
                logger.debug("Action unchanged", action=name)
                return

        self._actions_collection.upsert(
            ids=[name],
            documents=[embedding_text],
            metadatas=[
                {"docstring": docstring, "signature": signature, "content_hash": content_hash}
            ],
        )
        logger.debug("Action indexed", action=name)

    def sync_rules(self, rules_dir: Path) -> None:
        """Sync rules from directory, re-embedding only changed files.

        Also removes stale entries from chroma that no longer have files.
        """
        if not rules_dir.exists():
            return

        synced: dict[str, set[str]] = {}

        for rule_file in rules_dir.glob("*.rule.yaml"):
            rule = self._sync_rule_file(rule_file)
            if rule:
                synced.setdefault(rule.collection, set()).add(rule.name)

        for collection, rule_names in synced.items():
            self._cleanup_stale(collection, rule_names)

    def _cleanup_stale(self, collection: str, valid_names: set[str]) -> None:
        """Remove chroma entries that don't have corresponding files."""
        coll = self._get_collection(collection)
        existing = coll.get(include=[])
        stale = [name for name in existing["ids"] if name not in valid_names]
        if stale:
            coll.delete(ids=stale)
            logger.debug("Removed stale rules", collection=collection, rules=stale)

    def _sync_rule_file(self, path: Path) -> Rule | None:
        """Sync a single rule file. Returns the rule for tracking."""
        rule = Rule.from_yaml(path)
        collection = self._get_collection(rule.collection)

        existing = collection.get(ids=[rule.name], include=["metadatas"])
        metadatas = existing.get("metadatas")

        if existing["ids"] and metadatas:
            stored_hash = metadatas[0].get("content_hash", "")
            if stored_hash == rule.content_hash():
                logger.debug("Rule unchanged", rule=rule.name)
                return rule

        self.index_rule(rule)
        return rule

    def get_rule(self, collection: str, name: str) -> Rule | None:
        """Get a rule by name from a collection."""
        coll = self._get_collection(collection)
        result = coll.get(ids=[name], include=["documents", "metadatas"])

        if not result["ids"]:
            return None

        # We don't store full rule YAML in Chroma, only embeddings
        # Rules must be loaded from files. this just confirms existence
        return None

    def query_rules(
        self,
        collection: str,
        query_text: str,
        metadata_filter: dict[str, Any] | None = None,
        n_results: int = 10,
    ) -> list[tuple[str, float, dict[str, Any]]]:
        """Query rules by semantic similarity.

        Returns list of (rule_name, distance, metadata).
        """
        coll = self._get_collection(collection)

        where = metadata_filter if metadata_filter else None

        try:
            results = coll.query(
                query_texts=[query_text],
                where=where,
                n_results=n_results,
                include=["distances", "metadatas"],
            )
        except Exception as e:
            logger.warning("Chroma query failed", error=str(e))
            return []

        ids = results.get("ids")
        distances = results.get("distances")
        metadatas = results.get("metadatas")

        if not ids or not ids[0] or not distances or not metadatas:
            return []

        return [(name, dist, meta) for name, dist, meta in zip(ids[0], distances[0], metadatas[0])]

    def query_actions(self, query_text: str, n_results: int = 5) -> list[dict[str, Any]]:
        """Query actions by semantic similarity."""
        try:
            results = self._actions_collection.query(
                query_texts=[query_text],
                n_results=n_results,
                include=["metadatas"],
            )
        except Exception as e:
            logger.warning("Chroma query failed", error=str(e))
            return []

        ids = results.get("ids")
        metadatas = results.get("metadatas")

        if not ids or not ids[0] or not metadatas:
            return []

        return [{"name": name, **meta} for name, meta in zip(ids[0], metadatas[0])]

    def list_rules(self, collection: str) -> list[str]:
        """List all rule names in a collection."""
        coll = self._get_collection(collection)
        result = coll.get(include=[])
        return result["ids"]

    def list_actions(self) -> list[str]:
        """List all action names."""
        result = self._actions_collection.get(include=[])
        return result["ids"]

    def get_metadata_keys(self) -> set[str]:
        """Get known metadata keys for filtering."""
        return self._metadata_keys

    def update_rule_stats(
        self,
        collection: str,
        rule_name: str,
        success: bool,
    ) -> None:
        """Update success/fail counts for a rule."""
        coll = self._get_collection(collection)

        existing = coll.get(ids=[rule_name], include=["metadatas", "documents"])
        metadatas = existing.get("metadatas")
        if not existing["ids"] or not metadatas:
            return

        old_meta = metadatas[0]
        new_meta: dict[str, Any] = dict(old_meta)

        if success:
            new_meta["success_count"] = int(new_meta.get("success_count", 0)) + 1
        else:
            new_meta["fail_count"] = int(new_meta.get("fail_count", 0)) + 1

        coll.update(
            ids=[rule_name],
            metadatas=[new_meta],
        )

    def get_all_rules_with_stats(self) -> list[dict[str, Any]]:
        """Get all rules with their stats for meow()."""
        all_stats = []

        for coll_name in self._client.list_collections():
            if coll_name.name == "theow-actions":
                continue

            coll = self._client.get_collection(coll_name.name)
            result = coll.get(include=["metadatas"])
            metadatas = result.get("metadatas")
            if not metadatas:
                continue

            for name, meta in zip(result["ids"], metadatas):
                all_stats.append(
                    {
                        "name": name,
                        "collection": coll_name.name,
                        "success_count": meta.get("success_count", 0),
                        "fail_count": meta.get("fail_count", 0),
                        "explored": meta.get("explored", False),
                    }
                )

        return all_stats
