from __future__ import annotations

from collections import defaultdict

from lib.models.improvement_input import ImprovementInput


_MISMATCH_BUCKET = {
    'wrong_routing': 'workflow',
    'wrong_path': 'knowledge',
    'wrong_scope': 'knowledge',
    'stopped_too_early': 'output',
    'overclaimed_root_cause': 'output',
    'missing_known_issue': 'knowledge',
    'missing_required_evidence': 'workflow',
}


def _merge_readiness_from_case_count(case_count: int) -> str:
    if case_count >= 2:
        return 'eligible_now'
    return 'review_only'


def _build_mismatch_improvement(mismatch_type: str, case_ids: list[str]) -> ImprovementInput | None:
    bucket = _MISMATCH_BUCKET.get(mismatch_type) or 'do_not_change_yet'
    merge_readiness = _merge_readiness_from_case_count(len(case_ids))

    if mismatch_type == 'stopped_too_early':
        return ImprovementInput(
            bucket=bucket,
            title='Tighten proof-threshold promotion after first-pass routing',
            problem='Replay can reach the correct TiKV route, but it still stops at a provisional conclusion even when later owner-side evidence becomes explicit.',
            proposed_change='Add a stronger promotion rule for io_jitter and similar routes so the output names the exact evidence needed for durable root-cause promotion instead of stopping at a generic uncertainty statement.',
            evidence_case_ids=case_ids,
            merge_readiness=merge_readiness,
            mismatch_types=[mismatch_type],
        )
    if mismatch_type == 'wrong_routing':
        return ImprovementInput(
            bucket=bucket,
            title='Refine entry routing from initial ticket-visible cues',
            problem='Replay selected the wrong entry mode compared with the eventual human conclusion.',
            proposed_change='Refine the initial route classifier so stronger TiKV-specific keywords override weaker generic latency wording when they co-occur in the same ticket-visible inputs.',
            evidence_case_ids=case_ids,
            merge_readiness=merge_readiness,
            mismatch_types=[mismatch_type],
        )
    if mismatch_type == 'wrong_path':
        return ImprovementInput(
            bucket=bucket,
            title='Refine path split after route selection',
            problem='Replay recognized the rough problem family, but split into the wrong TiKV path.',
            proposed_change='Improve path-specific cue mapping so store-local, network, snapshot, and generic write-path symptoms are separated more reliably before deeper diagnosis starts.',
            evidence_case_ids=case_ids,
            merge_readiness=merge_readiness,
            mismatch_types=[mismatch_type],
        )
    if mismatch_type == 'wrong_scope':
        return ImprovementInput(
            bucket=bucket,
            title='Improve scope classification for cluster-wide vs subset issues',
            problem='Replay inferred the wrong scope compared with the eventual human conclusion.',
            proposed_change='Add stronger scope cues for single-store, hotspot, and subset incidents so the skill distinguishes them from cluster-wide slowness earlier.',
            evidence_case_ids=case_ids,
            merge_readiness=merge_readiness,
            mismatch_types=[mismatch_type],
        )
    if mismatch_type == 'missing_required_evidence':
        return ImprovementInput(
            bucket=bucket,
            title='Standardize missing-evidence blockers in evaluator output',
            problem='Replay did not make required evidence blockers explicit enough for later review.',
            proposed_change='Normalize missing evidence into explicit blocker reasons so reviewers can tell whether the gap is per-store, per-hop, lower-layer, or monitoring-window related.',
            evidence_case_ids=case_ids,
            merge_readiness=merge_readiness,
            mismatch_types=[mismatch_type],
        )
    return None


def aggregate_case_review_improvements(case_reviews: list[dict]) -> dict[str, list[dict]]:
    grouped: dict[str, list[str]] = defaultdict(list)

    for review in case_reviews:
        case_id = str(review.get('incident_id') or '').strip()
        mismatch_types = [str(item or '').strip() for item in (review.get('mismatch_type') or []) if str(item or '').strip()]
        for mismatch_type in mismatch_types:
            if case_id and case_id not in grouped[mismatch_type]:
                grouped[mismatch_type].append(case_id)

    buckets = {
        'workflow_improvements': [],
        'knowledge_improvements': [],
        'output_improvements': [],
        'do_not_change_yet': [],
    }

    for mismatch_type, case_ids in grouped.items():
        item = _build_mismatch_improvement(mismatch_type, case_ids)
        if not item:
            continue
        if item.bucket == 'workflow':
            buckets['workflow_improvements'].append(item.to_dict())
        elif item.bucket == 'knowledge':
            buckets['knowledge_improvements'].append(item.to_dict())
        elif item.bucket == 'output':
            buckets['output_improvements'].append(item.to_dict())
        else:
            buckets['do_not_change_yet'].append(item.to_dict())

    return buckets