# fitz_ai/vector_db/loader.py
"""
Vector DB plugin loader.

Fitz uses PostgreSQL + pgvector for unified storage. This module loads
the pgvector plugin which handles both local (pgserver) and external
PostgreSQL deployments.

Usage:
    from fitz_ai.vector_db import get_vector_db_plugin

    # Default: local mode (embedded PostgreSQL via pgserver)
    db = get_vector_db_plugin("pgvector")

    # External PostgreSQL
    db = get_vector_db_plugin("pgvector", mode="external",
                              connection_string="postgresql://...")

    # Standard interface
    db.upsert("collection", points)
    results = db.search("collection", vector, limit=10)
    count = db.count("collection")
"""

from __future__ import annotations

import importlib
import os
import uuid
from pathlib import Path
from typing import Any, Dict, List, Optional

import httpx
import yaml
from jinja2 import Template

from fitz_ai.core.utils import extract_path
from fitz_ai.vector_db.base import SearchResult


def _string_to_uuid(s: str) -> str:
    """
    Convert an arbitrary string ID to a deterministic UUID.

    Uses UUID5 with a fixed namespace for determinism - same input
    always produces the same UUID.
    """
    from fitz_ai.core.constants import UUID_NAMESPACE_DNS

    return str(uuid.uuid5(UUID_NAMESPACE_DNS, s))


class VectorDBSpec:
    """Parsed vector DB specification from YAML file."""

    def __init__(self, yaml_path: Path):
        with open(yaml_path) as f:
            self.spec = yaml.safe_load(f)

        self.name = self.spec["name"]
        self.type = self.spec["type"]
        self.description = self.spec.get("description", "")
        self.connection = self.spec["connection"]
        self.operations = self.spec["operations"]
        self.features = self.spec.get("features", {})

    def is_local(self) -> bool:
        """Check if this is a local (non-HTTP) plugin."""
        return self.connection.get("type") == "local"

    def get_local_class_path(self) -> Optional[str]:
        """Get Python class path for local implementations."""
        return self.operations.get("python_class")

    def requires_uuid_ids(self) -> bool:
        """Check if this vector DB requires UUID IDs (from YAML spec)."""
        return self.features.get("requires_uuid_ids", False)

    def get_auto_detect_service(self) -> Optional[str]:
        """Get the service name for auto-detection (from YAML spec)."""
        return self.features.get("auto_detect")

    def supports_namespaces(self) -> bool:
        """Check if this vector DB uses namespaces (like Pinecone)."""
        return self.features.get("supports_namespaces", False)

    def build_base_url(self, **kwargs) -> str:
        """Build base URL from template and kwargs."""
        template = self.connection.get("base_url", "")

        # Apply defaults from YAML spec
        context = {}
        if "default_host" in self.connection:
            context["host"] = kwargs.get("host", self.connection["default_host"])
        if "default_port" in self.connection:
            context["port"] = kwargs.get("port", self.connection["default_port"])
        if "default_environment" in self.connection:
            context["environment"] = kwargs.get(
                "environment", self.connection["default_environment"]
            )

        # Add all other kwargs (for Pinecone: index_name, project_id, etc.)
        context.update(kwargs)

        return Template(template).render(context)

    def get_auth_headers(self) -> Dict[str, str]:
        """Get authentication headers from environment."""
        if "auth" not in self.connection:
            return {}

        auth = self.connection["auth"]

        # Optional auth - skip if not present
        if auth.get("optional") and not os.getenv(auth.get("env_var", "")):
            return {}

        env_var = auth.get("env_var", "")
        api_key = os.getenv(env_var) if env_var else None

        if not api_key:
            if not auth.get("optional"):
                raise ValueError(f"{env_var} not set")
            return {}

        auth_type = auth.get("type", "bearer")
        header = auth.get("header", "Authorization")

        if auth_type == "bearer":
            scheme = auth.get("scheme", "Bearer")
            if scheme:
                return {header: f"{scheme} {api_key}"}
            else:
                # No scheme (like Qdrant's api-key header)
                return {header: api_key}
        elif auth_type == "custom":
            return {header: api_key}

        return {}

    def render_template(self, template: Any, context: Dict[str, Any]) -> Any:
        """Render Jinja2 templates in values recursively."""
        if isinstance(template, str):
            if "{{" in template:
                stripped = template.strip()
                # Direct variable substitution (no string conversion)
                if stripped.startswith("{{") and stripped.endswith("}}"):
                    var_name = stripped[2:-2].strip()
                    if var_name in context:
                        return context[var_name]
                return Template(template).render(context)
            return template
        elif isinstance(template, dict):
            return {k: self.render_template(v, context) for k, v in template.items()}
        elif isinstance(template, list):
            return [self.render_template(item, context) for item in template]
        else:
            return template

    def transform_points(
        self, points: List[Dict], operation: str, context: Dict[str, Any] | None = None
    ) -> List[Dict]:
        """Transform standard points format to provider-specific format."""
        op_spec = self.operations.get(operation, {})
        transform = op_spec.get("point_transform", {})

        if not transform or transform.get("identity"):
            return points

        context = context or {}

        # Check for flatten_payload (Milvus style)
        flatten_payload = transform.pop("flatten_payload", False)

        transformed = []
        for point in points:
            new_point = {}

            # Apply field mapping
            for target_field, source_field in transform.items():
                if source_field in point:
                    new_point[target_field] = point[source_field]

            # Flatten payload if requested
            if flatten_payload and "payload" in point:
                payload = point.get("payload", {})
                if isinstance(payload, dict):
                    for k, v in payload.items():
                        if k not in new_point:
                            new_point[k] = v

            # Add collection if specified in transform
            if "class" in transform and "collection" in context:
                new_point["class"] = context["collection"]

            transformed.append(new_point)

        return transformed


class GenericVectorDBPlugin:
    """
    Generic vector DB plugin that executes YAML specifications.

    This class handles all HTTP-based vector databases by reading their
    YAML specs and executing the defined operations. No provider-specific
    code needed - just drop a YAML file.

    Implements the VectorDBPlugin protocol.
    """

    plugin_type = "vector_db"

    def __init__(self, spec: VectorDBSpec, **kwargs):
        self.spec = spec
        self.plugin_name = spec.name
        self.kwargs = kwargs

        base_url = spec.build_base_url(**kwargs)
        headers = spec.get_auth_headers()

        self.client = httpx.Client(
            base_url=base_url,
            headers=headers,
            timeout=30.0,
        )

        self._vector_dim: Optional[int] = None

    def _convert_point_ids(self, points: List[Dict]) -> List[Dict]:
        """Convert string IDs to UUIDs if required by the vector DB (from YAML spec)."""
        if not self.spec.requires_uuid_ids():
            return points

        converted = []
        for point in points:
            new_point = dict(point)
            original_id = point.get("id")

            if original_id is not None and isinstance(original_id, str):
                try:
                    uuid.UUID(original_id)
                except ValueError:
                    new_point["id"] = _string_to_uuid(original_id)
                    if "payload" not in new_point:
                        new_point["payload"] = {}
                    new_point["payload"]["_original_id"] = original_id

            converted.append(new_point)

        return converted

    def search(
        self,
        collection_name: str,
        query_vector: List[float],
        limit: int,
        with_payload: bool = True,
    ) -> List[SearchResult]:
        """Search for similar vectors in collection."""
        if "search" not in self.spec.operations:
            raise NotImplementedError(f"{self.plugin_name} does not support search")

        op = self.spec.operations["search"]

        context = {
            "collection": collection_name,
            "query_vector": query_vector,
            "limit": limit,
            "with_payload": with_payload,
            **self.kwargs,
        }

        endpoint = Template(op["endpoint"]).render(context)
        body = self.spec.render_template(op.get("body", {}), context)

        response = self.client.request(
            method=op["method"],
            url=endpoint,
            json=body if body else None,
        )
        response.raise_for_status()

        data = response.json()
        results_path = op["response"]["results_path"]
        results = extract_path(data, results_path, default=[], strict=False)

        if not results:
            return []

        mapping = op["response"]["mapping"]
        search_results = []

        for item in results:
            result_id = extract_path(item, mapping["id"], strict=False)
            result_score = extract_path(item, mapping.get("score", ""), strict=False)
            result_payload = extract_path(
                item, mapping.get("payload", ""), default={}, strict=False
            )

            # Restore original ID if we converted it
            if result_payload and "_original_id" in result_payload:
                result_id = result_payload["_original_id"]

            search_result = SearchResult(
                id=str(result_id),
                score=float(result_score) if result_score is not None else None,
                payload=result_payload if result_payload else {},
            )
            search_results.append(search_result)

        return search_results

    def upsert(self, collection: str, points: List[Dict], defer_persist: bool = False) -> None:
        """Insert or update points in collection."""
        # Note: defer_persist is ignored for YAML-based plugins (no local persistence)
        if "upsert" not in self.spec.operations:
            raise NotImplementedError(f"{self.plugin_name} does not support upsert")

        if points and "vector" in points[0]:
            self._vector_dim = len(points[0]["vector"])

        op = self.spec.operations["upsert"]

        converted_points = self._convert_point_ids(points)

        # Build context for template rendering and point transformation
        context = {
            "collection": collection,
            "vector_dim": self._vector_dim,
            **self.kwargs,
        }

        transformed_points = self.spec.transform_points(converted_points, "upsert", context)
        context["points"] = transformed_points

        endpoint = Template(op["endpoint"]).render(context)
        body = self.spec.render_template(op.get("body", {}), context)

        response = self.client.request(
            method=op["method"],
            url=endpoint,
            json=body,
        )

        # Auto-create collection on 404 if configured
        if response.status_code == 404 and op.get("auto_create_collection"):
            self._auto_create_collection(collection, op, context)
            response = self.client.request(
                method=op["method"],
                url=endpoint,
                json=body,
            )

        response.raise_for_status()

    def _auto_create_collection(
        self,
        collection: str,
        op: Dict[str, Any],
        context: Dict[str, Any],
    ) -> None:
        """Auto-create a collection before upsert."""
        create_endpoint = op.get("create_collection_endpoint")
        create_method = op.get("create_collection_method", "PUT")
        create_body_template = op.get("create_collection_body", {})

        if not create_endpoint:
            raise ValueError(
                f"auto_create_collection is true but create_collection_endpoint "
                f"is not specified in {self.plugin_name} YAML"
            )

        endpoint = Template(create_endpoint).render(context)
        body = self.spec.render_template(create_body_template, context)

        response = self.client.request(
            method=create_method,
            url=endpoint,
            json=body,
        )

        # Accept 200, 201, 409 (already exists)
        if response.status_code not in (200, 201, 409):
            response.raise_for_status()

    def count(self, collection: str) -> int:
        """Get the number of points in a collection."""
        if "count" not in self.spec.operations:
            # Fall back to get_stats if count not available
            if "get_stats" in self.spec.operations:
                stats = self.get_collection_stats(collection)
                return stats.get("points_count", stats.get("vectorCount", 0))
            raise NotImplementedError(f"{self.plugin_name} does not support count")

        op = self.spec.operations["count"]
        context = {"collection": collection, **self.kwargs}

        endpoint = Template(op["endpoint"]).render(context)
        body = self.spec.render_template(op.get("body", {}), context)

        if op["method"] == "GET":
            response = self.client.get(endpoint)
        else:
            response = self.client.request(
                method=op["method"],
                url=endpoint,
                json=body if body else None,
            )
        response.raise_for_status()

        data = response.json()
        count_path = op["response"]["count_path"]

        # Handle dynamic path (e.g., namespaces.{{collection}}.vectorCount)
        resolved_path = Template(count_path).render(context)
        count = extract_path(data, resolved_path, default=0, strict=False)

        return int(count) if count else 0

    def create_collection(self, name: str, vector_size: int) -> None:
        """Create a new collection."""
        if "create_collection" not in self.spec.operations:
            raise NotImplementedError(f"{self.plugin_name} does not support create_collection")

        op = self.spec.operations["create_collection"]
        context = {
            "collection": name,
            "vector_size": vector_size,
            **self.kwargs,
        }

        endpoint = Template(op["endpoint"]).render(context)
        body = self.spec.render_template(op.get("body", {}), context)

        response = self.client.request(
            method=op["method"],
            url=endpoint,
            json=body,
        )
        response.raise_for_status()

    def delete_collection(self, name: str) -> None:
        """Delete a collection."""
        if "delete_collection" not in self.spec.operations:
            raise NotImplementedError(f"{self.plugin_name} does not support delete_collection")

        op = self.spec.operations["delete_collection"]
        context = {"collection": name, **self.kwargs}

        endpoint = Template(op["endpoint"]).render(context)
        body = self.spec.render_template(op.get("body", {}), context)

        if body:
            response = self.client.request(
                method=op["method"],
                url=endpoint,
                json=body,
            )
        else:
            response = self.client.request(
                method=op["method"],
                url=endpoint,
            )

        # Check for custom success codes
        success_codes = op.get("response", {}).get("success_codes", [200])
        if response.status_code not in success_codes:
            response.raise_for_status()

    def list_collections(self) -> List[str]:
        """List all collections."""
        if "list_collections" not in self.spec.operations:
            raise NotImplementedError(f"{self.plugin_name} does not support list_collections")

        op = self.spec.operations["list_collections"]
        context = {**self.kwargs}

        endpoint = Template(op["endpoint"]).render(context)
        body = self.spec.render_template(op.get("body", {}), context)

        if op["method"] == "GET":
            response = self.client.get(endpoint)
        else:
            response = self.client.request(
                method=op["method"],
                url=endpoint,
                json=body if body else None,
            )
        response.raise_for_status()

        data = response.json()
        collections_path = op["response"]["collections_path"]
        collections = extract_path(data, collections_path, default=[], strict=False)

        if not collections:
            return []

        # Handle different response formats
        response_config = op.get("response", {})

        # Dict with keys as collection names (Pinecone namespaces)
        if response_config.get("extract_keys"):
            if isinstance(collections, dict):
                return list(collections.keys())

        # List of strings (Milvus)
        if response_config.get("is_string_list"):
            return [str(c) for c in collections]

        # List of objects with name field (default)
        name_field = response_config.get("name_field", "name")
        return [c[name_field] if isinstance(c, dict) else str(c) for c in collections]

    def retrieve(
        self,
        collection_name: str,
        ids: List[str],
        with_payload: bool = True,
    ) -> List[Dict[str, Any]]:
        """
        Retrieve points by their IDs.

        Args:
            collection_name: Name of the collection
            ids: List of point IDs to retrieve
            with_payload: Whether to include payload in response

        Returns:
            List of dicts with 'id' and 'payload' keys
        """
        if "retrieve" not in self.spec.operations:
            raise NotImplementedError(f"{self.plugin_name} does not support retrieve")

        if not ids:
            return []

        op = self.spec.operations["retrieve"]

        # Convert IDs if required (e.g., string to UUID for Qdrant)
        converted_ids = ids
        if self.spec.requires_uuid_ids():
            converted_ids = []
            for point_id in ids:
                try:
                    uuid.UUID(point_id)
                    converted_ids.append(point_id)
                except ValueError:
                    converted_ids.append(_string_to_uuid(point_id))

        context = {
            "collection": collection_name,
            "ids": converted_ids,
            "with_payload": with_payload,
            **self.kwargs,
        }

        endpoint = Template(op["endpoint"]).render(context)
        body = self.spec.render_template(op.get("body", {}), context)

        if op["method"] == "GET":
            response = self.client.get(endpoint)
        else:
            response = self.client.request(
                method=op["method"],
                url=endpoint,
                json=body if body else None,
            )
        response.raise_for_status()

        data = response.json()
        results_path = op["response"]["results_path"]
        results = extract_path(data, results_path, default=[], strict=False)

        if not results:
            return []

        mapping = op["response"]["mapping"]
        retrieved = []

        for item in results:
            result_id = extract_path(item, mapping["id"], strict=False)
            result_payload = extract_path(
                item, mapping.get("payload", ""), default={}, strict=False
            )

            # Restore original ID if we converted it
            if result_payload and "_original_id" in result_payload:
                result_id = result_payload["_original_id"]

            retrieved.append(
                {
                    "id": str(result_id),
                    "payload": result_payload if result_payload else {},
                }
            )

        return retrieved

    def get_collection_stats(self, collection: str) -> Dict[str, Any]:
        """Get statistics for a collection."""
        if "get_stats" not in self.spec.operations:
            raise NotImplementedError(f"{self.plugin_name} does not support get_collection_stats")

        op = self.spec.operations["get_stats"]
        context = {"collection": collection, **self.kwargs}

        endpoint = Template(op["endpoint"]).render(context)
        body = self.spec.render_template(op.get("body", {}), context)

        if op["method"] == "GET":
            response = self.client.get(endpoint)
        else:
            response = self.client.request(
                method=op["method"],
                url=endpoint,
                json=body if body else None,
            )
        response.raise_for_status()

        data = response.json()
        stats_path = op["response"]["stats_path"]

        if not stats_path:
            return data

        # Handle dynamic path
        resolved_path = Template(stats_path).render(context)
        return extract_path(data, resolved_path, default={}, strict=False)

    def __del__(self):
        """Cleanup HTTP client."""
        if hasattr(self, "client"):
            try:
                self.client.close()
            except Exception:
                pass


def _get_vector_db_plugin_dirs() -> list[Path]:
    """
    Get all directories to search for vector DB plugins.

    Returns paths in priority order:
    1. User plugins (~/.fitz/plugins/vector_db/) - highest priority
    2. Package plugins (fitz_ai/vector_db/plugins/)
    """
    from fitz_ai.core.paths import FitzPaths

    dirs = []

    # User plugins (highest priority)
    user_dir = FitzPaths.user_vector_db_plugins()
    if user_dir.exists():
        dirs.append(user_dir)

    # Package plugins
    dirs.append(Path(__file__).parent / "plugins")

    return dirs


def load_vector_db_spec(plugin_name: str) -> VectorDBSpec:
    """
    Load vector DB specification from YAML file.

    Searches user plugins first, then package plugins.
    """
    for plugins_dir in _get_vector_db_plugin_dirs():
        yaml_path = plugins_dir / f"{plugin_name}.yaml"
        if yaml_path.exists():
            return VectorDBSpec(yaml_path)

    # Not found
    raise ValueError(
        f"Vector DB plugin '{plugin_name}' not found. "
        f"Searched: {[str(d) for d in _get_vector_db_plugin_dirs()]}"
    )


def available_vector_db_plugins() -> list[str]:
    """
    List available vector DB plugins.

    Includes both package plugins and user plugins from ~/.fitz/plugins/.
    """
    plugins = set()

    for plugins_dir in _get_vector_db_plugin_dirs():
        if plugins_dir.exists():
            plugins.update(f.stem for f in plugins_dir.glob("*.yaml"))

    return sorted(plugins)


def _get_auto_detected_kwargs(spec: VectorDBSpec, kwargs: Dict[str, Any]) -> Dict[str, Any]:
    """
    Auto-detect connection parameters based on YAML spec.

    Uses the 'features.auto_detect' field to determine which service
    detection function to call from fitz_ai.core.detect.

    This is PROVIDER-AGNOSTIC - the detection logic is driven by YAML.
    """
    result = dict(kwargs)

    auto_detect_service = spec.get_auto_detect_service()
    if not auto_detect_service:
        return result

    # Only auto-detect if host/port not explicitly provided
    if "host" in result and "port" in result:
        return result

    try:
        from fitz_ai.core import detect

        # Map service name to detection function
        detection_functions = {
            "ollama": detect.get_ollama_connection,
        }

        detect_func = detection_functions.get(auto_detect_service)
        if detect_func:
            detected_host, detected_port = detect_func()

            if "host" not in result:
                result["host"] = detected_host
            if "port" not in result:
                result["port"] = detected_port

    except (ImportError, AttributeError):
        # Detection not available, fall back to YAML defaults
        pass

    return result


def create_vector_db_plugin(plugin_name: str, **kwargs):
    """
    Create the pgvector plugin instance.

    Fitz uses PostgreSQL + pgvector exclusively for unified storage.

    Args:
        plugin_name: Must be 'pgvector'
        **kwargs: Plugin configuration (mode, connection_string, hnsw_m, etc.)

    Returns:
        PgVectorDB plugin instance

    Examples:
        # Local mode (default) - embedded PostgreSQL via pgserver
        db = create_vector_db_plugin("pgvector")

        # External PostgreSQL
        db = create_vector_db_plugin(
            "pgvector",
            mode="external",
            connection_string="postgresql://user:pass@host:5432/db"
        )
    """
    spec = load_vector_db_spec(plugin_name)

    # Handle local implementations (e.g., FAISS)
    if spec.is_local():
        class_path = spec.get_local_class_path()
        if not class_path:
            raise ValueError(f"Local plugin '{plugin_name}' missing python_class specification")

        module_path, class_name = class_path.rsplit(".", 1)
        module = importlib.import_module(module_path)
        PluginClass = getattr(module, class_name)

        return PluginClass(**kwargs)

    # For HTTP-based plugins, auto-detect connection based on YAML spec
    resolved_kwargs = _get_auto_detected_kwargs(spec, kwargs)

    return GenericVectorDBPlugin(spec, **resolved_kwargs)


# Convenience alias
get_vector_db_plugin = create_vector_db_plugin


__all__ = [
    "VectorDBSpec",
    "GenericVectorDBPlugin",
    "load_vector_db_spec",
    "create_vector_db_plugin",
    "get_vector_db_plugin",
    "available_vector_db_plugins",
]
