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

import asyncio
import re
import time
from typing import Any

from fastapi import APIRouter, Depends, Query, Request

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

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."""
    where = ["e.merged_into is null"]
    params: dict[str, Any] = {}
    if q.entity_type:
        where.append("e.entity_type = :etype")
        params["etype"] = q.entity_type
    if q.openness == "open":
        where.append("(e.attributes->>'openness' in ('open-weights','open-source','open') or e.attributes->>'weights_availability' = 'open')")
    elif q.openness == "proprietary":
        where.append("e.attributes->>'openness' in ('proprietary','closed')")
    if q.year_from:
        where.append("left(e.attributes->>'release_date', 4) >= :yf")
        params["yf"] = str(q.year_from)
    if q.year_to:
        where.append("left(e.attributes->>'release_date', 4) <= :yt")
        params["yt"] = str(q.year_to)
    if q.params_min:
        where.append("(e.attributes->>'parameter_count')::double precision >= :pmin")
        params["pmin"] = float(q.params_min)
    if q.params_max:
        where.append("(e.attributes->>'parameter_count')::double precision <= :pmax")
        params["pmax"] = float(q.params_max)
    if q.context_min:
        where.append("(e.attributes->>'context_length')::double precision >= :cmin")
        params["cmin"] = float(q.context_min)
    for i, mod in enumerate(q.modalities):
        where.append(f"e.attributes->'modalities' ? :mod{i}")
        params[f"mod{i}"] = mod
    if q.organization:
        where.append("exists (select 1 from entities o where o.id = e.organization_id and o.canonical_name ilike :org)")
        params["org"] = f"%{q.organization}%"
    text = (q.filters.get("residual") or "").strip() or ("" if any([q.entity_type, q.openness, q.year_from, q.params_min, q.context_min, q.modalities]) else q.text)
    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")
    embedding = None
    residual = (compiled.filters.get("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)
    async with connection() as conn:
        hits = await search_entities(conn, compiled, limit=limit, offset=offset, embedding=embedding)
        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)
    return {"query": {**compiled.as_dict(), "semantic": embedding is not None}, "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]}