"""Search: Postgres FTS + trigram (+ pgvector when embeddings exist) and a deterministic natural-language → filter compiler (v2).

    "open models over 100B released this year"      → type model · openness open · params ≥ 1e11 · year_from = year_to = <this year>
    "cheapest models with 1M context"               → type model · context ≥ 1 000 000 · sort cheapest
    "reasoning models under $1/M tokens"            → type model · reasoning · max_output_price 1.0
    "models that fit in 64GB"                       → type model · memory_gb 64 (→ ESTIMATED params bound at 4-bit)
    "Anthropic models released since 2025"          → type model · organization Anthropic (verified by the API against org aliases) · year_from 2025
    "open vision models with Apache license"        → type model · openness open · modality image · license Apache-2.0
    "papers introducing MoE models"                 → type paper · residual "introducing MoE"

Every recognised fragment is listed in `compiled` (filter, label, value, source_span) so the UI can show "Compiled as …"; the rest of the
text is `residual` (used for full-text search) and its words are echoed in `unrecognised`."""
from __future__ import annotations

import re
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any

from sqlalchemy.ext.asyncio import AsyncConnection

from aiatlas.db import fetch_all
from aiatlas.ontology.licenses import normalize_license
from aiatlas.sdk.extract.numbers import parse_context_length, parse_param_count
from aiatlas.services import hardware_fit as hf

TYPE_WORDS = {
    "model": ("model", "models", "llm", "llms", "language model", "language models"),
    "company": ("company", "companies", "startup", "startups", "lab", "labs", "organization", "organizations", "organisation"),
    "paper": ("paper", "papers", "research", "publication", "publications", "preprint", "preprints", "arxiv"),
    "provider": ("provider", "providers", "inference provider", "api provider", "inference providers"),
    "benchmark": ("benchmark", "benchmarks", "leaderboard", "leaderboards", "eval", "evals"),
    "hardware": ("gpu", "gpus", "chip", "chips", "hardware", "accelerator", "accelerators", "npu", "tpu"),
    "framework": ("framework", "frameworks", "library", "libraries", "runtime", "runtimes"),
    "dataset": ("dataset", "datasets"),
    "tool": ("tool", "tools", "agent", "agents", "app", "application"),
    "repository": ("repo", "repos", "repository", "repositories"),
}
MODALITY_WORDS = {"vision": "image", "multimodal": "image", "image": "image", "images": "image", "audio": "audio", "speech": "audio", "voice": "audio",
                  "video": "video", "embedding": "embedding", "embeddings": "embedding"}
LICENSE_WORDS = {"apache": "Apache-2.0", "apache 2.0": "Apache-2.0", "apache-2.0": "Apache-2.0", "mit": "MIT", "mit license": "MIT", "bsd": "BSD-3-Clause",
                 "cc-by": "CC-BY-4.0", "cc by": "CC-BY-4.0", "gpl": "GPL-3.0", "llama license": "Llama-3.1-Community", "gemma license": "Gemma-Terms"}
SORT_WORDS = {"cheapest": "cheapest", "cheap": "cheapest", "lowest price": "cheapest", "newest": "newest", "latest": "newest", "most recent": "newest",
              "recent": "newest", "largest": "largest", "biggest": "largest", "smallest": "smallest", "best": "best"}
STOP = {"the", "a", "an", "with", "and", "or", "of", "for", "that", "which", "in", "on", "to", "by", "from", "at", "my", "me", "i", "are", "is", "show", "find", "list", "all", "any"}
_SIZE = r"(\d+(?:\.\d+)?)\s*([bBmMtT])\b(?:\s*param(?:eter)?s?)?"


@dataclass
class Query:
    text: str = ""
    entity_type: str | None = None
    openness: str | None = None
    year_from: int | None = None
    year_to: int | None = None
    params_min: int | None = None
    params_max: int | None = None
    context_min: int | None = None
    modalities: list[str] = field(default_factory=list)
    organization: str | None = None
    filters: dict[str, Any] = field(default_factory=dict)
    # v2
    reasoning: bool | None = None
    max_input_price: float | None = None
    max_output_price: float | None = None
    memory_gb: float | None = None
    license_key: str | None = None
    commercial_use: bool | None = None
    days_back: int | None = None
    benchmark: str | None = None
    provider: str | None = None
    sort: str | None = None
    compiled: list[dict[str, Any]] = field(default_factory=list)
    residual: str = ""
    unrecognised: list[str] = field(default_factory=list)

    def as_dict(self) -> dict[str, Any]:
        out = {k: v for k, v in self.__dict__.items() if v not in (None, "", [], {}) and k not in ("compiled", "residual", "unrecognised")}
        out["compiled"] = self.compiled
        out["residual"] = self.residual
        out["unrecognised"] = self.unrecognised
        out["filters"] = {**self.filters, "residual": self.residual}
        return out

    @property
    def has_structure(self) -> bool:
        return any([self.entity_type, self.openness, self.year_from, self.year_to, self.params_min, self.params_max, self.context_min, self.modalities,
                    self.organization, self.reasoning, self.max_input_price, self.max_output_price, self.memory_gb, self.license_key, self.commercial_use,
                    self.days_back, self.benchmark, self.provider])


class _Compiler:
    def __init__(self, text: str):
        self.text = text
        self.low = text.lower()
        self.consumed: list[tuple[int, int]] = []
        self.q = Query(text=text)

    def take(self, m: re.Match[str], filter_: str, label: str, value: Any) -> None:
        self.consumed.append((m.start(), m.end()))
        self.q.compiled.append({"filter": filter_, "label": label, "value": value, "source_span": self.text[m.start():m.end()]})

    def search(self, pattern: str) -> re.Match[str] | None:
        for m in re.finditer(pattern, self.low, flags=re.IGNORECASE):
            if not any(s <= m.start() < e or s < m.end() <= e for s, e in self.consumed):
                return m
        return None

    def residual(self) -> str:
        chars = list(self.text)
        for s, e in self.consumed:
            for i in range(s, e):
                chars[i] = " "
        return " ".join("".join(chars).split())


def compile_query(q: str) -> Query:
    c = _Compiler(q.strip())
    out = c.q
    year_now = datetime.now(UTC).year

    # --- explicit "benchmark <name>" / "provider <name>"
    m = c.search(r"\bbenchmark\s+([a-z0-9][\w.\- ]{1,40}?)(?=\s+(?:leaderboard|results|scores|models)\b|$)")
    if m:
        out.benchmark = m.group(1).strip()
        out.entity_type = "benchmark"
        c.take(m, "benchmark", "Benchmark", out.benchmark)
    m = c.search(r"\bprovider\s+([a-z0-9][\w.\- ]{1,40}?)(?=\s+(?:prices?|models)\b|$)")
    if m:
        out.provider = m.group(1).strip()
        out.entity_type = "provider"
        c.take(m, "provider", "Provider", out.provider)

    # --- price bounds
    m = c.search(r"(?:under|below|cheaper than|less than|at most|<|≤)\s*\$?\s*(\d+(?:\.\d+)?)\s*(?:usd)?\s*(?:/|per)\s*(?:1\s*)?m(?:illion)?(?:\s*(?:tok(?:ens?)?|output(?:\s*tokens?)?|input(?:\s*tokens?)?))?")
    if m:
        which = "input" if "input" in m.group(0) else "output"
        val = float(m.group(1))
        if which == "input":
            out.max_input_price = val
        else:
            out.max_output_price = val
        c.take(m, f"max_{which}_price", f"Cheapest {which} price ≤ ${val:g} / 1M tokens", val)
    m = c.search(r"\$\s*(\d+(?:\.\d+)?)\s*(?:/|per)\s*(?:1\s*)?m(?:illion)?\b(?:\s*tok(?:ens?)?)?")
    if m and out.max_output_price is None and out.max_input_price is None:
        out.max_output_price = float(m.group(1))
        c.take(m, "max_output_price", f"Cheapest output price ≤ ${out.max_output_price:g} / 1M tokens", out.max_output_price)

    # --- memory / hardware fit
    m = c.search(r"(?:fits?\s+(?:in|on)|runs?\s+on|run\s+locally\s+on|on\s+my)\s*(?:a\s+|an\s+)?(?:mac(?:book)?(?:\s+\w+)?\s+with\s+)?(\d+(?:\.\d+)?)\s*gb\b(?:\s*(?:of\s+)?(?:ram|memory|vram|unified memory))?")
    if not m:
        m = c.search(r"\b(\d+(?:\.\d+)?)\s*gb\s*(?:of\s+)?(?:ram|memory|vram)\b")
    if m:
        out.memory_gb = float(m.group(1))
        c.take(m, "memory_gb", f"Fits in {out.memory_gb:g} GB (estimated at 4-bit, 8K context)", out.memory_gb)
    elif c.search(r"\bon my mac\b"):
        m2 = c.search(r"\bon my mac\b")
        assert m2 is not None
        out.openness = "open"
        c.take(m2, "openness", "Runs locally → open weights", "open")

    # --- context window
    m = c.search(r"(?:context(?:\s+window)?\s*(?:of\s+)?(?:>=|≥|>|over|above|at least|min(?:imum)?)?\s*(\d+(?:\.\d+)?\s*[kKmM])\b(?:\s*tokens?)?)")
    if not m:
        m = c.search(r"(?:with\s+)?(\d+(?:\.\d+)?\s*[kKmM])\b\s*(?:\+\s*)?(?:tokens?\s+)?(?:of\s+)?context(?:\s+window)?")
    if not m:
        m = c.search(r"\blong[\s-]context\b")
        if m:
            out.context_min = 128_000
            c.take(m, "context_min", "Context ≥ 128K tokens", 128_000)
            m = None
    if m:
        out.context_min = parse_context_length(m.group(1))
        c.take(m, "context_min", f"Context ≥ {out.context_min:,} tokens", out.context_min)

    # --- parameter counts
    m = c.search(r"\bbetween\s+" + _SIZE + r"\s+and\s+" + _SIZE)
    if m:
        out.params_min = parse_param_count(m.group(1) + m.group(2) + " params")
        out.params_max = parse_param_count(m.group(3) + m.group(4) + " params")
        c.take(m, "params_range", f"Parameters between {m.group(1)}{m.group(2).upper()} and {m.group(3)}{m.group(4).upper()}", [out.params_min, out.params_max])
    m = c.search(r"(?:more than|over|above|at least|>=|≥|>|larger than|bigger than)\s*" + _SIZE)
    if m:
        out.params_min = parse_param_count(m.group(1) + m.group(2) + " params")
        c.take(m, "params_min", f"Parameters ≥ {m.group(1)}{m.group(2).upper()}", out.params_min)
    m = c.search(r"(?:less than|under|below|at most|<=|≤|<|smaller than)\s*" + _SIZE)
    if m:
        out.params_max = parse_param_count(m.group(1) + m.group(2) + " params")
        c.take(m, "params_max", f"Parameters ≤ {m.group(1)}{m.group(2).upper()}", out.params_max)
    m = c.search(r"\b" + _SIZE + r"\s*(?:models?|parameters?)\b")
    if m and out.params_min is None and out.params_max is None:
        n = parse_param_count(m.group(1) + m.group(2) + " params")
        if n:
            out.params_min, out.params_max = int(n * 0.85), int(n * 1.15)
            c.take(m, "params_about", f"About {m.group(1)}{m.group(2).upper()} parameters (±15%)", n)

    # --- dates
    m = c.search(r"\b(?:(?:released|launched|published|announced)\s+)?(?:in\s+the\s+)?(?:last|past)\s+(\d+)\s+(day|week|month|year)s?\b")
    if m:
        n, unit = int(m.group(1)), m.group(2)
        out.days_back = n * {"day": 1, "week": 7, "month": 30, "year": 365}[unit]
        c.take(m, "days_back", f"Released in the last {n} {unit}{'s' if n > 1 else ''}", out.days_back)
    m = c.search(r"\b(?:released|launched|published|announced)?\s*(?:this|the current)\s+year\b")
    if m:
        out.year_from = out.year_to = year_now
        c.take(m, "year", f"Released in {year_now}", year_now)
    m = c.search(r"\b(?:released|launched|published|announced)?\s*(?:last|previous)\s+year\b")
    if m:
        out.year_from = out.year_to = year_now - 1
        c.take(m, "year", f"Released in {year_now - 1}", year_now - 1)
    m = c.search(r"\b(?:(?:released|launched|published|announced)\s+)?(?:since|after|from)\s+(20\d\d)\b")
    if m:
        out.year_from = int(m.group(1))
        c.take(m, "year_from", f"Released since {out.year_from}", out.year_from)
    m = c.search(r"\b(?:(?:released|launched|published|announced)\s+)?(?:before|until|up to)\s+(20\d\d)\b")
    if m:
        out.year_to = int(m.group(1))
        c.take(m, "year_to", f"Released before {out.year_to}", out.year_to)
    m = c.search(r"\b(?:(?:released|launched|published|announced)\s+)?in\s+(20\d\d)\b")
    if m:
        out.year_from = out.year_to = int(m.group(1))
        c.take(m, "year", f"Released in {m.group(1)}", int(m.group(1)))
    m = c.search(r"\b(20\d\d)\b")
    if m and out.year_from is None and out.year_to is None:
        out.year_from = out.year_to = int(m.group(1))
        c.take(m, "year", f"Released in {m.group(1)}", int(m.group(1)))

    # --- licence
    m = c.search(r"\b(apache(?:[\s-]*2(?:\.0)?)?|mit|bsd|cc[\s-]by|gpl|llama|gemma)\s*(?:licen[cs]e[d]?)\b")
    if m:
        raw = m.group(1).strip()
        key = normalize_license(raw) or LICENSE_WORDS.get(raw.lower())
        if key:
            out.license_key = key
            c.take(m, "license", f"Licence {key}", key)
    m = c.search(r"\b(?:for\s+)?commercial(?:\s+use|ly usable|-use)?\b(?:\s+allowed|\s+ok)?")
    if m:
        out.commercial_use = True
        c.take(m, "commercial_use", "Licence allows commercial use", True)

    # --- openness
    m = c.search(r"\b(open[- ]?(?:weights?|source|sourced)|open)\b(?=\s+(?:models?|llms?|vision|reasoning|coding|multimodal|\w+\s+models?)\b)") or c.search(r"\bopen[- ]?(?:weights?|source|sourced)\b")
    if m:
        out.openness = "open"
        c.take(m, "openness", "Open weights / open source", "open")
    else:
        m = c.search(r"\b(proprietary|closed(?:[- ]source)?|api[- ]only)\b")
        if m:
            out.openness = "proprietary"
            c.take(m, "openness", "Proprietary (API only)", "proprietary")

    # --- reasoning
    m = c.search(r"\b(reasoning|thinking)\b(?=\s+(?:models?|llms?)\b)") or c.search(r"\bwith\s+(reasoning|thinking)\b")
    if m:
        out.reasoning = True
        c.take(m, "reasoning", "Reasoning / thinking models", True)

    # --- modalities
    for w, mod in MODALITY_WORDS.items():
        m = c.search(rf"\b{w}\b")
        if m and mod not in out.modalities:
            out.modalities.append(mod)
            c.take(m, "modality", f"Modality: {mod}", mod)

    # --- entity type: the earliest surviving type word wins ("papers introducing MoE models" → paper)
    if out.entity_type is None:
        best: tuple[int, str, re.Match[str]] | None = None
        for etype, words in TYPE_WORDS.items():
            for w in sorted(words, key=len, reverse=True):
                m = c.search(rf"\b{re.escape(w)}\b")
                if m and (best is None or m.start() < best[0]):
                    best = (m.start(), etype, m)
                    break
        if best:
            out.entity_type = best[1]
            c.take(best[2], "entity_type", f"Type: {best[1]}", best[1])
    if out.entity_type is None and (out.openness or out.params_min or out.params_max or out.context_min or out.reasoning or out.memory_gb or out.max_output_price or out.max_input_price):
        out.entity_type = "model"
        out.compiled.append({"filter": "entity_type", "label": "Type: model (implied)", "value": "model", "source_span": None})

    # --- sort hints
    for w, s in sorted(SORT_WORDS.items(), key=lambda kv: -len(kv[0])):
        m = c.search(rf"\b{re.escape(w)}\b")
        if m:
            out.sort = s
            c.take(m, "sort", f"Sort: {s}", s)
            break

    # --- organization: "by <Org>" / "from <Org>" / "<Org> models" (verified against the organization alias table by the API layer)
    m = c.search(r"\b(?:by|from)\s+([A-Za-z][\w.&-]+(?:\s+[A-Z][\w.&-]+)?)")
    if m and m.group(1).lower() not in STOP and m.group(1).lower() not in {w for ws in TYPE_WORDS.values() for w in ws}:
        out.organization = q[m.start(1):m.end(1)]
        c.take(m, "organization", f"Organization: {out.organization}", out.organization)
    else:
        # "<Org> models …" only at the very start of the query — a capitalised word in the middle ("introducing MoE models") is free text
        m = re.match(r"\s*(?:(?:all|the|show|list|find)\s+)?([A-Z][\w.&-]*(?:\s+[A-Z][\w.&-]*)?)\s+(?:models?|llms?|papers?|releases?)\b", q)
        if m and not any(s <= m.start(1) < e for s, e in c.consumed) and m.group(1).lower() not in STOP and m.group(1).lower() not in {w for ws in TYPE_WORDS.values() for w in ws} \
           and m.group(1).lower() not in ("open", "reasoning", "vision", "proprietary", "multimodal", "cheapest", "newest", "largest", "best", "new", "moe", "small", "large"):
            out.organization = m.group(1)
            c.consumed.append((m.start(1), m.end(1)))
            out.compiled.append({"filter": "organization", "label": f"Organization: {out.organization}", "value": out.organization, "source_span": m.group(1)})

    residual = c.residual()
    residual = re.sub(r"\b(" + "|".join(sorted(STOP, key=len, reverse=True)) + r")\b", " ", residual, flags=re.IGNORECASE)
    residual = re.sub(r"[<>≤≥$]", " ", residual)
    out.residual = " ".join(residual.split())
    out.filters["residual"] = out.residual
    out.unrecognised = [w for w in re.findall(r"[\w.\-]+", out.residual) if len(w) > 1] if out.has_structure else []
    return out


def params_bound_from_memory(memory_gb: float, quant: str = "4bit") -> int:
    """ESTIMATED largest parameter count that fits `memory_gb` at 8K context (inverse of services.hardware_fit)."""
    kv = hf.KV_GB_PER_8K
    avail = max(0.0, memory_gb - hf.RESERVED_GB - kv)
    return int(avail * 1e9 / (hf.BYTES_PER_PARAM.get(quant, 0.5) * hf.OVERHEAD))


# ------------------------------------------------------------------------------------------------------------------ SQL


def _num(path: str) -> str:
    return f"(case when {path} ~ '^-?[0-9]+(\\.[0-9]+)?$' then ({path})::double precision end)"


def where_for(q: Query) -> tuple[list[str], dict[str, Any]]:
    """Filters shared by `search_entities` and the count query. Every numeric cast is regex-guarded."""
    where = ["e.merged_into is null"]
    params: dict[str, Any] = {}
    if q.entity_type:
        where.append("e.entity_type = :etype" if q.entity_type != "model" else "e.entity_type = 'model'")
        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')")
    date_col = "coalesce(e.attributes->>'release_date', e.attributes->>'published_at')"
    if q.year_from:
        where.append(f"left({date_col}, 4) >= :yf")
        params["yf"] = str(q.year_from)
    if q.year_to:
        where.append(f"left({date_col}, 4) <= :yt")
        params["yt"] = str(q.year_to)
    if q.days_back:
        where.append(f"({date_col} >= to_char((now() at time zone 'UTC') - make_interval(days => :days), 'YYYY-MM-DD') or e.first_seen_at >= now() - make_interval(days => :days))")
        params["days"] = int(q.days_back)
    pmax = q.params_max
    if q.memory_gb is not None:
        bound = params_bound_from_memory(q.memory_gb)
        pmax = min(pmax, bound) if pmax else bound
        where.append("e.attributes ? 'parameter_count'")
    if q.params_min:
        where.append(f"{_num('e.attributes->>' + chr(39) + 'parameter_count' + chr(39))} >= :pmin")
        params["pmin"] = float(q.params_min)
    if pmax:
        where.append(f"{_num('e.attributes->>' + chr(39) + 'parameter_count' + chr(39))} <= :pmax")
        params["pmax"] = float(pmax)
    if q.context_min:
        where.append(f"{_num('e.attributes->>' + chr(39) + 'context_length' + chr(39))} >= :cmin")
        params["cmin"] = float(q.context_min)
    for i, mod in enumerate(q.modalities):
        where.append(f"(e.attributes->'modalities' ? :mod{i} or e.attributes->'modalities_input' ? :mod{i} or e.attributes->'modalities_output' ? :mod{i}"
                     + (" or e.attributes->>'vision' = 'true'" if mod == "image" else "") + ")")
        params[f"mod{i}"] = mod
    if q.reasoning:
        where.append("e.attributes->>'reasoning' = 'true'")
    if q.max_output_price is not None:
        where.append("exists (select 1 from prices p where p.model_id = e.id and p.valid_to is null and p.output_per_mtok > 0 and p.output_per_mtok <= :maxout)")
        params["maxout"] = float(q.max_output_price)
    if q.max_input_price is not None:
        where.append("exists (select 1 from prices p where p.model_id = e.id and p.valid_to is null and p.input_per_mtok > 0 and p.input_per_mtok <= :maxin)")
        params["maxin"] = float(q.max_input_price)
    if q.license_key:
        where.append("(e.attributes->>'license_key' = :lic or lower(e.attributes->>'license') = any(cast(:lic_raw as text[])))")
        params["lic"] = q.license_key
        params["lic_raw"] = _raw_license_labels(q.license_key)
    if q.commercial_use:
        keys = _commercial_keys()
        where.append("(e.attributes->>'license_key' = any(cast(:ckeys as text[])) or lower(e.attributes->>'license') = any(cast(:craw as text[])))")
        params["ckeys"] = keys
        params["craw"] = [lbl for k in keys for lbl in _raw_license_labels(k)]
    if q.organization:
        where.append("exists (select 1 from entities o where o.id = e.organization_id and (o.canonical_name ilike :org or o.slug = :orgslug "
                     "or exists (select 1 from entity_aliases a where a.entity_id = o.id and a.alias ilike :org)))")
        params["org"] = q.organization
        params["orgslug"] = q.organization.lower().replace(" ", "-")
    return where, params


def _raw_license_labels(key: str) -> list[str]:
    from aiatlas.ontology.licenses import LICENSES

    info = LICENSES.get(key)
    if not info:
        return [key.lower()]
    return sorted({key.lower(), *(a.lower() for a in info.aliases), *( [info.spdx.lower()] if info.spdx else [])})


def _commercial_keys() -> list[str]:
    from aiatlas.ontology.licenses import LICENSES

    return sorted(k for k, v in LICENSES.items() if v.commercial_use is True and v.weights_downloadable)


SORT_SQL = {
    "cheapest": "(select min(p.output_per_mtok) from prices p where p.model_id = e.id and p.valid_to is null and p.output_per_mtok > 0) asc nulls last",
    "newest": "coalesce(e.attributes->>'release_date', e.attributes->>'published_at', '') desc",
    "largest": _num("e.attributes->>'parameter_count'") + " desc nulls last",
    "smallest": _num("e.attributes->>'parameter_count'") + " asc nulls last",
}


async def search_entities(conn: AsyncConnection, q: Query, *, limit: int = 30, offset: int = 0, embedding: list[float] | None = None) -> list[dict[str, Any]]:
    where, params = where_for(q)
    params.update({"limit": limit, "offset": offset})
    text = (q.residual or q.filters.get("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
    rank = "coalesce((e.quality->>'score')::float, 0) / 100.0"
    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))")
        rank = ("(ts_rank_cd(e.search, to_tsquery('simple', :qprefix)) * 2 + similarity(e.canonical_name, :q) * 3 "
                "+ case when e.canonical_name ilike :qlike then 1 else 0 end + coalesce((e.quality->>'score')::float, 0) / 200.0 "
                "+ case e.entity_type when 'model' then 0.3 when 'company' then 0.3 when 'provider' then 0.2 else 0 end)")
    if embedding is not None:
        params["vec"] = "[" + ",".join(f"{x:.6f}" for x in embedding) + "]"
        rank = f"({rank}) + coalesce(1 - (x.embedding <=> cast(:vec as vector)), 0) * 2"
        join = "left join entity_embeddings x on x.entity_id = e.id"
    else:
        join = ""
    order = SORT_SQL.get(q.sort or "", None)
    order_sql = f"{order}, rank desc" if order and q.entity_type in (None, "model") else "rank desc, e.updated_at desc"
    sql = f"""select e.id, e.entity_type, e.canonical_name, e.slug, left(e.description, 240) as description, e.status, e.attributes, e.quality,
                     o.canonical_name as organization_name, o.slug as organization_slug, {rank} as rank
              from entities e left join entities o on o.id = e.organization_id {join}
              where {' and '.join(where)} order by {order_sql} limit :limit offset :offset"""
    return await fetch_all(conn, sql, **params)


async def verify_organization(conn: AsyncConnection, name: str | None) -> dict[str, Any] | None:
    """The compiler only *proposes* an organization; it counts only when it matches an organization-like entity or one of its aliases."""
    if not name:
        return None
    rows = await fetch_all(conn, """select e.id, e.slug, e.canonical_name from entities e where e.entity_type in ('company','organization','lab','university') and e.merged_into is null
                                    and (e.canonical_name ilike :n or e.slug = :s or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :n))
                                    order by length(e.canonical_name) limit 1""", n=name, s=name.lower().replace(" ", "-"))
    return rows[0] if rows else None


async def suggest(conn: AsyncConnection, prefix: str, *, limit: int = 8) -> list[dict[str, Any]]:
    return await fetch_all(conn, """select e.id, e.entity_type, e.canonical_name, e.slug, o.canonical_name as organization_name
                                    from entities e left join entities o on o.id = e.organization_id
                                    where e.merged_into is null and (e.canonical_name ilike :p or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :p))
                                    order by case e.entity_type when 'model' then 0 when 'company' then 1 when 'provider' then 2 else 3 end,
                                             coalesce((e.quality->>'score')::float, 0) desc, length(e.canonical_name) limit :n""", p=f"{prefix}%", n=limit)


__all__ = ["Query", "compile_query", "params_bound_from_memory", "search_entities", "suggest", "verify_organization", "where_for"]