import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Dict, List, Optional, Tuple

import dataiku

from semantic_models.utils.spec_utils import should_use_distinct_values

logger = logging.getLogger(__name__)

# Thresholds for matching
FUZZY_THRESHOLD = 80.0  # 0-100, rapidfuzz WRatio score
SEMANTIC_MIN_SCORE = 75.0  # 0-100, cosine similarity percentage (raised from 70 to reduce false positives)


# =============================================================================
# In-Memory ChromaDB Singleton
# =============================================================================

_ephemeral_client: Optional[object] = None
_semantic_index_ready: Dict[str, bool] = {}  # Track per-collection readiness


def _get_ephemeral_chroma_client() -> Optional[object]:
    """Get or create the shared in-memory ChromaDB client."""
    global _ephemeral_client
    if _ephemeral_client is not None:
        return _ephemeral_client
    try:
        import chromadb

        _ephemeral_client = chromadb.EphemeralClient()
        return _ephemeral_client
    except ImportError:
        logger.warning("ChromaDB not installed, semantic search disabled")
        return None
    except AttributeError:
        # EphemeralClient may not exist in older ChromaDB versions
        logger.warning("ChromaDB version doesn't support EphemeralClient, semantic search disabled")
        return None
    except Exception as e:
        logger.warning("Failed to initialize ChromaDB: %s", e)
        return None


def is_semantic_index_ready(collection_name: str) -> bool:
    """Check if semantic index is populated for a collection."""
    return _semantic_index_ready.get(collection_name, False)


def _mark_semantic_index_ready(collection_name: str, ready: bool) -> None:
    """Mark semantic index as ready/not ready."""
    _semantic_index_ready[collection_name] = ready


# =============================================================================
# In-Memory Value Index Store
# =============================================================================

_value_index_store: Dict[str, Dict[str, List[str]]] = {}  # collection_name -> {column -> values}
_value_index_meta_store: Dict[str, "IndexMeta"] = {}  # collection_name -> IndexMeta


def get_value_index_data(collection_name: str) -> Optional[Dict[str, List[str]]]:
    """Get value index data from in-memory store."""
    return _value_index_store.get(collection_name)


def set_value_index_data(collection_name: str, data: Dict[str, List[str]], meta: "IndexMeta") -> None:
    """Store value index data in memory."""
    _value_index_store[collection_name] = data
    _value_index_meta_store[collection_name] = meta


def get_value_index_meta(collection_name: str) -> Optional["IndexMeta"]:
    """Get value index metadata from in-memory store."""
    return _value_index_meta_store.get(collection_name)


def clear_value_index_store() -> None:
    global _value_index_store, _value_index_meta_store, _semantic_index_ready

    # Clear value index stores
    _value_index_store.clear()
    _value_index_meta_store.clear()

    # Clear semantic index readiness flags
    _semantic_index_ready.clear()

    # Delete all Chroma collections
    client = _get_ephemeral_chroma_client()
    if client:
        try:
            for collection in client.list_collections():
                client.delete_collection(collection.name)
        except Exception as e:
            logger.warning("Failed to clear Chroma collections: %s", e)


def clear_source_index(source: "SourceKey") -> None:
    """Clear cached value + semantic index for a specific source."""
    collection_name = source.to_collection_name()
    _value_index_store.pop(collection_name, None)
    _value_index_meta_store.pop(collection_name, None)
    _semantic_index_ready.pop(collection_name, None)

    client = _get_ephemeral_chroma_client()
    if client:
        try:
            client.delete_collection(collection_name)
        except Exception:
            pass


class IndexStatus(str, Enum):
    """Status of the value index."""

    NOT_BUILT = "not_built"
    BUILDING = "building"
    READY = "ready"
    ERROR = "error"


@dataclass
class IndexMeta:
    """Metadata about a source's value index."""

    status: IndexStatus
    built_at: Optional[datetime] = None
    columns: List[str] = field(default_factory=list)
    value_count: int = 0
    error_message: Optional[str] = None

    def to_dict(self) -> dict:
        return {
            "status": self.status.value,
            "built_at": self.built_at.isoformat() if self.built_at else None,
            "columns": self.columns,
            "value_count": self.value_count,
            "error_message": self.error_message,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "IndexMeta":
        return cls(
            status=IndexStatus(data["status"]),
            built_at=datetime.fromisoformat(data["built_at"]) if data.get("built_at") else None,
            columns=data.get("columns", []),
            value_count=data.get("value_count", 0),
            error_message=data.get("error_message"),
        )


# =============================================================================
# Source Key Helpers
# =============================================================================


@dataclass
class SourceKey:
    """Identifies a semantic model entity within a specific version."""

    semantic_model_id: str
    version_id: str
    entity_name: str

    @classmethod
    def from_entity(cls, semantic_model_id: str, version_id: str, entity: dict) -> "SourceKey":
        """Create a SourceKey from an entity dict and version context."""
        entity_name = entity.get("name")
        if not entity_name:
            raise ValueError("Missing entity name for indexing")
        return cls(
            semantic_model_id=semantic_model_id,
            version_id=version_id,
            entity_name=entity_name,
        )

    def to_path_parts(self) -> List[str]:
        """Get path components for this source."""
        return [
            "semantic-models",
            self.semantic_model_id,
            "versions",
            self.version_id,
            "entities",
            self.entity_name,
        ]

    def to_collection_name(self) -> str:
        import hashlib

        base = f"idx-sm-{self.semantic_model_id}-{self.version_id}-{self.entity_name}"
        hash_input = f"{self.semantic_model_id}:{self.version_id}:{self.entity_name}"

        sanitized = "".join(c if c.isalnum() or c in "-_" else "-" for c in base)
        hash_suffix = hashlib.sha256(hash_input.encode()).hexdigest()[:8]
        prefix = sanitized[:54]

        return f"{prefix}-{hash_suffix}"

    def __hash__(self):
        return hash((self.semantic_model_id, self.version_id, self.entity_name))

    def __eq__(self, other):
        if not isinstance(other, SourceKey):
            return False
        return (
            self.semantic_model_id == other.semantic_model_id
            and self.version_id == other.version_id
            and self.entity_name == other.entity_name
        )


# =============================================================================
# Computed Attribute Detection
# =============================================================================


def _is_simple_column(attr: dict) -> bool:
    name = attr.get("name", "")
    attr_type = attr.get("type")
    if attr_type == "COLUMN":
        column = attr.get("column") or name
        return column == name
    expression = attr.get("pseudoSQLExpression")
    if not expression:
        return True
    return expression == name


# =============================================================================
# Source Index Class
# =============================================================================


@dataclass
class SourceIndex:
    source: SourceKey
    data: Dict[str, List[str]]
    meta: IndexMeta

    # Optional embedding LLM ID for semantic search (agent tool passes explicitly)
    embedding_llm_id: Optional[str] = field(default=None, repr=False, compare=False)

    # Private fields for lazy ChromaDB initialization
    _chroma_client: Optional[object] = field(default=None, repr=False, compare=False)
    _collection: Optional[object] = field(default=None, repr=False, compare=False)

    @property
    def collection_name(self) -> str:
        """Generate a valid ChromaDB collection name."""
        return self.source.to_collection_name()

    def _get_chroma_client(self) -> Optional[object]:
        """Get the shared in-memory ChromaDB client."""
        if self._chroma_client is not None:
            return self._chroma_client
        self._chroma_client = _get_ephemeral_chroma_client()
        return self._chroma_client

    def _get_collection(self) -> Optional[object]:
        """Lazily get or create the ChromaDB collection."""
        if self._collection is not None:
            return self._collection

        client = self._get_chroma_client()
        if client is None:
            return None

        try:
            self._collection = client.get_or_create_collection(
                name=self.collection_name,
                metadata={"hnsw:space": "cosine"},
            )
            return self._collection
        except Exception as e:
            logger.warning("Failed to get ChromaDB collection: %s", e)
            return None

    def get_values(self, column_name: str) -> List[str]:
        """Get values for a specific column."""
        return self.data.get(column_name, [])

    def search_best_value(self, column_name: str, term: str) -> Optional[Tuple[str, float, str]]:
        from rapidfuzz import fuzz, process

        values = self.get_values(column_name)
        if not values:
            return None

        term_norm = term.strip()
        if not term_norm:
            return None

        # 1) Exact match (case-insensitive)
        for v in values:
            if v.lower() == term_norm.lower():
                logger.debug("Exact match for '%s' on %s -> '%s'", term, column_name, v)
                return v, 100.0, "exact"

        # 2) Fuzzy match
        try:
            fuzzy_result = process.extractOne(term_norm, values, scorer=fuzz.WRatio, score_cutoff=FUZZY_THRESHOLD)
        except Exception as e:
            logger.warning("Fuzzy matching failed for '%s': %s", term_norm, e)
            fuzzy_result = None

        if fuzzy_result:
            matched_value, score, _ = fuzzy_result
            logger.debug("Fuzzy match for '%s' on %s -> '%s' (score=%.1f)", term, column_name, matched_value, score)
            return matched_value, float(score), "fuzzy"

        # 3) Semantic search (only if fuzzy didn't match)
        semantic_match = self._semantic_search(column_name, term_norm)
        if semantic_match is not None:
            return semantic_match

        logger.debug("No match found for '%s' on %s", term_norm, column_name)
        return None

    def _semantic_search(self, column_name: str, term: str) -> Optional[Tuple[str, float, str]]:
        """Semantic search using ChromaDB for a given term within one column."""
        embeddings = _embed_texts([term], self.embedding_llm_id)
        if not embeddings:
            return None

        collection = self._get_collection()
        if collection is None:
            return None

        try:
            results = collection.query(
                query_embeddings=embeddings,
                n_results=1,
                where={"column": column_name},
                include=["metadatas", "distances"],
            )
        except Exception as e:
            logger.warning("ChromaDB query failed for '%s': %s", term, e)
            return None

        metadatas = results.get("metadatas") or []
        distances = results.get("distances") or []

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

        value_meta = metadatas[0][0]
        value_str = value_meta.get("value")
        if value_str is None:
            return None

        # Cosine distance: [0, 2] where 0 = identical, 2 = opposite
        # Convert to score: similarity = 1 - (distance / 2), then * 100
        dist = float(distances[0][0]) if distances and distances[0] else 1.0
        score = max(0, (1 - dist / 2) * 100)

        if score < SEMANTIC_MIN_SCORE:
            logger.debug(
                "Semantic match too weak for '%s' on %s: '%s' (dist=%.3f, score=%.1f)",
                term,
                column_name,
                value_str,
                dist,
                score,
            )
            return None

        logger.debug(
            "Semantic match for '%s' on %s -> '%s' (dist=%.3f, score=%.1f)",
            term,
            column_name,
            value_str,
            dist,
            score,
        )
        return value_str, score, "semantic"


def _embed_texts(texts: List[str], embedding_llm_id: Optional[str] = None) -> List[List[float]]:
    if not texts:
        return []

    if not embedding_llm_id:
        logger.debug("No embedding LLM provided, semantic search disabled")
        return []

    llm_id = embedding_llm_id

    if not llm_id:
        logger.debug("No embedding LLM configured, semantic search disabled")
        return []

    try:
        client = dataiku.api_client()
        project = client.get_default_project()
        llm = project.get_llm(llm_id)
        BATCH_SIZE = 50  # Conservative batch size
        all_embeddings: List[List[float]] = []

        for i in range(0, len(texts), BATCH_SIZE):
            batch = texts[i : i + BATCH_SIZE]
            query = llm.new_embeddings()
            for t in batch:
                query.add_text(t)
            response = query.execute()
            batch_embeddings = response.get_embeddings()
            all_embeddings.extend(batch_embeddings)

        return all_embeddings

    except Exception as e:
        logger.warning("Failed to generate embeddings: %s", e)
        return []


def get_source_index_status(source: SourceKey) -> IndexMeta:
    collection_name = source.to_collection_name()
    meta = get_value_index_meta(collection_name)
    return meta if meta else IndexMeta(status=IndexStatus.NOT_BUILT)


def load_source_index(
    source: SourceKey,
    embedding_llm_id: Optional[str] = None,
) -> Optional[SourceIndex]:
    collection_name = source.to_collection_name()

    data = get_value_index_data(collection_name)
    meta = get_value_index_meta(collection_name)

    if data is None or meta is None or meta.status != IndexStatus.READY:
        return None

    return SourceIndex(
        source=source,
        data=data,
        meta=meta,
        embedding_llm_id=embedding_llm_id,
    )


def load_indices_for_spec(
    spec: dict,
    *,
    semantic_model_id: str,
    version_id: str,
    embedding_llm_id: Optional[str] = None,
) -> Dict[str, SourceIndex]:
    entity_indices: Dict[str, SourceIndex] = {}
    loaded_sources: Dict[SourceKey, SourceIndex] = {}

    for entity in spec.get("entities", []):
        entity_id = entity["name"]
        if entity.get("type") != "DATASET":
            continue

        # Check if entity has any resolvable simple columns
        has_resolvable_simple = any(
            attr.get("resolveInUserRequests", False) and _is_simple_column(attr) and should_use_distinct_values(attr)
            for attr in entity.get("attributes", [])
        )
        if not has_resolvable_simple:
            continue

        dataset_ref = entity.get("datasetRef") or ""
        if "." not in dataset_ref:
            continue

        source = SourceKey.from_entity(semantic_model_id, version_id, entity)

        # Reuse already-loaded index for same source
        if source in loaded_sources:
            entity_indices[entity_id] = loaded_sources[source]
            continue

        # Load or build the source index on demand
        index = load_source_index(source, embedding_llm_id)
        if not index:
            try:
                from semantic_models.services.column_values import fetch_column_values

                manual_values_by_column: Dict[str, List[str]] = {}
                for attr in entity.get("attributes", []):
                    attr_name = attr.get("name")
                    if not attr_name:
                        continue
                    values = attr.get("manualValues") or []
                    if values:
                        manual_values_by_column[attr_name] = values

                meta = build_source_index_in_memory(
                    source=source,
                    columns=get_resolvable_columns_for_entity(entity),
                    fetch_values_fn=fetch_column_values,
                    limit=1000,
                    embedding_llm_id=embedding_llm_id,
                    manual_values_by_column=manual_values_by_column,
                )
                if meta.status == IndexStatus.READY:
                    index = load_source_index(source, embedding_llm_id)
            except Exception as e:
                logger.warning("Failed to build index for %s: %s", source.to_path_parts(), e)

        if index:
            loaded_sources[source] = index
            entity_indices[entity_id] = index

    logger.info("Loaded %d source indices for %d entities", len(loaded_sources), len(entity_indices))
    return entity_indices


# =============================================================================
# Index Building
# =============================================================================


def _populate_chroma_for_source(
    source: SourceKey,
    data: Dict[str, List[str]],
    embedding_llm_id: Optional[str] = None,
) -> None:
    collection_name = source.to_collection_name()

    _mark_semantic_index_ready(collection_name, False)

    client = _get_ephemeral_chroma_client()
    if client is None:
        logger.info("ChromaDB unavailable, skipping semantic index")
        return  # Stays marked as not ready

    try:
        client.delete_collection(name=collection_name)
        logger.debug("Deleted existing ChromaDB collection '%s'", collection_name)
    except Exception:
        pass  # Collection didn't exist

    try:
        collection = client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},
        )
    except Exception as e:
        logger.warning("Failed to create ChromaDB collection: %s", e)
        return  # Stays marked as not ready

    # Collect all values
    ids: List[str] = []
    documents: List[str] = []
    metadatas: List[dict] = []

    for column_name, values in data.items():
        for idx, value in enumerate(values):
            doc_id = f"{column_name}:{idx}"
            ids.append(doc_id)
            documents.append(value)
            metadatas.append(
                {
                    "column": column_name,
                    "value": value,
                }
            )

    if not ids:
        # Empty source is still considered "indexed" - just nothing to search
        logger.info("No values to populate in ChromaDB for %s", collection_name)
        _mark_semantic_index_ready(collection_name, True)  # Empty but ready
        return

    logger.info("Populating ChromaDB collection '%s' with %d values", collection_name, len(ids))

    # Generate embeddings
    embeddings = _embed_texts(documents, embedding_llm_id)

    # If embeddings failed, semantic search won't work - don't mark as ready
    if not embeddings:
        logger.warning(
            "Embedding generation failed for %s - semantic search disabled for this source",
            collection_name,
        )
        return  # Stays marked as not ready

    # Add to collection in batches
    BATCH_SIZE = 50
    for start in range(0, len(ids), BATCH_SIZE):
        end = start + BATCH_SIZE
        batch_ids = ids[start:end]
        batch_docs = documents[start:end]
        batch_meta = metadatas[start:end]
        batch_emb = embeddings[start:end]

        try:
            collection.add(
                ids=batch_ids,
                documents=batch_docs,
                metadatas=batch_meta,
                embeddings=batch_emb,
            )
        except Exception as e:
            logger.warning("Failed to add batch to ChromaDB: %s", e)
            return  # Stays marked as not ready

    # ONLY mark ready after successful completion with embeddings
    _mark_semantic_index_ready(collection_name, True)
    logger.info("ChromaDB population complete for %s", collection_name)


def build_source_index_in_memory(
    source: SourceKey,
    columns: List[str],
    fetch_values_fn: callable,
    limit: int = 500,
    embedding_llm_id: Optional[str] = None,
    manual_values_by_column: Optional[Dict[str, List[str]]] = None,
) -> IndexMeta:
    collection_name = source.to_collection_name()

    # Reset semantic readiness BEFORE building to avoid stale state on errors
    _mark_semantic_index_ready(collection_name, False)

    # Delete existing Chroma collection to ensure clean state
    client = _get_ephemeral_chroma_client()
    if client:
        try:
            client.delete_collection(name=collection_name)
        except Exception:
            pass  # Collection may not exist

    try:
        data: Dict[str, List[str]] = {}
        total_values = 0
        indexed_columns: List[str] = []

        manual_values_by_column = manual_values_by_column or {}

        for column_name in columns:
            logger.info("Fetching values for column '%s'", column_name)

            try:
                if column_name in manual_values_by_column and manual_values_by_column[column_name]:
                    values = manual_values_by_column[column_name]
                    logger.info("  -> using MANUAL values for '%s' (%d values)", column_name, len(values))
                else:
                    values = fetch_values_fn(source, column_name, limit)
                    if values:
                        logger.info("  -> using AUTO fetched values for '%s'", column_name)
                if values:
                    data[column_name] = values
                    indexed_columns.append(column_name)
                    total_values += len(values)
                    logger.info("  -> %d values", len(values))
            except Exception as e:
                logger.warning("Failed to fetch values for column '%s': %s", column_name, e)

        # Store in memory
        ready_meta = IndexMeta(
            status=IndexStatus.READY,
            built_at=datetime.now(timezone.utc),
            columns=indexed_columns,
            value_count=total_values,
        )
        set_value_index_data(collection_name, data, ready_meta)

        # Populate ChromaDB (semantic index) - failure is non-fatal
        try:
            _populate_chroma_for_source(source, data, embedding_llm_id)
        except Exception as e:
            logger.warning("ChromaDB population failed for %s: %s", collection_name, e)

        logger.info(
            "Source index built in memory for %s: %d columns, %d values",
            source.to_path_parts(),
            len(indexed_columns),
            total_values,
        )
        return ready_meta

    except Exception as e:
        logger.error("Failed to build index for %s: %s", collection_name, e)
        error_meta = IndexMeta(status=IndexStatus.ERROR, error_message=str(e))
        set_value_index_data(collection_name, {}, error_meta)
        return error_meta


# =============================================================================
# Utility Functions
# =============================================================================


def get_resolvable_columns_for_entity(entity: dict) -> List[str]:
    """
    Get list of resolvable attribute names for an entity.

    TODO: Support indexing computed attributes (expression != column name)?
    """
    columns = []
    for attr in entity.get("attributes", []):
        if not attr.get("resolveInUserRequests", False):
            continue
        if not should_use_distinct_values(attr):
            logger.info(
                "Skipping %s.%s for value indexing (distinct values disabled)",
                entity.get("name", "?"),
                attr.get("name", "?"),
            )
            continue
        if not _is_simple_column(attr):
            logger.info(
                "Skipping computed attribute %s.%s (expression != name)", entity.get("name", "?"), attr.get("name", "?")
            )
            continue
        columns.append(attr.get("name"))
    return columns


def collect_sources_from_spec(
    semantic_model_id: str,
    version_id: str,
    spec_dict: dict,
) -> List[Tuple[SourceKey, List[str]]]:
    """
    Collect unique data sources and their resolvable columns from a spec.

    Args:
        spec_dict: SemanticModelSpec as dict

    Returns:
        List of (SourceKey, columns) tuples for indexing
    """
    from typing import Set

    source_columns: Dict[SourceKey, Set[str]] = {}

    for entity in spec_dict.get("entities", []):
        if entity.get("type") != "DATASET":
            continue
        resolvable_cols = get_resolvable_columns_for_entity(entity)
        if resolvable_cols:
            dataset_ref = entity.get("datasetRef") or ""
            if "." not in dataset_ref:
                continue
            source = SourceKey.from_entity(semantic_model_id, version_id, entity)
            if source not in source_columns:
                source_columns[source] = set()
            source_columns[source].update(resolvable_cols)

    return [(source, list(cols)) for source, cols in source_columns.items()]


def initialize_value_indices_from_sources(
    sources: List[Tuple[SourceKey, List[str]]],
    embedding_llm_id: Optional[str] = None,
    limit: int = 1000,
    spec_dict: Optional[dict] = None,
) -> Dict[str, bool]:
    from semantic_models.services.column_values import fetch_column_values

    results: Dict[str, bool] = {}
    seen: set[str] = set()

    for source, columns in sources:
        collection_name = source.to_collection_name()

        # Skip duplicates
        if collection_name in seen:
            continue
        seen.add(collection_name)

        # Check if already loaded with sufficient columns
        existing_meta = get_value_index_meta(collection_name)
        if existing_meta and existing_meta.status == IndexStatus.READY:
            # Compare requested columns vs stored columns
            existing_columns = set(existing_meta.columns)
            requested_columns = set(columns)
            if requested_columns.issubset(existing_columns):
                # All requested columns already indexed
                results[collection_name] = True
                continue
            # New columns needed - rebuild with union of columns
            new_columns = requested_columns - existing_columns
            logger.info(
                "Source %s needs additional columns %s, rebuilding with all columns", collection_name, new_columns
            )
            columns = list(existing_columns | requested_columns)

        # Build from data source
        manual_values_by_column: Dict[str, List[str]] = {}
        if spec_dict:
            entity = next((e for e in spec_dict.get("entities", []) if e.get("name") == source.entity_name), None)
            if entity:
                for attr in entity.get("attributes", []):
                    attr_name = attr.get("name")
                    if not attr_name:
                        continue
                    values = attr.get("manualValues") or []
                    if values:
                        manual_values_by_column[attr_name] = values

        try:
            meta = build_source_index_in_memory(
                source=source,
                columns=columns,
                fetch_values_fn=fetch_column_values,
                limit=limit,
                embedding_llm_id=embedding_llm_id,
                manual_values_by_column=manual_values_by_column,
            )
            results[collection_name] = meta.status == IndexStatus.READY
        except Exception as e:
            logger.error("Failed to build index for %s: %s", collection_name, e)
            results[collection_name] = False

    return results
