"""/companies listing (company, organization, lab, university) and /companies/{slug} alias."""
from __future__ import annotations

from typing import Any

from fastapi import APIRouter, Query, Request

from aiatlas.api.common import (
    COMPANY_TYPES,
    ENTITY_COLS,
    ENTITY_FROM,
    PAGINATION,
    ApiError,
    Pagination,
    cached,
    entity_summary,
    page,
)
from aiatlas.api.routers.entities import detail_for_type
from aiatlas.db import connection, fetch_all, fetch_val

router = APIRouter(prefix="/api/v1/companies", tags=["companies"])

MODEL_COUNT = "(select count(*) from entities m where m.organization_id = e.id and m.entity_type = 'model' and m.merged_into is null)"
PAPER_COUNT = ("(select count(*) from entities m where m.organization_id = e.id and m.entity_type = 'paper' and m.merged_into is null) + "
               "(select count(*) from relations r join entities pp on pp.id = r.subject_id where r.object_id = e.id and r.valid_to is null and pp.entity_type = 'paper' and pp.organization_id is distinct from e.id)")
SORTS = {"models": "model_count desc, e.canonical_name", "name": "e.canonical_name asc", "updated": "e.updated_at desc",
         "quality": "coalesce((e.quality->>'score')::float, 0) desc, e.canonical_name", "papers": "paper_count desc, e.canonical_name"}


@router.get("")
@cached(300)
async def list_companies(request: Request, p: Pagination = PAGINATION, q: str | None = Query(None, max_length=200), country: str | None = None, kind: str | None = None,
                         sort: str = "models", facets: int = Query(0, ge=0, le=1)) -> dict[str, Any]:
    if sort not in SORTS:
        raise ApiError(400, f"sort must be one of {', '.join(SORTS)}")
    where = ["e.entity_type = any(cast(:types as text[]))", "e.merged_into is null"]
    params: dict[str, Any] = {"types": list(COMPANY_TYPES)}
    if q:
        where.append("(e.canonical_name ilike :qlike or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :qlike))")
        params["qlike"] = f"%{q}%"
    if country:
        where.append("e.attributes->>'country' ilike :country")
        params["country"] = country
    if kind:
        where.append("(e.attributes->>'org_kind' ilike :kind or e.entity_type = :kind)")
        params["kind"] = kind
    where_sql = " and ".join(where)
    async with connection() as conn:
        rows = await fetch_all(conn, f"select {ENTITY_COLS}, {MODEL_COUNT} as model_count, {PAPER_COUNT} as paper_count from {ENTITY_FROM} where {where_sql} "
                                     f"order by {SORTS[sort]}, e.id limit :lim offset :off", lim=p.limit, off=p.offset, **params)
        total = await fetch_val(conn, f"select count(*) from {ENTITY_FROM} where {where_sql}", **params)
        out = page([{**(entity_summary(r) or {}), "model_count": int(r["model_count"]), "paper_count": int(r["paper_count"])} for r in rows], int(total or 0), p)
        if facets:
            countries = await fetch_all(conn, f"select e.attributes->>'country' as value, count(*) as count from {ENTITY_FROM} where {where_sql} and e.attributes ? 'country' group by 1 order by 2 desc limit 60", **params)
            kinds = await fetch_all(conn, f"select coalesce(e.attributes->>'org_kind', e.entity_type) as value, count(*) as count from {ENTITY_FROM} where {where_sql} group by 1 order by 2 desc", **params)
            out["facets"] = {"countries": [{"value": r["value"], "count": int(r["count"])} for r in countries], "kinds": [{"value": r["value"], "count": int(r["count"])} for r in kinds]}
    return out


@router.get("/{slug}")
@cached(300)
async def get_company(request: Request, slug: str) -> dict[str, Any]:
    return await detail_for_type(slug, COMPANY_TYPES)