from typing import Iterable, List, Optional, Dict, Callable
from datetime import datetime

from analogai.foundation.modules.belief.belief import Belief
from analogai.foundation.modules.shared.statement_base import StatementBase


class BeliefAssembler:
    def build_belief(
        self,
        statements: Iterable[StatementBase],
        belief_id: Optional[str] = None,
    ) -> Belief:
        statements_list = list(statements)
        if not statements_list:
            raise ValueError("Cannot build belief without statements")
        ordered = sorted(
            statements_list,
            key=lambda stmt: getattr(stmt, "created_at", datetime.min),
        )
        first = ordered[0]
        belief = Belief(
            agent_id=first.agent_id,
            id=belief_id or first.id,
        )
        for statement in ordered:
            belief.add_statement(statement)
        created_candidates: List[datetime] = []
        for stmt in ordered:
            created_at = getattr(stmt, "created_at", None)
            if isinstance(created_at, datetime):
                created_candidates.append(created_at)
        created_at = min(created_candidates) if created_candidates else None
        updated_at = max(created_candidates) if created_candidates else None
        if created_at:
            belief.created_at = created_at
        if updated_at:
            belief.updated_at = updated_at
        return belief

    def group_statements(
        self,
        statements: Iterable[StatementBase],
        key_builder: Callable[[StatementBase], object],
    ) -> Dict[object, List[StatementBase]]:
        groups: Dict[object, List[StatementBase]] = {}
        for statement in statements:
            key = key_builder(statement)
            groups.setdefault(key, []).append(statement)
        return groups
