"""组合规则 tagger + LLM tagger 的总入口。"""

from __future__ import annotations
import asyncio
from fastapi_modules.fastapi_leaudit.govdoc_engine.models import Document
from fastapi_modules.fastapi_leaudit.govdoc_engine.parser.role_tagger_rule import RuleBasedTagger
from fastapi_modules.fastapi_leaudit.govdoc_engine.parser.role_tagger_llm import LlmTagger
from fastapi_modules.fastapi_leaudit.govdoc_engine.llm.client import LlmClient


class RoleTagger:
    """两段式：先规则打标，置信度 < threshold 的段落送 LLM 兜底。"""

    def __init__(
        self,
        llm_client: LlmClient | None = None,
        threshold: float = 0.8,
    ):
        self.rule = RuleBasedTagger()
        self.llm = LlmTagger(llm_client) if llm_client else None
        self.threshold = threshold

    def _low_conf_indices(self, doc: Document) -> list[int]:
        return [
            i for i, p in enumerate(doc.paragraphs)
            if p.role_confidence < self.threshold
        ]

    def tag(self, doc: Document) -> None:
        self.rule.tag(doc)
        if self.llm is None:
            return
        for i in self._low_conf_indices(doc):
            role, conf = self.llm.disambiguate(doc, i)
            doc.paragraphs[i].role = role
            doc.paragraphs[i].role_confidence = conf

    async def tag_async(self, doc: Document) -> None:
        self.rule.tag(doc)
        if self.llm is None:
            return
        targets = self._low_conf_indices(doc)
        if not targets:
            return
        results = await asyncio.gather(
            *(self.llm.disambiguate_async(doc, i) for i in targets)
        )
        for i, (role, conf) in zip(targets, results):
            doc.paragraphs[i].role = role
            doc.paragraphs[i].role_confidence = conf
