"""Search service — business logic layer over OpenSearch client.

Provides graph_id-scoped search and retrieval, mapping OpenSearch
responses to Pydantic models. Default mode is BM25-only (fast);
hybrid BM25 + KNN is opt-in via semantic=True on SearchRequest.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from robosystems.config import env
from robosystems.logger import logger
from robosystems.models.api.search import (
  DocumentListItem,
  DocumentListResponse,
  DocumentSection,
  DocumentUploadRequest,
  DocumentUploadResponse,
  SearchHit,
  SearchRequest,
  SearchResponse,
)

from .client import OpenSearchClient

if TYPE_CHECKING:
  from .embeddings import EmbeddingService

# Snippet fallback length when no highlights available
SNIPPET_FALLBACK_LENGTH = 300


class SearchService:
  """Search service scoped by graph_id. BM25 default, hybrid opt-in."""

  def __init__(self, client: OpenSearchClient) -> None:
    self.client = client
    self._embedding_service: EmbeddingService | None = None

  @property
  def embedding_service(self) -> EmbeddingService:
    """Lazy-load the platform embedding service."""
    if self._embedding_service is None:
      from .embeddings import get_embedding_service

      self._embedding_service = get_embedding_service()
      logger.info("Loaded EmbeddingService for search")
    return self._embedding_service

  def search_documents(self, graph_id: str, request: SearchRequest) -> SearchResponse:
    """Search documents with graph_id isolation.

    Default mode is BM25-only (fast keyword search). When request.semantic
    is True, uses hybrid BM25 + KNN search which adds vector similarity
    scoring at the cost of higher latency on large corpora.
    """
    filters: dict[str, Any] = {}
    if request.entity:
      filters["entity"] = request.entity
    if request.form_type:
      filters["form_type"] = request.form_type
    if request.section:
      filters["section"] = request.section
    if request.element:
      filters["element"] = request.element
    if request.source_type:
      filters["source_type"] = request.source_type
    if request.fiscal_year:
      filters["fiscal_year"] = request.fiscal_year
    if request.date_from:
      filters["date_from"] = request.date_from
    if request.date_to:
      filters["date_to"] = request.date_to

    if request.semantic:
      query_embedding = self.embedding_service.embed_single(request.query)
      result = self.client.search_hybrid(
        query=request.query,
        query_embedding=query_embedding,
        graph_id=graph_id,
        filters=filters if filters else None,
        size=request.size,
        offset=request.offset,
      )
    else:
      result = self.client.search(
        query=request.query,
        graph_id=graph_id,
        filters=filters if filters else None,
        size=request.size,
        offset=request.offset,
      )

    hits = []
    for hit in result.get("hits", {}).get("hits", []):
      source = hit.get("_source", {})
      highlight = hit.get("highlight", {})

      # Build snippet from highlights or fallback
      content_fragments = highlight.get("content", [])
      if content_fragments:
        snippet = " ... ".join(content_fragments)
      else:
        # No highlight — we excluded content from _source in search,
        # so use section_label as minimal context
        snippet = source.get("section_label", "")

      # Extract PG document ID for user docs: "udoc_{pg_doc_id}_{idx}" → pg_doc_id
      doc_id = hit.get("_id", source.get("document_id", ""))
      parent_doc_id = None
      if doc_id.startswith("udoc_"):
        parent_doc_id = doc_id.rsplit("_", 1)[0].removeprefix("udoc_") or None

      hits.append(
        SearchHit(
          document_id=doc_id,
          score=hit.get("_score", 0.0),
          source_type=source.get("source_type", ""),
          parent_document_id=parent_doc_id,
          entity_ticker=source.get("entity_ticker"),
          entity_name=source.get("entity_name"),
          section_label=source.get("section_label"),
          section_id=source.get("section_id"),
          element_qname=source.get("element_qname"),
          xbrl_elements=source.get("xbrl_elements"),
          filing_date=source.get("filing_date"),
          fiscal_year=source.get("fiscal_year") or None,
          form_type=source.get("form_type"),
          snippet=snippet,
          content_length=source.get("content_length", 0),
          content_url=source.get("content_url"),
          document_title=source.get("document_title"),
          tags=source.get("tags"),
          folder=source.get("folder"),
        )
      )

    total = result.get("hits", {}).get("total", {})
    total_count = total.get("value", 0) if isinstance(total, dict) else total

    return SearchResponse(
      total=total_count,
      hits=hits,
      query=request.query,
      graph_id=graph_id,
    )

  def get_document_section(
    self, graph_id: str, document_id: str
  ) -> DocumentSection | None:
    """Retrieve full document section by ID with graph_id verification."""
    doc = self.client.get_document(document_id, graph_id)
    if doc is None:
      return None

    return DocumentSection(
      document_id=document_id,
      graph_id=doc.get("graph_id", graph_id),
      source_type=doc.get("source_type", ""),
      entity_ticker=doc.get("entity_ticker"),
      entity_name=doc.get("entity_name"),
      entity_cik=doc.get("entity_cik"),
      section_label=doc.get("section_label"),
      section_id=doc.get("section_id"),
      element_qname=doc.get("element_qname"),
      filing_date=doc.get("filing_date"),
      fiscal_year=doc.get("fiscal_year") or None,
      fiscal_period=doc.get("fiscal_period"),
      form_type=doc.get("form_type"),
      accession_number=doc.get("accession_number"),
      content=doc.get("content", ""),
      content_url=doc.get("content_url"),
      content_length=doc.get("content_length", 0),
      document_title=doc.get("document_title"),
      tags=doc.get("tags"),
      folder=doc.get("folder"),
    )

  def upload_document(
    self,
    graph_id: str,
    request: DocumentUploadRequest,
  ) -> DocumentUploadResponse:
    """Upload a markdown document: parse, section, embed, and index.

    Args:
        graph_id: Target graph for tenant isolation.
        request: Document upload request with title, content, etc.

    Returns:
        DocumentUploadResponse with section IDs and counts.
    """
    from .embeddings import EMBEDDING_MODEL_ID
    from .markdown_parser import content_hash, parse_document

    # Parse markdown
    metadata, sections = parse_document(request.content, request.title)
    if not sections:
      raise ValueError("Document produced no indexable sections")

    # Generate base document ID
    # User docs use "udoc_" prefix so the PG document ID is trivially extractable
    # from search results: "udoc_{pg_doc_id}_{section_idx}" → rsplit("_", 1)[0][5:]
    # SEC/pipeline docs use "doc_" prefix (no PG document table).
    if request.external_id:
      base_id = f"udoc_{request.external_id}"
    else:
      base_id = f"udoc_{content_hash(request.content)}"

    # Remove old sections if re-uploading
    self.client.delete_by_document_prefix(graph_id, base_id)

    # Use frontmatter values, falling back to request fields
    doc_title = metadata.get("title", request.title)
    doc_tags = metadata.get("tags", request.tags)
    doc_folder = metadata.get("folder", request.folder)

    # Embed all section contents
    section_texts = [s["content"] for s in sections]
    embeddings = self.embedding_service.embed_batch(section_texts)

    # Build OpenSearch documents
    os_docs: list[dict[str, Any]] = []
    section_ids: list[str] = []
    total_length = 0

    for idx, (section, embedding) in enumerate(zip(sections, embeddings, strict=True)):
      doc_id = f"{base_id}_{idx}"
      section_ids.append(doc_id)
      total_length += len(section["content"])

      os_docs.append(
        {
          "graph_id": graph_id,
          "document_id": doc_id,
          "source_type": "uploaded_doc",
          "document_title": doc_title,
          "section_id": section["section_id"],
          "section_label": section["section_label"],
          "content": section["content"],
          "content_length": len(section["content"]),
          "tags": doc_tags,
          "folder": doc_folder,
          "embedding": embedding,
          "embedding_model": EMBEDDING_MODEL_ID,
        }
      )

    self.client.bulk_index(os_docs)

    return DocumentUploadResponse(
      id="",
      document_id=base_id,
      sections_indexed=len(sections),
      total_content_length=total_length,
      section_ids=section_ids,
    )

  def delete_document(self, graph_id: str, document_id: str) -> bool:
    """Delete a document and all its sections."""
    # Try exact match first
    if self.client.delete_document(document_id, graph_id):
      return True
    # Try prefix match (delete all sections of a multi-section doc)
    deleted = self.client.delete_by_document_prefix(graph_id, document_id)
    return deleted > 0

  def list_documents(
    self, graph_id: str, source_type: str | None = None
  ) -> DocumentListResponse:
    """List indexed documents grouped by title."""
    result = self.client.list_documents(graph_id, source_type)
    aggs = result.get("aggregations", {})
    doc_buckets = aggs.get("documents", {}).get("buckets", [])

    documents: list[DocumentListItem] = []
    for bucket in doc_buckets:
      # Extract nested aggregation values
      source_types = bucket.get("source_type", {}).get("buckets", [])
      st = source_types[0]["key"] if source_types else "uploaded_doc"
      folders = bucket.get("folder", {}).get("buckets", [])
      folder = folders[0]["key"] if folders else None
      tags_buckets = bucket.get("tags", {}).get("buckets", [])
      tags = [t["key"] for t in tags_buckets] if tags_buckets else None
      last_indexed = bucket.get("last_indexed", {}).get("value_as_string")

      documents.append(
        DocumentListItem(
          document_title=bucket["key"],
          section_count=bucket["doc_count"],
          source_type=st,
          folder=folder,
          tags=tags,
          last_indexed=last_indexed,
        )
      )

    return DocumentListResponse(
      total=len(documents),
      documents=documents,
      graph_id=graph_id,
    )

  def index_documents(
    self, graph_id: str, documents: list[dict[str, Any]]
  ) -> dict[str, int]:
    """Index documents for a graph_id (used by SEC pipeline)."""
    for doc in documents:
      doc["graph_id"] = graph_id
    return self.client.bulk_index(documents)

  def health(self) -> dict[str, Any]:
    """Check search service health."""
    return self.client.health()


# Lazy singleton
_service: SearchService | None = None


def get_search_service() -> SearchService | None:
  """Get the search service singleton. Returns None if SEMANTIC_SEARCH_ENABLED is false."""
  global _service
  if _service is not None:
    return _service

  if not env.SEMANTIC_SEARCH_ENABLED:
    return None

  try:
    client = OpenSearchClient(env.OPENSEARCH_URL, env.OPENSEARCH_INDEX)
    client.create_index_if_not_exists()
    _service = SearchService(client)
    logger.info("Search service initialized")
    return _service
  except Exception as e:
    logger.warning(f"Failed to initialize search service: {e}")
    return None
