"""
correction_engine.py
Centralized correction and policy-evidence extraction module
for the CXR Reasoning Kernel (v3.2+)

Integrated into v3.1 architecture per user requirements
"""

import re
from typing import List, Dict

# ------------------------------------------------------------
# 1. Core Correction Engine
# ------------------------------------------------------------
class CorrectionEngine:
    def __init__(self):
        # keyword templates and associated actions (ordered by priority)
        self.patterns = [
            (r"modifier\s*59", self._add_modifier_59),
            (r"modifier\s*(xe|xp|xs|xu)", self._add_x_modifiers),
            (r"separately\s+(identify|report|bill)", self._add_modifier_59),
            (r"separate\s+(procedure|service)", self._add_modifier_59),
            (r"distinct\s+(procedure|service)", self._add_modifier_59),
            (r"not\s+separately\s+billable", self._remove_conflicting_cpt),
            (r"bundled\s+(into|with)", self._remove_conflicting_cpt),
            (r"component\s+of", self._remove_conflicting_cpt),
            (r"covered\s+when", self._add_covered_diagnosis),
            (r"documentation\s+(required|must|should)", self._flag_documentation_review),
            (r"(requires|require)\s+prior\s+authorization", self._flag_prior_auth),
            (r"(additional|supplemental)\s+service", self._add_modifier_for_additional)
        ]

    # --------------------------------------------------------
    # 2. Main correction dispatcher
    # --------------------------------------------------------
    def generate(self, docs: List[Dict], decision: str, claim: Dict = None) -> List[Dict]:
        """
        Scans retrieved Qdrant docs for policy patterns.
        Returns structured correction suggestions.
        
        Args:
            docs: Semantic docs from Qdrant
            decision: PASS/FAIL/REVIEW
            claim: Original claim (for context)
        """
        if decision not in ("FAIL", "REVIEW", "ESCALATE"):
            return []

        corrections = []
        for d in docs:
            text = d.get("text", "").lower()
            score = d.get("score", 0.0)
            
            for pattern, handler in self.patterns:
                if re.search(pattern, text):
                    result = handler(text, score, d)
                    if result:
                        corrections.append(result)
                        break  # One correction per doc

        # remove duplicates and keep top 5 by confidence
        seen = set()
        unique = []
        for c in sorted(corrections, key=lambda x: x["confidence"], reverse=True):
            if c["action"] not in seen:
                seen.add(c["action"])
                unique.append(c)
        return unique[:5]

    # --------------------------------------------------------
    # 3. Pattern Handlers
    # --------------------------------------------------------
    def _add_modifier_59(self, text, score, doc):
        return {
            "action": "Add modifier 59 for distinct diagnostic service",
            "code": "59",
            "confidence": score,
            "rationale": text[:250],
            "source": doc.get("source", "Policy")
        }

    def _add_x_modifiers(self, text, score, doc):
        match = re.search(r"modifier\s*(xe|xp|xs|xu)", text)
        mod = match.group(1).upper() if match else "XE"
        return {
            "action": f"Add modifier {mod} for distinct procedural context",
            "code": mod,
            "confidence": score,
            "rationale": text[:250],
            "source": doc.get("source", "Policy")
        }

    def _remove_conflicting_cpt(self, text, score, doc):
        return {
            "action": "Remove conflicting CPT — not separately billable",
            "code": None,
            "confidence": score,
            "rationale": text[:250],
            "source": doc.get("source", "Policy")
        }

    def _add_covered_diagnosis(self, text, score, doc):
        # try to extract ICD pattern if present
        icd_match = re.search(r"([A-Z]\d{2}(\.\d+)?)", text)
        icd_code = icd_match.group(1) if icd_match else None
        return {
            "action": "Add covered ICD-10 diagnosis from policy list",
            "code": icd_code,
            "confidence": score,
            "rationale": text[:250],
            "source": doc.get("source", "Policy")
        }

    def _flag_documentation_review(self, text, score, doc):
        return {
            "action": "Flag for manual review — documentation required",
            "code": None,
            "confidence": score,
            "rationale": text[:250],
            "source": doc.get("source", "Policy")
        }

    def _flag_prior_auth(self, text, score, doc):
        return {
            "action": "Verify prior authorization requirement",
            "code": None,
            "confidence": score,
            "rationale": text[:250],
            "source": doc.get("source", "Policy")
        }
    
    def _add_modifier_for_additional(self, text, score, doc):
        return {
            "action": "Consider modifier for additional/supplemental service",
            "code": "59",
            "confidence": score,
            "rationale": text[:250],
            "source": doc.get("source", "Policy")
        }

# ------------------------------------------------------------
# 4. Policy Snippet Extractor
# ------------------------------------------------------------
def extract_policy_snippets(docs: List[Dict], top_n: int = 5) -> List[Dict]:
    """
    Extracts short policy excerpts with source metadata.
    
    Args:
        docs: Semantic docs from Qdrant
        top_n: Number of snippets to extract
    """
    snippets = []
    for d in sorted(docs, key=lambda x: x.get("score", 0), reverse=True)[:top_n]:
        source = d.get("source", "Unknown")
        
        # Build richer source reference
        source_parts = [source]
        if d.get('chapter'):
            source_parts.append(f"Ch.{d.get('chapter')}")
        if d.get('section'):
            source_parts.append(f"§{d.get('section')}")
        source_ref = ' '.join(source_parts)
        
        snippets.append({
            "source": source_ref,
            "excerpt": d.get("text", "")[:400],
            "confidence": d.get("score", 0.0),
            "collection": d.get("collection", "unknown"),
            "metadata": {
                "cpt_codes": d.get("cpt_codes", []),
                "modifiers": d.get("modifiers", []),
                "page": d.get("page")
            }
        })
    return snippets

