"""Build company-linked evidence groups from entities and evidence spans.

Groups evidence spans by the company they are associated with.
A span can belong to multiple groups if it mentions multiple companies.
Relations can add additional evidence linkage (e.g., inferred exposure).
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Protocol

from services.intelligence_pipeline_v3.sentiment.models import EvidenceGroup


class EntityLike(Protocol):
    """Protocol for entity objects that link to companies and evidence."""

    @property
    def company_id(self) -> str | None: ...

    @property
    def evidence_id(self) -> str: ...


@dataclass(frozen=True)
class EvidenceSpanInput:
    """Minimal evidence span input for grouping."""

    id: str
    text: str


def build_evidence_groups(
    entities: list[dict[str, str | None]],
    evidence_spans: dict[str, str],
    relations: list[dict[str, str | None]] | None = None,
) -> dict[str, EvidenceGroup]:
    """Build company-linked evidence groups from entity-company associations.

    Groups evidence spans by the company they relate to, using both
    direct entity associations and relation-based linkages. A span
    can appear in multiple groups when it mentions multiple companies.

    Parameters
    ----------
    entities
        List of dicts with keys: company_id (str or None), evidence_id (str).
        Each entity associates an evidence span with a resolved company.
        Entities without a company_id are skipped.
    evidence_spans
        Mapping of evidence_id -> text content for each evidence span.
    relations
        Optional list of dicts with keys: company_id (str or None),
        evidence_id (str or None), relation_type (str or None).
        Relations link additional evidence to companies (e.g., via
        directly_affects or inferred_exposure edges).

    Returns
    -------
    dict[str, EvidenceGroup]
        Mapping of company_id -> EvidenceGroup containing all evidence
        associated with that company.
    """
    # Accumulate evidence IDs per company
    company_evidence: dict[str, list[str]] = {}

    for entity in entities:
        company_id = entity.get("company_id")
        evidence_id = entity.get("evidence_id")

        if company_id is None or evidence_id is None:
            continue

        if company_id not in company_evidence:
            company_evidence[company_id] = []

        # Avoid duplicate evidence IDs per company
        if evidence_id not in company_evidence[company_id]:
            company_evidence[company_id].append(evidence_id)

    # Process relations for additional evidence linkage
    if relations:
        for relation in relations:
            company_id = relation.get("company_id")
            evidence_id = relation.get("evidence_id")

            if company_id is None or evidence_id is None:
                continue

            if company_id not in company_evidence:
                company_evidence[company_id] = []

            if evidence_id not in company_evidence[company_id]:
                company_evidence[company_id].append(evidence_id)

    # Build EvidenceGroup objects
    groups: dict[str, EvidenceGroup] = {}
    for company_id, evidence_ids in company_evidence.items():
        texts = []
        valid_ids = []
        for eid in evidence_ids:
            text = evidence_spans.get(eid)
            if text is not None:
                valid_ids.append(eid)
                texts.append(text)

        if valid_ids:
            groups[company_id] = EvidenceGroup(
                company_id=company_id,
                evidence_ids=valid_ids,
                texts=texts,
            )

    return groups