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

    "open models released in 2026 with more than 100B parameters and 128k context"
      → {entity_type: model, openness: open, year_from: 2026, params_min: 1e11, context_min: 128000}
"""
from __future__ import annotations

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

from sqlalchemy.ext.asyncio import AsyncConnection

from aiatlas.db import fetch_all
from aiatlas.sdk.extract.numbers import parse_context_length, parse_param_count

TYPE_WORDS = {
    "model": ("model", "models", "llm", "llms", "language model"), "company": ("company", "companies", "startup", "startups", "lab", "labs", "organization"),
    "paper": ("paper", "papers", "research", "publication", "preprint", "arxiv"), "provider": ("provider", "providers", "inference provider", "api provider"),
    "benchmark": ("benchmark", "benchmarks", "leaderboard", "eval", "evals"), "hardware": ("gpu", "gpus", "chip", "chips", "hardware", "accelerator", "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": "multimodal", "image": "image", "audio": "audio", "speech": "audio", "video": "video", "code": "code",
                  "coding": "code", "embedding": "embedding", "embeddings": "embedding"}


@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)

    def as_dict(self) -> dict[str, Any]:
        return {k: v for k, v in self.__dict__.items() if v not in (None, "", [], {})}


def compile_query(q: str) -> Query:
    s = q.strip()
    low = f" {s.lower()} "
    out = Query(text=s)
    for etype, words in TYPE_WORDS.items():
        if any(f" {w} " in low for w in words):
            out.entity_type = etype
            break
    if re.search(r"\b(open[- ]?(weight|weights|source)|open models?)\b", low):
        out.openness = "open"
    elif re.search(r"\b(proprietary|closed)\b", low):
        out.openness = "proprietary"
    m = re.search(r"\b(?:released|launched|published|from|since|after|in)\s+(20\d\d)\b", low)
    if m:
        out.year_from = int(m.group(1))
        if re.search(r"\bin\s+" + m.group(1), low):
            out.year_to = out.year_from
    m = re.search(r"\b(?:before|until)\s+(20\d\d)\b", low)
    if m:
        out.year_to = int(m.group(1))
    m = re.search(r"(?:more than|over|above|>|at least|≥)\s*(\d+(?:\.\d+)?\s*[bBmMtT])\b(?:\s*param)?", s)
    if m:
        out.params_min = parse_param_count(m.group(1) + " params")
    m = re.search(r"(?:less than|under|below|<|at most|≤)\s*(\d+(?:\.\d+)?\s*[bBmMtT])\b", s)
    if m:
        out.params_max = parse_param_count(m.group(1) + " params")
    m = re.search(r"(?:context|ctx)\s*(?:window\s*)?(?:>|over|above|of at least|at least|≥)?\s*(\d+\s*[kKmM])", s) or \
        re.search(r"(\d+\s*[kKmM])\s*(?:\+\s*)?(?:token\s*)?context", s)
    if m:
        out.context_min = parse_context_length(m.group(1))
    for w, mod in MODALITY_WORDS.items():
        if f" {w} " in low and mod not in out.modalities:
            out.modalities.append(mod)
    m = re.search(r"\b(?:by|from)\s+([A-Z][\w.&-]+(?:\s+[A-Z][\w.&-]+)?)", s)
    if m:
        out.organization = m.group(1)
    # residual free-text (remove the structured bits)
    residual = re.sub(r"\b(more than|over|above|less than|under|below|at least|at most|released|launched|published|since|after|before|until|with|and|in|from|by|the|a|an|context|window|params?|parameters)\b", " ", s, flags=re.IGNORECASE)
    residual = re.sub(r"\d+(?:\.\d+)?\s*[bBmMkKtT]\b|\b20\d\d\b|[<>≤≥]", " ", residual)
    for words in TYPE_WORDS.values():
        for w in words:
            residual = re.sub(rf"\b{re.escape(w)}\b", " ", residual, flags=re.IGNORECASE)
    residual = re.sub(r"\b(open[- ]?(weight|weights|source)|proprietary|closed)\b", " ", residual, flags=re.IGNORECASE)
    out.filters["residual"] = " ".join(residual.split())
    return out


async def search_entities(conn: AsyncConnection, q: Query, *, limit: int = 30, offset: int = 0, embedding: list[float] | None = None) -> list[dict[str, Any]]:
    where = ["e.merged_into is null"]
    params: dict[str, Any] = {"limit": limit, "offset": offset}
    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)
    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 = ""
    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 rank desc, e.updated_at desc limit :limit offset :offset"""
    return await fetch_all(conn, sql, **params)


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", "search_entities", "suggest"]