"""/search · /search/suggest — natural-language compiler v2 + FTS/trigram (+ embeddings when the gateway is reachable)."""
from __future__ import annotations

import asyncio
import time
from typing import Any

from fastapi import APIRouter, Depends, Query, Request
from sqlalchemy.exc import DBAPIError

from aiatlas.api.common import ENTITY_COLS, ENTITY_FROM, ApiError, cached, entity_summary, rate_limit
from aiatlas.db import connection, fetch_all, fetch_val
from aiatlas.services import cache
from aiatlas.services.embeddings import embed_query
from aiatlas.services.llm import gateway
from aiatlas.services.search import Query as SearchQuery
from aiatlas.services.search import compile_query, search_entities, suggest, verify_organization, where_for

router = APIRouter(prefix="/api/v1/search", tags=["search"])
TOTAL_CAP = 10_000
EMBED_TIMEOUT_S = 2.0
EMBED_BACKOFF_S = 300
EMBED_DEGRADED_KEY = "search:embed-degraded"


def _count_sql(q: SearchQuery) -> tuple[str, dict[str, Any]]:
    """Same filters as `search_entities`, without ranking — capped estimate. All casts are regex-guarded in `where_for`."""
    import re

    where, params = where_for(q)
    text = (q.residual or "").strip() or ("" if q.has_structure else q.text)
    if q.benchmark and q.entity_type == "benchmark":
        text = q.benchmark
    if q.provider and q.entity_type == "provider":
        text = q.provider
    if text:
        params["q"] = text
        params["qlike"] = f"%{text}%"
        params["qprefix"] = " & ".join(f"{w}:*" for w in re.findall(r"\w+", text)[:8]) or text
        where.append("(e.search @@ to_tsquery('simple', :qprefix) or e.canonical_name ilike :qlike or e.canonical_name % :q "
                     "or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :qlike))")
    return f"select count(*) from (select 1 from entities e where {' and '.join(where)} limit {TOTAL_CAP}) t", params


@router.get("", dependencies=[Depends(rate_limit("search"))])
@cached(120)
async def search(request: Request, q: str = Query("", max_length=300), type: str | None = Query(None, alias="type"),
                 limit: int = Query(30, ge=1, le=200), offset: int = Query(0, ge=0)) -> dict[str, Any]:
    compiled = compile_query(q)
    if type:
        compiled.entity_type = type
    if not q.strip() and not type:
        raise ApiError(400, "q is required")
    async with connection() as conn:
        org_ok = None
        if compiled.organization:
            org_ok = await verify_organization(conn, compiled.organization)
            if org_ok is None:
                # not an organization we know: demote to free text and say so
                name = compiled.organization
                compiled.organization = None
                compiled.compiled = [c for c in compiled.compiled if c["filter"] != "organization"]
                compiled.residual = " ".join(x for x in [compiled.residual, name] if x)
                compiled.filters["residual"] = compiled.residual
                compiled.unrecognised = sorted(set(compiled.unrecognised) | {name})
            else:
                compiled.organization = org_ok["canonical_name"]
                for c in compiled.compiled:
                    if c["filter"] == "organization":
                        c["value"] = {"id": org_ok["id"], "slug": org_ok["slug"], "name": org_ok["canonical_name"]}
                        c["label"] = f"Organization: {org_ok['canonical_name']}"
        embedding = None
        residual = (compiled.residual or "").strip()
        if gateway.available and residual and not await cache.cache_get(EMBED_DEGRADED_KEY):
            try:
                embedding = await asyncio.wait_for(embed_query(residual), timeout=EMBED_TIMEOUT_S)
            except Exception:  # noqa: BLE001 — never wait on the LLM
                embedding = None
            if embedding is None:  # back off: FTS-only for a while instead of paying the timeout on every query
                await cache.cache_set(EMBED_DEGRADED_KEY, {"since": time.time()}, EMBED_BACKOFF_S)
        if embedding is not None:
            # the vector join needs `entity_embeddings` (pgvector); when the table is missing the statement aborts the connection's
            # transaction, so the FTS-only fallback runs on a fresh pooled connection and embeddings are backed off for a while
            try:
                async with connection() as vconn:
                    hits = await search_entities(vconn, compiled, limit=limit, offset=offset, embedding=embedding)
            except DBAPIError:
                embedding = None
                await cache.cache_set(EMBED_DEGRADED_KEY, {"since": time.time(), "reason": "entity_embeddings unavailable"}, EMBED_BACKOFF_S)
        if embedding is None:
            hits = await search_entities(conn, compiled, limit=limit, offset=offset, embedding=None)
        ids = [h["id"] for h in hits]
        full = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where e.id = any(cast(:ids as text[]))", ids=ids) if ids else []
        count_sql, count_params = _count_sql(compiled)
        total = await fetch_val(conn, count_sql, **count_params)
    by_id = {r["id"]: r for r in full}
    items = []
    for h in hits:
        row = by_id.get(h["id"])
        s = entity_summary(row) if row else None
        if s:
            s["rank"] = float(h["rank"] or 0)
            items.append(s)
    query = {**compiled.as_dict(), "semantic": embedding is not None, "version": 2}
    if compiled.memory_gb is not None:
        query["note"] = "memory_gb is converted to an ESTIMATED parameter bound (4-bit weights, 8K context, 2 GB reserved) — see /methodology hardware_fit"
    return {"query": query, "items": items, "total": int(total or 0), "limit": limit, "offset": offset}


@router.get("/suggest", dependencies=[Depends(rate_limit("search"))])
@cached(120)
async def search_suggest(request: Request, q: str = Query("", max_length=120), limit: int = Query(8, ge=1, le=20)) -> dict[str, Any]:
    prefix = q.strip()
    if len(prefix) < 1:
        return {"items": []}
    async with connection() as conn:
        rows = await suggest(conn, prefix, limit=limit)
    return {"items": [{"id": r["id"], "entity_type": r["entity_type"], "slug": r["slug"], "name": r["canonical_name"], "organization_name": r["organization_name"]} for r in rows]}