"""Business-Unit (BU) resolution helpers.

A BU is the top level of the org tree (e.g. "AI Innovation",
"Engineering Group", "Finance"). Helix doesn't store a dedicated
``business_unit`` column on Person — instead each Person has an
``identity_org_memberships`` row whose ``org_path`` is something
like ``"PingCAP / Engineering Group / TiDB / Cloud SaaS"``. The
top-level segment (after stripping the tenant root "PingCAP") is
the BU.

This module centralises that derivation so the asset cost-stats,
AI spending dashboard, and any future BU-scoped query share one
truth. Callers should not re-implement the "first segment of
org_path minus PingCAP" rule in their own SQL — the dedup logic
that drops the tenant root and handles the PingCAP-ATI alias is
delicate enough that one place is enough.
"""
from __future__ import annotations

import logging
import uuid
from collections.abc import Iterable

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.identity.person_models import (
    IdentityOrgMembership,
    IdentityOrgNode,
    IdentityPerson,
)

logger = logging.getLogger(__name__)

# Tenant-root labels stripped from ``org_path`` before computing the
# BU. Mirrors the same set used by ``IdentityPerson.department_chain``.
_TENANT_ROOTS = {"pingcap", "pingcap-ati"}


def _strip_tenant_root(parts: list[str]) -> list[str]:
    if parts and parts[0].lower() in _TENANT_ROOTS:
        return parts[1:]
    return parts


def department_segments(org_path: str | None) -> list[str]:
    """Split an ``org_path`` into its department chain, tenant root stripped.

    ``org_path`` is stored slash-delimited as
    ``"PingCAP / Engineering Group / TiDB / Cloud SaaS"``. Returns
    ``["Engineering Group", "TiDB", "Cloud SaaS"]`` — the same
    tenant-root-stripped chain ``IdentityPerson.department_chain`` yields,
    kept here so callers that only hold an ``org_path`` string (not a
    loaded Person) share the one truth instead of re-splitting it. Empty
    list for a missing path or a path that is only the tenant root.
    """
    raw = (org_path or "").strip()
    if not raw:
        return []
    parts = [p.strip() for p in raw.split("/") if p.strip()]
    return _strip_tenant_root(parts)


def bu_from_org_path(org_path: str | None, fallback_name: str | None = None) -> str | None:
    """Pick the BU label out of a single ``org_path`` string.

    ``org_path`` is the slash-delimited department chain stored on
    ``identity_org_memberships``. The first segment after the tenant
    root is the BU. Returns ``None`` if the path is missing or only
    contains the tenant root.
    """
    raw = (org_path or "").strip()
    if not raw:
        return fallback_name or None
    parts = [p.strip() for p in raw.split("/") if p.strip()]
    parts = _strip_tenant_root(parts)
    if parts:
        return parts[0]
    return fallback_name or None


async def list_bus(db: AsyncSession) -> list[dict[str, object]]:
    """List every active BU.

    Two fallback strategies, in priority order:

      1. ``identity_org_nodes`` rows with no parent and ``node_status='active'``
         — this is the canonical seeded list and matches the org-tree
         page the IT admin sees.
      2. Distinct top segments of ``identity_org_memberships.org_path``
         for active primary memberships — covers tenants whose org-node
         seed lags behind the membership sync.

    Result rows: ``{name, member_count}``. Sorted alphabetically so the
    dropdown order is stable across calls.
    """
    # Path 1 — seeded org nodes.
    org_node_rows = (await db.execute(
        select(IdentityOrgNode.org_node_name)
        .where(IdentityOrgNode.parent_org_node_source_id.is_(None))
        .where(IdentityOrgNode.node_status == "active")
        .where(IdentityOrgNode.deleted_at.is_(None))
        .where(IdentityOrgNode.org_node_type == "department")
    )).scalars().all()
    bu_names: set[str] = {n for n in org_node_rows if n and n.lower() not in _TENANT_ROOTS}

    # Path 2 — top segment of membership org_path. Always run it and
    # union with path 1, so we never hide a BU just because the
    # org_node sync was slow. Filter by ``employment_status != 'left'``
    # so a BU staffed exclusively by ex-employees doesn't appear as a
    # ghost row in the dropdown (member_count would already render as 0
    # but the BU itself shouldn't be selectable).
    membership_paths = (await db.execute(
        select(IdentityOrgMembership.org_path)
        .join(IdentityPerson, IdentityPerson.id == IdentityOrgMembership.person_id)
        .where(IdentityOrgMembership.is_primary.is_(True))
        .where(IdentityOrgMembership.deleted_at.is_(None))
        .where(IdentityPerson.employment_status != "left")
    )).scalars().all()
    for path in membership_paths:
        bu = bu_from_org_path(path)
        if bu and bu.lower() not in _TENANT_ROOTS:
            bu_names.add(bu)

    if not bu_names:
        return []

    # Compute member_count by counting active persons whose primary
    # membership resolves to each BU.
    member_paths = (await db.execute(
        select(
            IdentityOrgMembership.person_id,
            IdentityOrgMembership.org_path,
            IdentityOrgMembership.org_node_name,
        )
        .join(IdentityPerson, IdentityPerson.id == IdentityOrgMembership.person_id)
        .where(IdentityOrgMembership.is_primary.is_(True))
        .where(IdentityOrgMembership.deleted_at.is_(None))
        .where(IdentityPerson.employment_status != "left")
    )).all()
    counts: dict[str, int] = {n: 0 for n in bu_names}
    for _person_id, path, leaf in member_paths:
        bu = bu_from_org_path(path, fallback_name=leaf)
        if bu and bu in counts:
            counts[bu] += 1

    return [
        {"name": name, "member_count": counts.get(name, 0)}
        for name in sorted(bu_names, key=lambda s: s.lower())
    ]


async def resolve_bu_user_ids(
    db: AsyncSession, bu_name: str,
) -> set[uuid.UUID]:
    """Return every active person's id whose primary membership's
    BU matches ``bu_name``.

    Empty set means the BU is unknown OR has no active members. The
    caller should treat empty as "no rows" — issuing an ``IN ()``
    against the result is safe (SQLAlchemy handles empty list as
    ``WHERE 0=1``).
    """
    target = bu_name.strip()
    if not target:
        return set()
    rows = (await db.execute(
        select(
            IdentityOrgMembership.person_id,
            IdentityOrgMembership.org_path,
            IdentityOrgMembership.org_node_name,
        )
        .join(IdentityPerson, IdentityPerson.id == IdentityOrgMembership.person_id)
        .where(IdentityOrgMembership.is_primary.is_(True))
        .where(IdentityOrgMembership.deleted_at.is_(None))
        .where(IdentityPerson.employment_status != "left")
    )).all()
    target_lower = target.lower()
    out: set[uuid.UUID] = set()
    for person_id, path, leaf in rows:
        bu = bu_from_org_path(path, fallback_name=leaf)
        if bu and bu.lower() == target_lower:
            out.add(person_id)
    return out


async def bu_for_user_ids(
    db: AsyncSession, user_ids: Iterable[uuid.UUID | str],
) -> dict[str, str]:
    """Map ``user_id (str)`` → BU name for the given persons.

    Accepts either ``uuid.UUID`` or string inputs — most callers reach
    this from ``DedupCostRow.dimensions["user_id"]`` which is
    already-stringified, so being permissive on the input side avoids
    silent KeyError downstream. Non-UUID strings are skipped (logged
    by SQLAlchemy if debug is on, but never propagated as an
    exception so one bad row doesn't crash the whole rollup).

    Persons without a primary membership are simply absent from the
    return dict — callers should treat absence as "unattributed".
    """
    parsed: list[uuid.UUID] = []
    for u in user_ids:
        if not u:
            continue
        if isinstance(u, uuid.UUID):
            parsed.append(u)
            continue
        try:
            parsed.append(uuid.UUID(str(u)))
        except (TypeError, ValueError):
            # Bad row — skip rather than crash the rollup. The
            # dimensions dict on a DedupCostRow shouldn't carry a
            # non-UUID user_id, but defending against it costs nothing.
            # Log so a real schema regression is observable instead
            # of silently corrupting BU rollups.
            logger.warning(
                "bu_for_user_ids skipping non-UUID user_id: %r", u,
            )
            continue
    # De-dup AFTER parsing (sets of UUID are equality-correct).
    parsed = list({u for u in parsed})
    if not parsed:
        return {}
    rows = (await db.execute(
        select(
            IdentityOrgMembership.person_id,
            IdentityOrgMembership.org_path,
            IdentityOrgMembership.org_node_name,
        )
        .where(IdentityOrgMembership.is_primary.is_(True))
        .where(IdentityOrgMembership.deleted_at.is_(None))
        .where(IdentityOrgMembership.person_id.in_(parsed))
    )).all()
    out: dict[str, str] = {}
    for person_id, path, leaf in rows:
        bu = bu_from_org_path(path, fallback_name=leaf)
        if bu:
            out[str(person_id)] = bu
    return out


__all__ = [
    "bu_for_user_ids",
    "bu_from_org_path",
    "department_segments",
    "list_bus",
    "resolve_bu_user_ids",
]