"""Local embeddings (via the LLM gateway's embedding endpoint) stored in pgvector. Semantic search is optional: FTS works without it."""
from __future__ import annotations

import hashlib
import logging
from typing import Any

from aiatlas.config import settings
from aiatlas.db import execute, fetch_all, fetch_one, transaction
from aiatlas.services.llm import LLMUnavailable, gateway

log = logging.getLogger(__name__)


def entity_text(row: dict[str, Any]) -> str:
    attrs = row.get("attributes") or {}
    bits = [row["canonical_name"], row["entity_type"].replace("_", " ")]
    for k in ("family", "architecture", "openness", "license", "modalities", "country", "kind", "category"):
        v = attrs.get(k)
        if v:
            bits.append(f"{k}: {', '.join(v) if isinstance(v, list) else v}")
    if row.get("organization_name"):
        bits.append(f"by {row['organization_name']}")
    if row.get("description"):
        bits.append(row["description"][:1200])
    return "\n".join(str(b) for b in bits)


async def embed_entities(entity_ids: list[str]) -> dict[str, Any]:
    if not gateway.available:
        raise LLMUnavailable("embedding engine not configured")
    async with transaction() as conn:
        has_vector = await fetch_one(conn, "select 1 from pg_extension where extname = 'vector'")
        if not has_vector:
            return {"skipped": "pgvector missing"}
        rows = await fetch_all(conn, """select e.id, e.entity_type, e.canonical_name, e.description, e.attributes, o.canonical_name as organization_name
                                        from entities e left join entities o on o.id = e.organization_id where e.id = any(cast(:ids as text[]))""", ids=entity_ids)
    todo = []
    for r in rows:
        text = entity_text(r)
        h = hashlib.sha256((settings.embedding_model + text).encode()).hexdigest()
        async with transaction() as conn:
            existing = await fetch_one(conn, "select text_hash from entity_embeddings where entity_id = :id", id=r["id"])
        if existing and existing["text_hash"] == h:
            continue
        todo.append((r["id"], text, h))
    done = 0
    for i in range(0, len(todo), 16):
        batch = todo[i:i + 16]
        vectors = await gateway.embed([t for _, t, _ in batch])
        async with transaction() as conn:
            for (eid, _, h), vec in zip(batch, vectors, strict=False):
                if len(vec) != settings.embedding_dim:
                    log.warning("embedding dimension mismatch", extra={"got": len(vec), "want": settings.embedding_dim})
                    continue
                await execute(conn, """insert into entity_embeddings (entity_id, model, embedding, text_hash) values (:id, :m, cast(:v as vector), :h)
                                       on conflict (entity_id) do update set model = excluded.model, embedding = excluded.embedding, text_hash = excluded.text_hash, created_at = now()""",
                              id=eid, m=settings.embedding_model, v="[" + ",".join(f"{x:.6f}" for x in vec) + "]", h=h)
                done += 1
    return {"embedded": done, "skipped": len(rows) - len(todo)}


async def embed_query(text: str) -> list[float] | None:
    if not gateway.available:
        return None
    try:
        return (await gateway.embed([text]))[0]
    except Exception as exc:  # noqa: BLE001
        log.info("query embedding failed", extra={"error": str(exc)})
        return None


async def pending_entity_ids(limit: int = 200) -> list[str]:
    async with transaction() as conn:
        if not await fetch_one(conn, "select 1 from pg_extension where extname = 'vector'"):
            return []
        rows = await fetch_all(conn, """select e.id from entities e left join entity_embeddings x on x.entity_id = e.id
                                        where e.merged_into is null and e.entity_type in ('model','company','paper','provider','benchmark','hardware','framework','tool','dataset')
                                        and (x.entity_id is null or x.created_at < e.updated_at - interval '1 day') order by e.updated_at desc limit :n""", n=limit)
    return [r["id"] for r in rows]


__all__ = ["embed_entities", "embed_query", "entity_text", "pending_entity_ids"]