from typing import TypedDict, List
import operator
from typing_extensions import Annotated
from langchain_core.messages import AnyMessage
from agents.auditor import AuditIssue, GlobalAuditOutput


def merge_dicts(a: dict, b: dict) -> dict:
    return {**(a or {}), **(b or {})}


class AuditorPanelState(TypedDict):
    project_id: str
    program_name: str
    content: str
    is_external_audit: bool
    # Agregacja błędów z poszczególnych ról
    issues: Annotated[List[AuditIssue], operator.add]
    perspectives_summary: Annotated[dict, merge_dicts]
    # Przechowuje scores do finalnego uśrednienia
    perspective_scores: Annotated[List[int], operator.add]

    # Zarządzanie Dynamic Query Routing dla Prawnika
    legal_attempts: int
    legal_queries: Annotated[List[str], operator.add]
    messages: Annotated[
        List[AnyMessage], operator.add
    ]  # służy do wymiany zapytań z narzędziami prawnika
    prawnik_done: bool

    # Zarządzanie Dynamic Query Routing dla Finansisty
    finansista_attempts: int
    finansista_queries: Annotated[List[str], operator.add]
    finansista_messages: Annotated[List[AnyMessage], operator.add]
    finansista_done: bool

    # Zarządzanie Dynamic Query Routing dla Innowatora
    innowator_attempts: int
    innowator_queries: Annotated[List[str], operator.add]
    innowator_messages: Annotated[List[AnyMessage], operator.add]
    innowator_done: bool

    # Wynik końcowy (wyliczony przez Zarządzającego)
    final_output: GlobalAuditOutput