"""Instruction retrieval tools."""

from __future__ import annotations

import json
from typing import Iterable

from ims_mcp.clients.document import DocumentClient
from ims_mcp.constants import TAG_BOOTSTRAP
from ims_mcp.context import CallContext
from ims_mcp.services.bundler import Bundler
from ims_mcp.services.query_builder import QueryBuilder


def _unique_docs(docs: Iterable) -> list:
    seen: set[str] = set()
    out = []
    for doc in docs:
        if doc.id not in seen:
            out.append(doc)
            seen.add(doc.id)
    return out


def _extract_tags(doc) -> set[str]:
    meta = getattr(doc, "meta_fields", {}) or {}
    if isinstance(meta, str):
        try:
            meta = json.loads(meta)
        except Exception:
            meta = {}
    if not isinstance(meta, dict) and hasattr(meta, "__dict__"):
        meta = {k: v for k, v in vars(meta).items() if k != "rag"}
    if isinstance(meta, dict):
        tags = meta.get("tags", [])
    else:
        tags = getattr(meta, "tags", [])
    return {str(tag).lower() for tag in tags if str(tag).strip()}


def _filter_docs_by_any_tag(docs: list, tags: list[str] | None) -> list:
    required = {tag.strip().lower() for tag in (tags or []) if tag and tag.strip()}
    if not required:
        return docs
    return [doc for doc in docs if required & _extract_tags(doc)]


async def get_context_instructions(
    call_ctx: CallContext,
    document_client: DocumentClient,
    bundler: Bundler,
    query_builder: QueryBuilder,
    topic: str | None = None,
) -> str:
    # Compatibility wrapper: get-context semantics are query-instructions
    # with predefined bootstrap tag.
    return await query_instructions(
        call_ctx=call_ctx,
        document_client=document_client,
        bundler=bundler,
        query_builder=query_builder,
        query=None,
        tags=[TAG_BOOTSTRAP],
        topic=topic,
    )


async def query_instructions(
    call_ctx: CallContext,
    document_client: DocumentClient,
    bundler: Bundler,
    query_builder: QueryBuilder,
    query: str | None = None,
    tags: list[str] | None = None,
    topic: str | None = None,
) -> str:
    if not (query and query.strip()) and not tags:
        return "Error: at least one of query or tags is required"

    dataset_name = call_ctx.config.instruction_dataset
    try:
        dataset_id = call_ctx.dataset_lookup.get_id(dataset_name)
    except Exception as exc:
        return f"Error: failed to resolve instruction dataset '{dataset_name}': {exc}"
    if not dataset_id:
        return f"Error: instruction dataset not found: {dataset_name}"

    try:
        dataset = call_ctx.ragflow.get_dataset(name=dataset_name)
    except Exception as exc:
        return f"Error: failed to open instruction dataset '{dataset_name}': {exc}"

    if not dataset:
        return f"Error: instruction dataset not found: {dataset_name}"

    docs = []

    # Keyword search: tags via metadata_condition, query via keywords.
    if tags or (query and query.strip()):
        try:
            docs.extend(
                document_client.list_docs(
                    dataset=dataset,
                    page_size=1000,
                    **query_builder.build_list_params(tags=tags, query=query),
                )
            )
        except Exception as exc:
            return f"Error: failed to list instruction documents: {exc}"

    # Semantic expansion via retrieve when topic is provided.
    if topic and topic.strip():
        try:
            semantic = call_ctx.ragflow.retrieve(
                **query_builder.build_retrieve_params(dataset_ids=[dataset_id], query=topic.strip(), tags=tags)
            )
            for chunk in semantic:
                if getattr(chunk, "document_id", ""):
                    docs.extend(document_client.list_docs(dataset=dataset, doc_id=chunk.document_id, page_size=1))
        except Exception:
            # Semantic expansion should not block keyword results.
            pass

    docs = _filter_docs_by_any_tag(_unique_docs(docs), tags)
    if not docs:
        return "No instructions found"
    return bundler.bundle(docs, dataset_name)
