"""Bounded-query spend snapshot for the organisation-wide API-key report.

The per-assignment spend API deliberately resolves one key at a time.  The
fleet report cannot call it N times: provider, binding, transfer and fact
lookups then grow linearly with the number of keys.  This module projects the
same attribution rules from a fixed set of bulk lookups and grouped fact
queries.  It returns only the fields the fleet report consumes; detail/daily
series remain owned by ``spend.compute_assignment_spend``.
"""
from __future__ import annotations

import uuid
from collections import defaultdict
from datetime import date, datetime
from decimal import Decimal
from typing import Any
from uuid import UUID

from sqlalchemy import func, or_, select, tuple_
from sqlalchemy.ext.asyncio import AsyncSession

from app.domains.ai_spending.whitelist import load_whitelist
from app.domains.apikey.models import ApiKeyAssignment
from app.domains.apikey.spend import (
    BUCKET_BAILIAN,
    BUCKET_BEDROCK,
    BUCKET_FOUNDRY,
    BUCKET_UNKNOWN,
    BUCKET_VERTEX,
    PER_KEY_PROFILE_AXES,
    _coalesced_epochs,
    _event_ts,
    _retag_state_for_event,
    _transfer_epoch_end_with_retag,
    _TransferEvent,
    classify_bucket,
    parse_iam_user_name,
)
from app.domains.bailian.models import (
    BAILIAN_AGGREGATE_MODEL,
    BailianApiKey,
    BailianUsageDaily,
)
from app.domains.cmdb.models import (
    CmdbBedrockProfileFact,
    CmdbCloudScope,
    CmdbCostFact,
    CmdbIamUser,
)
from app.identity.person_models import IdentityPerson

_LIVE = ("active", "approved", "expired", "pending_release")
_FORCE_ACTION = "apikey.transfer.force_complete"
_FORCE_RESOURCE = "api_key_assignment"


def _eligible(a: ApiKeyAssignment, month_start: date) -> bool:
    if a.status in _LIVE:
        return True
    if a.status != "revoked" or a.revoked_at is None:
        return False
    revoked = a.revoked_at.date() if isinstance(a.revoked_at, datetime) else a.revoked_at
    return revoked >= month_start


def _service_filter(patterns: list[str]):
    def esc(value: str) -> str:
        return value.replace("\\", "\\\\").replace("%", r"\%").replace("_", r"\_")

    return or_(*(CmdbCostFact.service.ilike(f"%{esc(p)}%", escape="\\") for p in patterns))


def _summary(bucket: str, amount: Decimal = Decimal(0), fact_date: date | None = None,
             *, available: bool = True) -> dict[str, Any]:
    return {
        "bucket": bucket,
        "mtd_usd": amount,
        "last_fact_date": fact_date.isoformat() if fact_date else None,
        "data_available": available,
    }


async def _transfer_events(
    session: AsyncSession,
    assignment_ids: list[UUID],
) -> tuple[dict[UUID, list[_TransferEvent]], set[UUID]]:
    """Load both transfer provenance stores once for the whole fleet."""
    events: dict[UUID, list[_TransferEvent]] = defaultdict(list)
    person_ids: set[UUID] = set()
    if not assignment_ids:
        return events, person_ids

    # Deferred cross-domain imports keep the API-key import graph acyclic,
    # matching the single-assignment implementation in ``spend.py``.
    from app.audit.models import AuditLog
    from app.domains.offboarding.models import (
        XFER_ACCEPTED,
        XFER_COMPLETED,
        AssetTransferRequest,
    )

    rows = (await session.execute(
        select(
            AssetTransferRequest.asset_ref_id,
            AssetTransferRequest.id,
            AssetTransferRequest.initiator_id,
            AssetTransferRequest.completed_at,
        ).where(
            AssetTransferRequest.asset_type == "api_key",
            AssetTransferRequest.asset_ref_id.in_(assignment_ids),
            AssetTransferRequest.status.in_((XFER_ACCEPTED, XFER_COMPLETED)),
            AssetTransferRequest.completed_at.is_not(None),
        )
    )).all()
    for assignment_id, transfer_id, initiator_id, completed_at in rows:
        events[assignment_id].append(_TransferEvent(
            completed_at=completed_at,
            from_person_id=initiator_id,
            transfer_request_id=transfer_id,
        ))
        if initiator_id:
            person_ids.add(initiator_id)

    audit_rows = (await session.execute(
        select(AuditLog.resource_id, AuditLog.detail, AuditLog.created_at).where(
            AuditLog.action == _FORCE_ACTION,
            AuditLog.resource_type == _FORCE_RESOURCE,
            AuditLog.resource_id.in_([str(i) for i in assignment_ids]),
            AuditLog.status == "success",
        )
    )).all()
    for raw_assignment_id, detail, created_at in audit_rows:
        if created_at is None:
            continue
        try:
            assignment_id = uuid.UUID(str(raw_assignment_id))
        except ValueError:
            continue
        raw_from = (detail or {}).get("from_person_id")
        try:
            from_id = uuid.UUID(str(raw_from)) if raw_from is not None else None
        except (AttributeError, ValueError):
            from_id = None
        events[assignment_id].append(_TransferEvent(
            completed_at=created_at,
            from_person_id=from_id,
            transfer_request_id=None,
        ))
        if from_id:
            person_ids.add(from_id)
    return events, person_ids


async def compute_assignments_spend_bulk(
    session: AsyncSession,
    *,
    assignments: list[ApiKeyAssignment],
    provider_names: dict[UUID, str],
    month_start: date,
    month_end: date,
) -> dict[UUID, dict[str, Any]]:
    """Return fleet-report spend keyed by assignment id with O(1) SQL.

    The number of statements depends on which provider buckets are present,
    never on assignment count.  All large-table reads aggregate in SQL.
    """
    if not assignments:
        return {}

    eligible = [a for a in assignments if _eligible(a, month_start)]
    buckets = {a.id: classify_bucket(provider_names.get(a.provider_id)) for a in eligible}
    result: dict[UUID, dict[str, Any]] = {}
    whitelist = await load_whitelist(session)

    # Scope-backed buckets share one binding lookup.
    scope_ids = {str(a.cloud_scope_id) for a in eligible if a.cloud_scope_id}
    scopes: dict[str, tuple[str, str]] = {}
    if scope_ids:
        scope_rows = (await session.execute(
            select(CmdbCloudScope.id, CmdbCloudScope.integration_id, CmdbCloudScope.external_id)
            .where(CmdbCloudScope.id.in_(scope_ids))
        )).all()
        scopes = {
            str(i): (str(integration_id), external_id)
            for i, integration_id, external_id in scope_rows
        }

    foundry = [a for a in eligible if buckets[a.id] == BUCKET_FOUNDRY]
    foundry_patterns = list(whitelist.get("azure") or [])
    foundry_bindings = {
        a.id: scopes.get(str(a.cloud_scope_id)) for a in foundry if a.cloud_scope_id
    }
    foundry_keys = {binding for binding in foundry_bindings.values() if binding and binding[1]}
    foundry_totals: dict[tuple[str, str], tuple[Decimal, date | None]] = {}
    if foundry_keys and foundry_patterns:
        foundry_rows = (await session.execute(
            select(
                CmdbCostFact.integration_id, CmdbCostFact.resource_group,
                func.coalesce(func.sum(CmdbCostFact.amount_usd), 0),
                func.max(CmdbCostFact.date),
            ).where(
                CmdbCostFact.date >= month_start,
                CmdbCostFact.date <= month_end,
                _service_filter(foundry_patterns),
                tuple_(CmdbCostFact.integration_id, CmdbCostFact.resource_group).in_(list(foundry_keys)),
            ).group_by(CmdbCostFact.integration_id, CmdbCostFact.resource_group)
        )).all()
        foundry_totals = {
            (str(i), rg): (Decimal(str(v or 0)), d)
            for i, rg, v, d in foundry_rows
        }
    for a in foundry:
        foundry_binding = foundry_bindings.get(a.id)
        if not foundry_binding or not foundry_binding[1] or not foundry_patterns:
            result[a.id] = _summary(BUCKET_FOUNDRY, available=False)
        else:
            amount, last = foundry_totals.get(foundry_binding, (Decimal(0), None))
            result[a.id] = _summary(BUCKET_FOUNDRY, amount, last)

    vertex = [a for a in eligible if buckets[a.id] == BUCKET_VERTEX]
    vertex_patterns = list(whitelist.get("gcp") or [])
    vertex_bindings = {
        a.id: (scopes.get(str(a.cloud_scope_id)), str(a.cloud_scope_id))
        for a in vertex if a.cloud_scope_id
    }
    vertex_keys = {
        (binding[0][0], binding[1]) for binding in vertex_bindings.values() if binding[0]
    }
    vertex_totals: dict[tuple[str, str], tuple[Decimal, date | None]] = {}
    if vertex_keys and vertex_patterns:
        vertex_rows = (await session.execute(
            select(
                CmdbCostFact.integration_id, CmdbCostFact.scope_id,
                func.coalesce(func.sum(CmdbCostFact.amount_usd), 0),
                func.max(CmdbCostFact.date),
            ).where(
                CmdbCostFact.date >= month_start,
                CmdbCostFact.date <= month_end,
                _service_filter(vertex_patterns),
                tuple_(CmdbCostFact.integration_id, CmdbCostFact.scope_id).in_(list(vertex_keys)),
            ).group_by(CmdbCostFact.integration_id, CmdbCostFact.scope_id)
        )).all()
        vertex_totals = {
            (str(i), str(s)): (Decimal(str(v or 0)), d)
            for i, s, v, d in vertex_rows
        }
    for a in vertex:
        vertex_binding = vertex_bindings.get(a.id)
        if not vertex_binding or not vertex_binding[0] or not vertex_patterns:
            result[a.id] = _summary(BUCKET_VERTEX, available=False)
        else:
            key = (vertex_binding[0][0], vertex_binding[1])
            amount, last = vertex_totals.get(key, (Decimal(0), None))
            result[a.id] = _summary(BUCKET_VERTEX, amount, last)

    # Bailian attribution is already assignment-native; aggregate through the
    # key mirror in one statement and separately identify linked zero-spend keys.
    bailian = [a for a in eligible if buckets[a.id] == BUCKET_BAILIAN]
    bailian_ids = [a.id for a in bailian]
    linked_bailian: set[UUID] = set()
    bailian_totals: dict[UUID, tuple[Decimal, date | None]] = {}
    if bailian_ids:
        linked_bailian = {
            assignment_id
            for assignment_id in (await session.execute(
                select(BailianApiKey.assignment_id).where(
                    BailianApiKey.assignment_id.in_(bailian_ids)
                ).distinct()
            )).scalars().all()
            if assignment_id is not None
        }
        bailian_rows = (await session.execute(
            select(
                BailianApiKey.assignment_id,
                func.coalesce(func.sum(BailianUsageDaily.cost_usd), 0),
                func.max(BailianUsageDaily.usage_date),
            ).join(BailianUsageDaily, BailianUsageDaily.api_key_pk == BailianApiKey.id)
            .where(
                BailianApiKey.assignment_id.in_(bailian_ids),
                BailianUsageDaily.usage_date >= month_start,
                BailianUsageDaily.usage_date <= month_end,
                BailianUsageDaily.model_name != BAILIAN_AGGREGATE_MODEL,
            ).group_by(BailianApiKey.assignment_id)
        )).all()
        bailian_totals = {
            i: (Decimal(str(v or 0)), d)
            for i, v, d in bailian_rows if i is not None
        }
    for a in bailian:
        if a.id not in linked_bailian:
            result[a.id] = _summary(BUCKET_BAILIAN, available=False)
        else:
            amount, last = bailian_totals.get(a.id, (Decimal(0), None))
            result[a.id] = _summary(BUCKET_BAILIAN, amount, last)

    await _compute_bedrock_bulk(
        session, assignments=[a for a in eligible if buckets[a.id] == BUCKET_BEDROCK],
        month_start=month_start, month_end=month_end,
        patterns=list(whitelist.get("aws") or []), result=result,
    )
    for a in eligible:
        if buckets[a.id] == BUCKET_UNKNOWN:
            result[a.id] = _summary(BUCKET_UNKNOWN, available=False)
    return result


async def _compute_bedrock_bulk(
    session: AsyncSession,
    *, assignments: list[ApiKeyAssignment], month_start: date, month_end: date,
    patterns: list[str], result: dict[UUID, dict[str, Any]],
) -> None:
    if not assignments:
        return
    bound = [a for a in assignments if a.iam_integration_id]
    unbound = [a for a in assignments if not a.iam_integration_id]
    for a in unbound:
        result[a.id] = _summary(BUCKET_BEDROCK, available=False)
    if not bound or not patterns:
        for a in bound:
            result[a.id] = _summary(BUCKET_BEDROCK, available=False)
        return

    iam_ids = {a.iam_user_id for a in bound if a.iam_user_id}
    iam_tags: dict[UUID, dict[str, Any]] = {}
    if iam_ids:
        iam_tags = {
            iam_id: tags or {}
            for iam_id, tags in (await session.execute(
                select(CmdbIamUser.id, CmdbIamUser.tags).where(
                    CmdbIamUser.id.in_(iam_ids),
                )
            )).all()
        }
    owner_person_ids = {a.assigned_to for a in bound}
    emails: dict[UUID, str | None] = {
        person_id: email
        for person_id, email in (await session.execute(
            select(IdentityPerson.id, IdentityPerson.work_email).where(
                IdentityPerson.id.in_(owner_person_ids)
            )
        )).all()
    }

    owner_by_id: dict[UUID, str | None] = {}
    for a in bound:
        tags = iam_tags.get(a.iam_user_id, {}) if a.iam_user_id else {}
        raw = tags.get("Owner") if isinstance(tags, dict) else None
        owner = raw.strip() if isinstance(raw, str) and raw.strip() != "(unallocated)" else None
        owner_by_id[a.id] = owner or (emails.get(a.assigned_to) or "").strip() or parse_iam_user_name(a.iam_user_arn)

    requested_profile_keys = {
        (str(a.iam_integration_id), (a.request_number or "").strip().lower())
        for a in bound
        if (a.request_number or "").strip()
    }
    per_key: dict[tuple[str, str], tuple[Decimal, date | None]] = {}
    if requested_profile_keys:
        profile_expr = func.lower(CmdbBedrockProfileFact.profile)
        profile_integrations = {key[0] for key in requested_profile_keys}
        profile_rows = (await session.execute(
            select(
                CmdbBedrockProfileFact.integration_id,
                profile_expr,
                func.coalesce(func.sum(CmdbBedrockProfileFact.amount_usd), 0),
                func.max(CmdbBedrockProfileFact.date),
            ).where(
                CmdbBedrockProfileFact.integration_id.in_(profile_integrations),
                func.lower(CmdbBedrockProfileFact.profile_tag_key).in_(PER_KEY_PROFILE_AXES),
                tuple_(
                    CmdbBedrockProfileFact.integration_id,
                    profile_expr,
                ).in_(list(requested_profile_keys)),
                CmdbBedrockProfileFact.date >= month_start,
                CmdbBedrockProfileFact.date <= month_end,
            ).group_by(
                CmdbBedrockProfileFact.integration_id,
                profile_expr,
            )
        )).all()
        per_key = {
            (str(i), str(ref)): (Decimal(str(v or 0)), d)
            for i, ref, v, d in profile_rows
        }

    per_key_by_id = {
        a.id: per_key.get((str(a.iam_integration_id), (a.request_number or "").strip().lower()))
        for a in bound
    }
    cohorts: dict[tuple[str, str], list[ApiKeyAssignment]] = defaultdict(list)
    for a in bound:
        owner = owner_by_id[a.id]
        if owner:
            cohorts[(str(a.iam_integration_id), owner.lower())].append(a)
    canonical_ids: set[UUID] = set()
    for cohort in cohorts.values():
        owner_axis = [a for a in cohort if per_key_by_id[a.id] is None]
        if owner_axis:
            canonical_ids.add(min(owner_axis, key=lambda a: (a.created_at, str(a.id))).id)

    events, historic_person_ids = await _transfer_events(session, [a.id for a in bound])
    if historic_person_ids:
        emails.update({
            person_id: email
            for person_id, email in (await session.execute(
                select(IdentityPerson.id, IdentityPerson.work_email).where(
                    IdentityPerson.id.in_(historic_person_ids)
                )
            )).all()
        })

    lower_resource_id = func.lower(CmdbCostFact.resource_id)
    owner_expr = func.substring_index(
        func.substring_index(lower_resource_id, "|", 1), "owner=", -1,
    )
    requested_owner_keys = {
        (str(a.iam_integration_id), owner.lower())
        for a in bound
        if (owner := owner_by_id[a.id])
    }
    for a in bound:
        for event in events.get(a.id, []):
            if not isinstance(event.from_person_id, UUID):
                continue
            historic_owner = (
                emails.get(event.from_person_id) or ""
            ).strip().lower()
            if historic_owner:
                requested_owner_keys.add(
                    (str(a.iam_integration_id), historic_owner)
                )
    owner_daily: dict[tuple[str, str], list[tuple[date, Decimal]]] = defaultdict(list)
    if requested_owner_keys:
        owner_integrations = {key[0] for key in requested_owner_keys}
        fact_rows = (await session.execute(
            select(
                CmdbCostFact.integration_id, owner_expr.label("owner"), CmdbCostFact.date,
                func.coalesce(func.sum(CmdbCostFact.amount_usd), 0),
            ).where(
                CmdbCostFact.integration_id.in_(owner_integrations),
                CmdbCostFact.date >= month_start,
                CmdbCostFact.date <= month_end,
                _service_filter(patterns),
                lower_resource_id.like("aws-principal:owner=%|%"),
                tuple_(CmdbCostFact.integration_id, owner_expr).in_(
                    list(requested_owner_keys)
                ),
            ).group_by(CmdbCostFact.integration_id, owner_expr, CmdbCostFact.date)
        )).all()
        for integration_id, owner, fact_date, value in fact_rows:
            owner_daily[(str(integration_id), str(owner))].append(
                (fact_date, Decimal(str(value or 0)))
            )

    current_claims: dict[tuple[str, str], set[UUID]] = defaultdict(set)
    for a in bound:
        owner = owner_by_id[a.id]
        if owner:
            current_claims[(str(a.iam_integration_id), owner.lower())].add(a.id)

    for a in bound:
        owner = owner_by_id[a.id]
        if not owner:
            result[a.id] = _summary(BUCKET_BEDROCK, available=False)
            continue
        isolated = per_key_by_id[a.id]
        if isolated is not None:
            result[a.id] = _summary(BUCKET_BEDROCK, isolated[0], isolated[1])
            continue
        if a.id not in canonical_ids:
            result[a.id] = _summary(BUCKET_BEDROCK, available=False)
            continue

        chain = sorted(events.get(a.id, []), key=lambda e: _event_ts(e.completed_at))
        epochs: list[tuple[str, date | None, date]] = []
        previous_end: date | None = None
        for idx, event in enumerate(chain):
            completed = event.completed_at.date() if isinstance(event.completed_at, datetime) else event.completed_at
            end = _transfer_epoch_end_with_retag(
                _retag_state_for_event(a, event=event, is_latest=idx == len(chain) - 1),
                completed_d=completed,
            )
            start, previous_end = previous_end, completed
            if event.from_person_id is None or end < month_start or (start and start > month_end):
                continue
            historic_owner = (emails.get(event.from_person_id) or "").strip()
            claims = current_claims.get((str(a.iam_integration_id), historic_owner.lower()), set())
            if historic_owner and historic_owner.lower() != owner.lower() and not (claims - {a.id}):
                epochs.append((historic_owner, start, end))
        epochs = _coalesced_epochs(epochs)

        amount = Decimal(0)
        last: date | None = None
        integration = str(a.iam_integration_id)
        matching_facts = list(owner_daily.get((integration, owner.lower()), []))
        for prior, start, end in epochs:
            matching_facts.extend(
                (fact_date, value)
                for fact_date, value in owner_daily.get(
                    (integration, prior.lower()), []
                )
                if fact_date <= end and (start is None or fact_date >= start)
            )
        for fact_date, value in matching_facts:
            amount += value
            last = max(last, fact_date) if last else fact_date
        sibling_total = sum(
            (per_key_by_id[sibling.id] or (Decimal(0), None))[0]
            for sibling in cohorts[(integration, owner.lower())]
            if sibling.id != a.id and per_key_by_id[sibling.id] is not None
        )
        result[a.id] = _summary(BUCKET_BEDROCK, max(Decimal(0), amount - sibling_total), last)


__all__ = ["compute_assignments_spend_bulk"]