from __future__ import annotations

from collections.abc import Sequence
from datetime import datetime

from sqlalchemy import or_

from govoplan_core.core.identity import IdentityDirectory
from govoplan_core.core.idm import (
    IdmDirectory,
    OrganizationFunctionAssignmentRef,
    OrganizationFunctionIncumbencyRef,
)
from govoplan_core.core.organizations import OrganizationDirectory
from govoplan_core.db.session import get_database
from govoplan_core.security.time import utc_now
from govoplan_idm.backend.db.models import IdmOrganizationFunctionAssignment


def _status(active: bool) -> str:
    return "active" if active else "inactive"


def _assignment_ref(item: IdmOrganizationFunctionAssignment) -> OrganizationFunctionAssignmentRef:
    return OrganizationFunctionAssignmentRef(
        id=item.id,
        tenant_id=item.tenant_id,
        identity_id=item.identity_id,
        account_id=item.account_id,
        function_id=item.function_id,
        organization_unit_id=item.organization_unit_id,
        applies_to_subunits=item.applies_to_subunits,
        source=item.source,  # type: ignore[arg-type]
        delegated_from_assignment_id=item.delegated_from_assignment_id,
        acting_for_account_id=item.acting_for_account_id,
        valid_from=item.valid_from,
        valid_until=item.valid_until,
        status=_status(item.is_active),  # type: ignore[arg-type]
    )


class SqlIdmDirectory(IdmDirectory):
    def __init__(self, *, identities: IdentityDirectory, organizations: OrganizationDirectory) -> None:
        self._identities = identities
        self._organizations = organizations

    def get_organization_function_assignment(self, assignment_id: str) -> OrganizationFunctionAssignmentRef | None:
        with get_database().session() as session:
            item = session.get(IdmOrganizationFunctionAssignment, assignment_id)
            return _assignment_ref(item) if item is not None else None

    def organization_function_assignments_for_identity(
        self,
        identity_id: str,
        *,
        tenant_id: str | None = None,
        effective_at: datetime | None = None,
    ) -> tuple[OrganizationFunctionAssignmentRef, ...]:
        return tuple(
            self.organization_function_assignments_for_identities(
                (identity_id,),
                tenant_id=tenant_id,
                effective_at=effective_at,
            ).get(identity_id, ())
        )

    def organization_function_assignments_for_account(
        self,
        account_id: str,
        *,
        tenant_id: str | None = None,
        effective_at: datetime | None = None,
    ) -> tuple[OrganizationFunctionAssignmentRef, ...]:
        return tuple(
            self.organization_function_assignments_for_accounts(
                (account_id,),
                tenant_id=tenant_id,
                effective_at=effective_at,
            ).get(account_id, ())
        )

    def organization_function_assignments_for_identities(
        self,
        identity_ids: Sequence[str],
        *,
        tenant_id: str | None = None,
        effective_at: datetime | None = None,
    ) -> dict[str, tuple[OrganizationFunctionAssignmentRef, ...]]:
        requested = tuple(dict.fromkeys(identity_ids))
        result: dict[str, list[OrganizationFunctionAssignmentRef]] = {
            identity_id: [] for identity_id in requested
        }
        if not requested:
            return {}
        with get_database().session() as session:
            items = self._effective_assignment_items(
                session,
                tenant_id=tenant_id,
                effective_at=effective_at or utc_now(),
                identity_ids=requested,
            )
        for item in items:
            result[item.identity_id].append(_assignment_ref(item))
        return {
            identity_id: tuple(assignments)
            for identity_id, assignments in result.items()
        }

    def organization_function_assignments_for_accounts(
        self,
        account_ids: Sequence[str],
        *,
        tenant_id: str | None = None,
        effective_at: datetime | None = None,
    ) -> dict[str, tuple[OrganizationFunctionAssignmentRef, ...]]:
        requested = tuple(dict.fromkeys(account_ids))
        result: dict[str, list[OrganizationFunctionAssignmentRef]] = {
            account_id: [] for account_id in requested
        }
        if not requested:
            return {}
        identities = self._identities.identities_for_accounts(requested)
        identity_ids_by_account: dict[str, set[str]] = {
            account_id: set() for account_id in requested
        }
        requested_set = set(requested)
        for identity in identities:
            linked_accounts = set(identity.account_ids)
            if identity.primary_account_id:
                linked_accounts.add(identity.primary_account_id)
            for account_id in linked_accounts & requested_set:
                identity_ids_by_account[account_id].add(identity.id)
        identity_ids = tuple(
            dict.fromkeys(
                identity_id
                for values in identity_ids_by_account.values()
                for identity_id in values
            )
        )
        if not identity_ids:
            return {
                account_id: tuple(assignments)
                for account_id, assignments in result.items()
            }
        with get_database().session() as session:
            items = self._effective_assignment_items(
                session,
                tenant_id=tenant_id,
                effective_at=effective_at or utc_now(),
                identity_ids=identity_ids,
            )
        for account_id, linked_identity_ids in identity_ids_by_account.items():
            result[account_id].extend(
                _assignment_ref(item)
                for item in items
                if item.identity_id in linked_identity_ids
                and item.account_id in {None, account_id}
            )
        return {
            account_id: tuple(assignments)
            for account_id, assignments in result.items()
        }

    def organization_function_assignments_for_function(
        self,
        function_id: str,
        *,
        tenant_id: str | None = None,
        effective_at: datetime | None = None,
    ) -> tuple[OrganizationFunctionAssignmentRef, ...]:
        function = self._organizations.get_function(function_id)
        if function is None:
            return ()
        if tenant_id is not None and function.tenant_id != tenant_id:
            raise ValueError("Organization function belongs to another tenant.")
        resolved_tenant_id = tenant_id or function.tenant_id
        return self.organization_function_incumbencies(
            (function_id,),
            tenant_id=resolved_tenant_id,
            effective_at=effective_at,
        )[function_id].assignments

    def organization_function_incumbencies(
        self,
        function_ids: Sequence[str],
        *,
        tenant_id: str,
        effective_at: datetime | None = None,
    ) -> dict[str, OrganizationFunctionIncumbencyRef]:
        requested = tuple(dict.fromkeys(function_ids))
        functions = {}
        for function_id in requested:
            function = self._organizations.get_function(function_id)
            if function is None:
                raise ValueError(
                    f"Organization function does not exist: {function_id}"
                )
            if function.tenant_id != tenant_id:
                raise ValueError(
                    "Organization function belongs to another tenant."
                )
            functions[function_id] = function
        if not requested:
            return {}
        with get_database().session() as session:
            items = self._effective_assignment_items(
                session,
                tenant_id=tenant_id,
                effective_at=effective_at or utc_now(),
                function_ids=requested,
            )
        assignments: dict[str, list[OrganizationFunctionAssignmentRef]] = {
            function_id: [] for function_id in requested
        }
        for item in items:
            assignments[item.function_id].append(_assignment_ref(item))
        return {
            function_id: OrganizationFunctionIncumbencyRef(
                tenant_id=tenant_id,
                function_id=function_id,
                assignments=tuple(assignments[function_id]),
                function_active=functions[function_id].status == "active",
            )
            for function_id in requested
        }

    def _effective_assignment_items(
        self,
        session,
        *,
        effective_at: datetime,
        tenant_id: str | None = None,
        identity_ids: Sequence[str] = (),
        function_ids: Sequence[str] = (),
    ) -> tuple[IdmOrganizationFunctionAssignment, ...]:
        query = (
            session.query(IdmOrganizationFunctionAssignment)
            .filter(
                IdmOrganizationFunctionAssignment.is_active.is_(True),
                or_(
                    IdmOrganizationFunctionAssignment.valid_from.is_(None),
                    IdmOrganizationFunctionAssignment.valid_from <= effective_at,
                ),
                or_(
                    IdmOrganizationFunctionAssignment.valid_until.is_(None),
                    IdmOrganizationFunctionAssignment.valid_until > effective_at,
                ),
            )
            .order_by(IdmOrganizationFunctionAssignment.created_at.asc())
        )
        if identity_ids:
            query = query.filter(
                    IdmOrganizationFunctionAssignment.identity_id.in_(identity_ids),
            )
        if function_ids:
            query = query.filter(
                IdmOrganizationFunctionAssignment.function_id.in_(function_ids),
            )
        if tenant_id is not None:
            query = query.filter(
                IdmOrganizationFunctionAssignment.tenant_id == tenant_id
            )
        items = query.all()
        source_ids = {
            item.delegated_from_assignment_id
            for item in items
            if item.source in {"delegated", "acting_for"}
            and item.delegated_from_assignment_id is not None
        }
        effective_sources = (
            session.query(IdmOrganizationFunctionAssignment)
            .filter(
                IdmOrganizationFunctionAssignment.id.in_(source_ids),
                IdmOrganizationFunctionAssignment.is_active.is_(True),
                or_(
                    IdmOrganizationFunctionAssignment.valid_from.is_(None),
                    IdmOrganizationFunctionAssignment.valid_from <= effective_at,
                ),
                or_(
                    IdmOrganizationFunctionAssignment.valid_until.is_(None),
                    IdmOrganizationFunctionAssignment.valid_until > effective_at,
                ),
            )
            .all()
            if source_ids
            else ()
        )
        effective_sources_by_id = {
            item.id: item for item in effective_sources
        }
        function_cache = {}
        active_items: list[IdmOrganizationFunctionAssignment] = []
        for item in items:
            if item.source in {"delegated", "acting_for"}:
                source = effective_sources_by_id.get(
                    item.delegated_from_assignment_id or ""
                )
                if (
                    source is None
                    or source.tenant_id != item.tenant_id
                    or source.function_id != item.function_id
                ):
                    continue
            if item.function_id not in function_cache:
                function_cache[item.function_id] = (
                    self._organizations.get_function(item.function_id)
                )
            function = function_cache[item.function_id]
            if (
                function is None
                or function.status != "active"
                or function.tenant_id != item.tenant_id
            ):
                continue
            active_items.append(item)
        return tuple(active_items)