"""Hybrid skill search engine (BM25 + embedding + lexical boost).

Implements the search pipeline:
  Phase 1: BM25 rough-rank over all candidates
  Phase 2: Vector scoring (embedding cosine similarity)
  Phase 3: Hybrid score = vector_score + lexical_boost
  Phase 4: Deduplication + limit

Used by MCP ``search_skills`` and other search interfaces.
"""

from __future__ import annotations

import asyncio
import logging
import re
from typing import Any, Dict, List, Optional

logger = logging.getLogger("openspace.cloud")
CLOUD_EMBEDDING_SEARCH_MAX_LIMIT = 300


def _check_safety(text: str) -> list[str]:
    """Lazy wrapper — avoids importing skill_engine at module load time."""
    from openspace.skill_engine.skill_utils import check_skill_safety
    return check_skill_safety(text)


def _is_safe(flags: list[str]) -> bool:
    from openspace.skill_engine.skill_utils import is_skill_safe
    return is_skill_safe(flags)

_WORD_RE = re.compile(r"[a-z0-9]+")


def _tokenize(value: str) -> list[str]:
    return _WORD_RE.findall(value.lower()) if value else []


def _lexical_boost(query_tokens: list[str], name: str, slug: str) -> float:
    """Compute lexical boost score based on exact/prefix token matching."""
    slug_tokens = _tokenize(slug)
    name_tokens = _tokenize(name)
    boost = 0.0

    # Slug exact / prefix
    if slug_tokens and all(
        any(ct == qt for ct in slug_tokens) for qt in query_tokens
    ):
        boost += 1.4
    elif slug_tokens and all(
        any(ct.startswith(qt) for ct in slug_tokens) for qt in query_tokens
    ):
        boost += 0.8

    # Name exact / prefix
    if name_tokens and all(
        any(ct == qt for ct in name_tokens) for qt in query_tokens
    ):
        boost += 1.1
    elif name_tokens and all(
        any(ct.startswith(qt) for ct in name_tokens) for qt in query_tokens
    ):
        boost += 0.6

    return boost


class SkillSearchEngine:
    """Hybrid BM25 + embedding search engine for skills.

    Usage::

        engine = SkillSearchEngine()
        results = engine.search(
            query="weather forecast",
            candidates=candidates,
            query_embedding=[...],  # optional
            limit=20,
        )
    """

    def search(
        self,
        query: str,
        candidates: List[Dict[str, Any]],
        *,
        query_embedding: Optional[List[float]] = None,
        limit: int = 20,
    ) -> List[Dict[str, Any]]:
        """Run the full search pipeline on candidates.

        Each candidate dict should have at minimum:
          - ``skill_id``, ``name``, ``description``
          - ``_embedding`` (optional): pre-computed embedding vector
          - ``source``: "openspace-local" | "cloud"

        Args:
            query: Search query text.
            candidates: Candidate dicts to rank.
            query_embedding: Pre-computed query embedding (if available).
            limit: Max results to return.

        Returns:
            Sorted list of result dicts (highest score first).
        """
        q = query.strip()
        if not q or not candidates:
            return []

        query_tokens = _tokenize(q)
        if not query_tokens:
            return []

        # Phase 1: BM25 rough-rank
        filtered = self._bm25_phase(q, candidates, limit)

        # Phase 2+3: Vector + lexical scoring
        scored = self._score_phase(filtered, query_tokens, query_embedding)

        # Phase 4: Deduplicate and limit
        return self._dedup_and_limit(scored, limit)

    def _bm25_phase(
        self,
        query: str,
        candidates: List[Dict[str, Any]],
        limit: int,
    ) -> List[Dict[str, Any]]:
        """BM25 rough-rank to keep top candidates for embedding stage."""
        from openspace.skill_engine.skill_ranker import SkillRanker, SkillCandidate

        ranker = SkillRanker(enable_cache=True)
        bm25_candidates = [
            SkillCandidate(
                skill_id=c.get("skill_id", ""),
                name=c.get("name", ""),
                description=(
                    f"{c.get('description', '')}\n"
                    f"{c.get('category', '')}\n"
                    f"{c.get('local_category_path', '')}"
                ),
                body="",
                metadata=c,
            )
            for c in candidates
        ]
        ranked = ranker.bm25_only(query, bm25_candidates, top_k=min(limit * 3, len(candidates)))

        ranked_ids = {sc.skill_id for sc in ranked}
        filtered = [c for c in candidates if c.get("skill_id") in ranked_ids]

        # If BM25 found nothing, fall back to all candidates
        return filtered if filtered else candidates

    def _score_phase(
        self,
        candidates: List[Dict[str, Any]],
        query_tokens: list[str],
        query_embedding: Optional[List[float]],
    ) -> List[Dict[str, Any]]:
        """Compute hybrid score = vector_score + lexical_boost."""
        from openspace.cloud.embedding import cosine_similarity

        scored = []
        for candidate in candidates:
            candidate_name = candidate.get("name", "")
            candidate_slug = candidate.get("skill_id", candidate_name).split("__")[0].replace(":", "-")

            # Vector score. If client-side query embeddings are unavailable,
            # reuse the server-side cloud rank so cloud results keep semantic signal.
            vector_score: Optional[float] = None
            ranking_signal_score = 0.0
            if query_embedding:
                candidate_embedding = candidate.get("_embedding")
                if candidate_embedding and isinstance(candidate_embedding, list):
                    vector_score = cosine_similarity(query_embedding, candidate_embedding)
                    ranking_signal_score = vector_score
            elif isinstance(candidate.get("_search_rank"), (int, float)):
                ranking_signal_score = float(candidate["_search_rank"])

            # Lexical boost
            lexical_boost = _lexical_boost(query_tokens, candidate_name, candidate_slug)
            category_slug = str(candidate.get("local_category_path") or "").replace("/", "-")
            if category_slug:
                lexical_boost += 0.35 * _lexical_boost(query_tokens, "", category_slug)

            final_score = ranking_signal_score + lexical_boost

            result_entry: Dict[str, Any] = {
                "skill_id": candidate.get("skill_id", ""),
                "name": candidate_name,
                "description": candidate.get("description", ""),
                "source": candidate.get("source", ""),
                "score": round(final_score, 4),
            }
            if vector_score is not None and vector_score > 0:
                result_entry["vector_score"] = round(vector_score, 4)
            if isinstance(candidate.get("_search_rank"), (int, float)):
                result_entry["server_search_rank"] = round(float(candidate["_search_rank"]), 4)
            # Include optional fields
            for key in (
                "path",
                "visibility",
                "created_by",
                "origin",
                "tags",
                "quality",
                "safety_flags",
                "category",
                "local_category_path",
            ):
                if candidate.get(key):
                    result_entry[key] = candidate[key]
            scored.append(result_entry)

        scored.sort(key=lambda x: -x["score"])
        return scored

    @staticmethod
    def _dedup_and_limit(
        scored: List[Dict[str, Any]],
        limit: int,
    ) -> List[Dict[str, Any]]:
        """Deduplicate by name and apply limit."""
        seen: set[str] = set()
        deduped = []
        for item in scored:
            name = item["name"]
            if name in seen:
                continue
            seen.add(name)
            deduped.append(item)
        return deduped[:limit]


def build_local_candidates(
    skills: list,
    store: Any = None,
) -> List[Dict[str, Any]]:
    """Build search candidate dicts from SkillRegistry skills.

    Args:
        skills: List of ``SkillMeta`` from ``registry.list_skills()``.
        store: Optional ``SkillStore`` instance for quality data enrichment.

    Returns:
        List of candidate dicts ready for ``SkillSearchEngine.search()``.
    """
    from openspace.cloud.embedding import build_skill_embedding_text

    candidates: List[Dict[str, Any]] = []
    classification_store = None
    if store is not None and getattr(store, "db_path", None):
        try:
            from openspace.cloud.local_mapping import CloudLocalMappingStore

            classification_store = CloudLocalMappingStore(getattr(store, "db_path"))
        except Exception as e:
            logger.debug(f"Local classification lookup unavailable: {e}")
    for s in skills:
        # Read SKILL.md body
        readme_body = ""
        try:
            raw = s.path.read_text(encoding="utf-8")
            m = re.match(r"^---\n.*?\n---\n?", raw, re.DOTALL)
            readme_body = raw[m.end():].strip() if m else raw
        except Exception:
            pass

        local_category = ""
        local_category_path = ""
        if classification_store is not None:
            try:
                classification = classification_store.get_skill_local_classification(s.skill_id)
                if classification is not None:
                    local_category = classification.category
                    local_category_path = classification.local_category_path
            except Exception as e:
                logger.debug(f"Local classification lookup failed for {s.skill_id}: {e}")

        category_text = (
            f"\n\nLocal category: {local_category}\nLocal category path: {local_category_path}"
            if local_category or local_category_path
            else ""
        )
        embedding_text = build_skill_embedding_text(
            s.name,
            s.description,
            f"{readme_body}{category_text}",
        )

        # Safety check
        flags = _check_safety(embedding_text)
        if not _is_safe(flags):
            logger.info(f"BLOCKED local skill {s.skill_id} — {flags}")
            continue

        candidates.append({
            "skill_id": s.skill_id,
            "name": s.name,
            "description": s.description,
            "source": "openspace-local",
                "path": str(s.path),
                "is_local": True,
                "category": local_category,
                "local_category_path": local_category_path,
                "safety_flags": flags if flags else None,
                "_embedding_text": embedding_text,
            })
    if classification_store is not None:
        classification_store.close()

    # Enrich with quality data
    if store and candidates:
        try:
            all_records = store.load_all(active_only=True)
            for c in candidates:
                rec = all_records.get(c["skill_id"])
                if rec:
                    c["quality"] = {
                        "total_selections": rec.total_selections,
                        "completion_rate": round(rec.completion_rate, 3),
                        "effective_rate": round(rec.effective_rate, 3),
                    }
                    c["tags"] = rec.tags
                    if not c.get("category"):
                        c["category"] = rec.category.value
        except Exception as e:
            logger.warning(f"Quality lookup failed: {e}")

    return candidates


def build_cloud_candidates(
    cloud_items: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
    """Build search candidate dicts from cloud metadata/search items.

    Args:
        cloud_items: Items from cloud metadata or embedding search endpoints.

    Returns:
        List of candidate dicts (with safety filtering applied).
    """
    candidates: List[Dict[str, Any]] = []
    for item in cloud_items:
        cloud_skill_id = item.get("cloud_skill_id") or item.get("skill_id") or ""
        candidate_name = item.get("name") or item.get("title") or item.get("skill_name") or ""
        candidate_description = (
            item.get("description")
            or item.get("summary")
            or item.get("preview_text")
            or item.get("summary_line")
            or ""
        )
        candidate_tags = item.get("tags", [])
        if not isinstance(candidate_tags, list):
            candidate_tags = []
        safety_text = f"{candidate_name}\n{candidate_description}\n{' '.join(candidate_tags)}"
        flags = _check_safety(safety_text)
        if not _is_safe(flags):
            continue

        candidate_entry: Dict[str, Any] = {
            "skill_id": cloud_skill_id,
            "name": candidate_name,
            "description": candidate_description,
            "source": "cloud",
            "visibility": item.get("visibility") or item.get("effective_visibility") or "public",
            "is_local": False,
            "created_by": item.get("created_by", ""),
            "origin": item.get("origin", ""),
            "tags": candidate_tags,
            "safety_flags": flags if flags else None,
        }
        for key in (
            "cloud_skill_id",
            "package_id",
            "package_path",
            "source_api",
            "snippet",
            "skill_search_id",
            "match_mode",
            "served_mode",
            "semantic_status",
            "fallback_reason",
            "artifact_state",
            "downloadable",
            "metadata_only",
        ):
            if item.get(key):
                candidate_entry[key] = item[key]
        # Carry pre-computed embedding
        server_embedding = item.get("embedding")
        if server_embedding and isinstance(server_embedding, list):
            candidate_entry["_embedding"] = server_embedding
        server_search_rank = item.get("search_rank")
        if isinstance(server_search_rank, (int, float)):
            candidate_entry["_search_rank"] = float(server_search_rank)
        candidates.append(candidate_entry)

    return candidates


def build_cloud_results(
    cloud_search_items: List[Dict[str, Any]],
    *,
    limit: int,
) -> List[Dict[str, Any]]:
    """Map server-ranked cloud search rows to MCP search result shape."""
    results: List[Dict[str, Any]] = []
    seen_names: set[str] = set()

    for candidate in build_cloud_candidates(cloud_search_items):
        candidate_name = candidate.get("name", "")
        dedupe_name = candidate_name or candidate.get("skill_id", "")
        if dedupe_name in seen_names:
            continue
        seen_names.add(dedupe_name)

        entry: Dict[str, Any] = {
            "skill_id": candidate.get("skill_id", ""),
            "name": candidate_name,
            "description": candidate.get("description", ""),
            "source": "cloud",
            "score": round(float(candidate.get("_search_rank", 0.0)), 4),
        }
        if isinstance(candidate.get("_search_rank"), (int, float)):
            entry["server_search_rank"] = round(float(candidate["_search_rank"]), 4)
        for key in (
            "visibility",
            "created_by",
            "origin",
            "tags",
            "safety_flags",
            "cloud_skill_id",
            "package_id",
            "package_path",
            "source_api",
            "snippet",
            "skill_search_id",
            "match_mode",
            "served_mode",
            "semantic_status",
            "fallback_reason",
            "artifact_state",
            "downloadable",
            "metadata_only",
        ):
            if candidate.get(key):
                entry[key] = candidate[key]
        results.append(entry)
        if len(results) >= limit:
            break

    return results


async def hybrid_search_skills(
    query: str,
    local_skills: list = None,
    store: Any = None,
    source: str = "all",
    limit: int = 20,
) -> List[Dict[str, Any]]:
    """Shared cloud+local skill search with graceful fallback.

    Builds candidates, generates embeddings, runs ``SkillSearchEngine``.
    Cloud is attempted when *source* includes it; failures are silently
    skipped so the caller always gets local results at minimum.

    Args:
        query: Free-text search query.
        local_skills: ``SkillMeta`` list (from ``registry.list_skills()``).
        store: Optional ``SkillStore`` for quality enrichment.
        source: ``"all"`` | ``"local"`` | ``"cloud"``.
        limit: Maximum results.

    Returns:
        Ranked result dicts (same format as ``SkillSearchEngine.search()``).
    """
    from openspace.cloud.embedding import generate_embedding

    normalized_query = query.strip()
    if not normalized_query:
        return []

    candidates: List[Dict[str, Any]] = []

    if source in ("all", "local") and local_skills:
        candidates.extend(build_local_candidates(local_skills, store))

    if source in ("all", "cloud"):
        try:
            from openspace.cloud.client import OpenSpaceClient
            from openspace.cloud.config import load_cloud_config

            config = load_cloud_config()
            if not config.enabled:
                logger.info("hybrid_search_skills: cloud disabled by OPENSPACE_CLOUD_MODE")
            elif not config.api_key:
                logger.warning("hybrid_search_skills: OPENSPACE_CLOUD_API_KEY is required for cloud search")
            else:
                cloud_client = OpenSpaceClient(config)
                cloud_result_limit = (
                    limit if source == "cloud" else CLOUD_EMBEDDING_SEARCH_MAX_LIMIT
                )
                cloud_search_items = await asyncio.to_thread(
                    cloud_client.search_cloud_skills,
                    query=normalized_query,
                    limit=min(cloud_result_limit, 50),
                )
                if source == "cloud":
                    return build_cloud_results(cloud_search_items, limit=limit)
                candidates.extend(build_cloud_candidates(cloud_search_items))
        except Exception as e:
            logger.warning(f"hybrid_search_skills: cloud unavailable: {e}")
            if source == "cloud":
                raise

    if not candidates:
        return []

    # query embedding (optional — key/URL resolved inside generate_embedding)
    query_embedding: Optional[List[float]] = None
    try:
        query_embedding = await asyncio.to_thread(generate_embedding, normalized_query)
        if query_embedding:
            for candidate in candidates:
                if not candidate.get("_embedding") and candidate.get("_embedding_text"):
                    candidate_embedding = await asyncio.to_thread(
                        generate_embedding, candidate["_embedding_text"],
                    )
                    if candidate_embedding:
                        candidate["_embedding"] = candidate_embedding
    except Exception:
        pass

    engine = SkillSearchEngine()
    return engine.search(normalized_query, candidates, query_embedding=query_embedding, limit=limit)
