"""
Policy-Driven Recommendation Extractor
======================================

Extracts specific, actionable recommendations from policy text using LLM.
Replaces hardcoded templates with policy-driven instructions.

Version: 1.0
Date: November 1, 2025
Author: CXR Labs
"""

import json
from typing import List, Dict, Any

try:
    import ollama
    OLLAMA_AVAILABLE = True
except ImportError:
    OLLAMA_AVAILABLE = False

try:
    import openai
    OPENAI_AVAILABLE = True
except ImportError:
    OPENAI_AVAILABLE = False


class PolicyRecommendationExtractor:
    """
    Extracts specific recommendations from policy documents using LLM
    
    Instead of hardcoded templates like "Add modifier 59 for distinct diagnostic service",
    extracts the actual instruction from the policy text itself.
    """
    
    def __init__(self, model=None, backend="auto"):
        """
        Initialize the extractor
        
        Args:
            model: Model name (auto-selects based on backend)
            backend: "auto", "ollama", or "openai"
        """
        self.backend = backend
        
        # Auto-select backend
        if self.backend == "auto":
            if OLLAMA_AVAILABLE:
                self.backend = "ollama"
                self.model = model or "mistral:instruct"
            elif OPENAI_AVAILABLE:
                self.backend = "openai"
                self.model = model or "gpt-4o-mini"
            else:
                self.backend = None
                self.model = None
        else:
            self.model = model or ("mistral:instruct" if backend == "ollama" else "gpt-4o-mini")
        
        # Initialize clients
        if self.backend == "openai" and OPENAI_AVAILABLE:
            self.client = openai.OpenAI()
        else:
            self.client = None
    
    def extract_recommendations(
        self, 
        policy_docs: List[Dict], 
        claim: Dict,
        archetype: str,
        archetype_metadata: Dict = None,
        max_recommendations: int = 3
    ) -> List[Dict]:
        """
        Extract specific, actionable recommendations from policy documents
        
        Args:
            policy_docs: List of policy documents from Qdrant
            claim: The claim being analyzed
            archetype: Detected archetype (e.g., "PTP_MUTUALLY_EXCLUSIVE")
            max_recommendations: Maximum number of recommendations to return
        
        Returns:
            List of recommendations with:
            - action: Specific action to take (extracted from policy)
            - rationale: Why this action is needed (extracted from policy)
            - policy_reference: Which policy document supports this
            - codes_affected: Which codes this applies to
            - confidence: How confident we are in this recommendation
        """
        if not self.backend or not policy_docs:
            return []
        
        # Extract claim context
        procedures = self._extract_procedures(claim)
        diagnoses = self._extract_diagnoses(claim)
        claim_id = claim.get('CLM_ID', claim.get('claim_id', 'unknown'))
        
        # Build context for LLM
        policy_context = []
        for i, doc in enumerate(policy_docs[:5], 1):  # Top 5 policies
            text = doc.get('text', '')[:500]  # First 500 chars
            source = doc.get('source', 'Unknown')
            score = doc.get('score', doc.get('policy_score', 0))
            
            policy_context.append({
                "number": i,
                "source": source,
                "relevance": f"{score:.0%}",
                "excerpt": text
            })
        
        # Create prompt with archetype context
        prompt = self._build_extraction_prompt(
            claim_id=claim_id,
            procedures=procedures,
            diagnoses=diagnoses,
            archetype=archetype,
            archetype_metadata=archetype_metadata or {},
            policy_context=policy_context,
            max_recommendations=max_recommendations
        )
        
        # Call LLM
        try:
            if self.backend == "ollama":
                response_text = self._call_ollama(prompt)
            elif self.backend == "openai":
                response_text = self._call_openai(prompt)
            else:
                return []
            
            # Parse response
            recommendations = self._parse_llm_response(response_text, policy_docs)
            
            # Filter out recommendations that don't apply to this claim's codes
            filtered_recommendations = self._filter_irrelevant_recommendations(
                recommendations, 
                procedures,
                diagnoses
            )
            
            if not filtered_recommendations and recommendations:
                print(f"[PolicyExtractor] Filtered out {len(recommendations)} irrelevant recommendations")
            
            return filtered_recommendations[:max_recommendations]
            
        except Exception as e:
            print(f"[PolicyExtractor] Error: {e}")
            return []
    
    def _build_extraction_prompt(
        self,
        claim_id: str,
        procedures: List[str],
        diagnoses: List[str],
        archetype: str,
        archetype_metadata: Dict,
        policy_context: List[Dict],
        max_recommendations: int
    ) -> str:
        """Build the prompt for LLM"""
        
        # Build archetype context from database detection logic
        archetype_context = ""
        if archetype_metadata and archetype_metadata.get('problematic_codes'):
            # We have codes from the detection logic
            prob_codes = archetype_metadata['problematic_codes']
            # Find which codes from detection are in this claim
            claim_probs = [c for c in procedures if c in prob_codes]
            
            if len(claim_probs) >= 2 and 'PTP' in archetype:
                archetype_context = f"\n\nARCHETYPE DETECTION RESULT:"
                archetype_context += f"\n- Database detected {archetype.replace('_', ' ')}"
                archetype_context += f"\n- Conflicting codes in THIS claim: {', '.join(claim_probs)}"
                archetype_context += f"\n- These codes were flagged as mutually exclusive by database detection logic"
                archetype_context += f"\n- Database recommendation: One of these codes should be removed"
            elif claim_probs and 'MUE' in archetype:
                archetype_context = f"\n\nARCHETYPE DETECTION RESULT:"
                archetype_context += f"\n- Database detected MUE violation for: {', '.join(claim_probs)}"
                archetype_context += f"\n- Database recommendation: Review units for these codes"
        
        prompt = f"""You are a medical coding expert analyzing a Medicare claim.

CLAIM CONTEXT:
- Claim ID: {claim_id}
- Procedures: {', '.join(procedures) if procedures else 'None'}
- Diagnoses: {', '.join(diagnoses) if diagnoses else 'None'}
- Detected Issue: {archetype.replace('_', ' ')}{archetype_context}

POLICY DOCUMENTS:
"""
        
        for policy in policy_context:
            prompt += f"\n[Policy {policy['number']}] {policy['source']} (Relevance: {policy['relevance']})\n"
            prompt += f'"{policy["excerpt"]}..."\n'
        
        prompt += f"""

TASK:
Extract up to {max_recommendations} specific, actionable recommendations.

PRIORITY ORDER (follow this strictly):
1. **If ARCHETYPE DETECTION shows specific conflicting codes:** Use those codes as the basis for your recommendation
2. **If policy text mentions those same codes:** Use policy language for the rationale
3. **If policy text mentions DIFFERENT codes:** Ignore the policy text, base recommendation solely on archetype detection
4. **If NO archetype detection context:** Only then extract from policy text (and only if it mentions claim codes)

For each recommendation, provide:
1. **action**: Specific instruction (e.g., "Remove procedure code 27447 as it conflicts with 27130 per NCCI PTP edits")
2. **rationale**: Brief reason (from archetype detection + policy if available, max 100 chars)
3. **policy_number**: Which policy document supports this (or 0 if based on archetype detection)
4. **codes_affected**: Exact codes from the claim this applies to (array)
5. **confidence**: 0.95 if from archetype detection, 0.8+ if policy confirms, <0.7 if policy-only

CRITICAL RULES:
- **ARCHETYPE DETECTION IS AUTHORITATIVE** - it's from database analysis of 790K claims
- If archetype detection says "codes X and Y conflict", your recommendation MUST address X and Y specifically
- DO NOT extract recommendations about codes that aren't in the claim
- DO NOT use policy text that mentions different codes than what archetype detection found
- For PTP: Recommend removing the LATER procedure code (e.g., second in the list)
- For MUE: Recommend reducing units for the flagged code
- **NEVER mention diagnosis codes (ICD) in the action field** - only procedure codes (CPT/HCPCS)
- Diagnosis codes are for context only - recommendations can ONLY modify procedure codes
- Keep rationale under 100 characters - be concise
- ONLY return recommendations that mention PROCEDURE codes in the action field

OUTPUT FORMAT (JSON only, no markdown):
{{
  "recommendations": [
    {{
      "action": "Specific action from policy",
      "rationale": "Policy reason extracted from text",
      "policy_number": 1,
      "codes_affected": ["27447", "27130"],
      "confidence": 0.85
    }}
  ]
}}

Respond ONLY with valid JSON. No explanations, no markdown, just the JSON object."""
        
        return prompt
    
    def _call_ollama(self, prompt: str) -> str:
        """Call Ollama local LLM"""
        response = ollama.chat(
            model=self.model,
            messages=[
                {
                    "role": "system",
                    "content": "You are a certified medical coder extracting specific recommendations from Medicare policy documents. Respond ONLY in valid JSON."
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            format="json",
            options={
                "temperature": 0.1,
                "num_predict": 2000
            }
        )
        return response['message']['content']
    
    def _call_openai(self, prompt: str) -> str:
        """Call OpenAI API"""
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {
                    "role": "system",
                    "content": "You are a certified medical coder extracting specific recommendations from Medicare policy documents. Respond ONLY in valid JSON."
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            response_format={"type": "json_object"},
            temperature=0.1,
            max_tokens=2000
        )
        return response.choices[0].message.content
    
    def _parse_llm_response(self, response_text: str, policy_docs: List[Dict]) -> List[Dict]:
        """Parse LLM response into structured recommendations"""
        try:
            data = json.loads(response_text)
            recommendations = data.get('recommendations', [])
            
            # Enrich with source policy info
            for rec in recommendations:
                policy_num_raw = rec.get('policy_number')
                if policy_num_raw is not None and isinstance(policy_num_raw, int):
                    policy_num = policy_num_raw - 1
                    if 0 <= policy_num < len(policy_docs):
                        rec['source'] = policy_docs[policy_num].get('source', 'Policy')
                        rec['score'] = policy_docs[policy_num].get('score', policy_docs[policy_num].get('policy_score', 0))
                    else:
                        rec['source'] = 'Policy'
                        rec['score'] = 0.5
                else:
                    # Default to first policy if not specified
                    rec['source'] = policy_docs[0].get('source', 'Policy') if policy_docs else 'Policy'
                    rec['score'] = policy_docs[0].get('score', 0.5) if policy_docs else 0.5
            
            return recommendations
            
        except json.JSONDecodeError as e:
            print(f"[PolicyExtractor] JSON parse error: {e}")
            print(f"[PolicyExtractor] Response: {response_text[:200]}")
            return []
    
    def _filter_irrelevant_recommendations(
        self, 
        recommendations: List[Dict], 
        procedure_codes: List[str],
        diagnosis_codes: List[str]
    ) -> List[Dict]:
        """
        Filter out recommendations that don't apply to this claim's codes
        
        CRITICAL: Only procedure codes (CPT/HCPCS) can be modified.
        Diagnosis codes (ICD) should never be modified.
        """
        filtered = []
        
        import re
        
        # Build sets for checking
        proc_set = set(str(c).upper().strip() for c in procedure_codes if c)
        diag_set = set(str(c).upper().strip() for c in diagnosis_codes if c)
        all_codes_set = proc_set | diag_set
        
        for rec in recommendations:
            action = rec.get('action', '').lower()
            rationale = rec.get('rationale', '').lower()
            codes_affected = rec.get('codes_affected', [])
            
            # CRITICAL: Only PROCEDURE codes can be modified
            # If recommendation tries to modify diagnosis codes, REJECT
            for code in codes_affected:
                code_str = str(code).upper().strip()
                
                # Check if this code is a diagnosis code
                if code_str in diag_set:
                    # Trying to modify a diagnosis code - REJECT
                    print(f"[PolicyExtractor] Filtered out: Cannot modify diagnosis code {code_str}")
                    break
            else:
                # No diagnosis codes in codes_affected - continue checking
                
                # Extract all codes mentioned in action text (including alphanumeric like I25.10)
                action_codes_numeric = set(re.findall(r'\b(\d{4,5})\b', action))
                rationale_codes_numeric = set(re.findall(r'\b(\d{4,5})\b', rationale))
                
                # Also check for ICD-10 format codes (letter + numbers + optional decimal)
                action_codes_alpha = set(re.findall(r'\b([A-Z]\d{1,3}(?:\.\d{1,2})?)\b', action.upper()))
                rationale_codes_alpha = set(re.findall(r'\b([A-Z]\d{1,3}(?:\.\d{1,2})?)\b', rationale.upper()))
                
                mentioned_codes = action_codes_numeric | rationale_codes_numeric
                mentioned_alpha_codes = action_codes_alpha | rationale_codes_alpha
                
                # Check if diagnosis codes are mentioned (both numeric and alpha formats)
                diag_set_upper = set(str(c).upper() for c in diag_set)
                diag_mentioned = (mentioned_codes & diag_set) | (mentioned_alpha_codes & diag_set_upper)
                
                if diag_mentioned:
                    # LLM mentioned diagnosis codes in action - REJECT
                    print(f"[PolicyExtractor] Filtered out: Action mentions diagnosis codes: {diag_mentioned}")
                    continue
                
                # Check for foreign codes (not in claim at all)
                foreign_codes = mentioned_codes - all_codes_set
                if foreign_codes:
                    # LLM mentioned codes not in the claim - REJECT
                    print(f"[PolicyExtractor] Filtered out: Action mentions codes not in claim: {foreign_codes}")
                    continue
                
                # Recommendation passed all checks
                filtered.append(rec)
                continue
            
            # If we broke from the inner loop, we skipped this recommendation
            continue
        
        return filtered
    
    def _extract_procedures(self, claim: Dict) -> List[str]:
        """Extract procedure codes from claim"""
        procedures = []
        
        if 'procedure_codes' in claim:
            pc = claim['procedure_codes']
            if isinstance(pc, dict):
                procedures = [str(v) for k, v in pc.items() if v and str(v) != 'None']
            elif isinstance(pc, list):
                procedures = [str(p) for p in pc if p and str(p) != 'None']
        
        return procedures
    
    def _extract_diagnoses(self, claim: Dict) -> List[str]:
        """Extract diagnosis codes from claim"""
        diagnoses = []
        
        if 'diagnosis_codes' in claim:
            dc = claim['diagnosis_codes']
            if isinstance(dc, dict):
                diagnoses = [str(v) for k, v in dc.items() if v and str(v) != 'None']
            elif isinstance(dc, list):
                diagnoses = [str(d) for d in dc if d and str(d) != 'None']
        
        return diagnoses

