import re
from datetime import datetime, timezone

import yaml
from typing import Dict, Any, Optional

import re
from typing import Dict, Any, Optional
from langfuse import Langfuse


class SovereignPromptRegistry:
    def __init__(self, langfuse_client: Langfuse):
        self.client = langfuse_client
        self.locked_regex = r"[A-Z]{5}[0-9]{4}[A-Z]{1}"

    # Inside SovereignPromptRegistry
    def _get_raw_template(self, name: str) -> Optional[Any]:
        try:
            # Try raw name first
            return self.client.get_prompt(name, label="production")
        except Exception:
            try:
                # 🎯 Fallback: Try with .md if not present
                if not name.endswith(".md"):
                    return self.client.get_prompt(f"{name}.md", label="production")
                return None
            except Exception:
                return None

    def get_mission_prompt(self,
                           task_key: str,
                           user_id: str,
                           agent_id: str,
                           context_vars: Dict[str, Any],  # 🎯 FACTS passed from Node
                           mission_template_id: Optional[str] = None) -> str:

        # 🎯 THE BASELINE TRUTH: System-wide variables that every prompt can use
        global_context = {
            "org_name": "YantramOps Sovereign",
            "current_date": datetime.now(timezone.utc).strftime("%Y-%m-%d"),
            "system_version": "Sutram_OS_v1.5",
            "user_id": user_id
        }
        merged_facts = {**global_context, **context_vars}


        # --- LAYER 1: The Sovereign Foundation ---
        base_dna = self._get_raw_template("templates/yantram-ops-dna-v1")
        # Compile base DNA (usually static, but good to have context)
        base_dna_str = base_dna.compile(**context_vars) if base_dna else ""

        # --- LAYER 2: The User Identity ---
        persona_path = f"templates/user-{user_id.lower()}"
        user_persona = self._get_raw_template(persona_path)
        user_persona_str = user_persona.compile(**context_vars) if user_persona else ""

        # --- LAYER 3 & 4: TASK RESOLUTION & NATIVE COMPILE ---
        # Find the most specific template available
        task_prompt_obj = (
                self._get_raw_template(f"templates/{mission_template_id}/{task_key}") or
                self._get_raw_template(f"agents/{agent_id}/{task_key}") or
                self._get_raw_template(f"core/{task_key}") or
                # 🛡️ THE FALLBACK: For your specific Langfuse slug naming convention
                self._get_raw_template(f"core/{task_key}-v1") or
                self._get_raw_template(f"{task_key}.md")
        )

        if not task_prompt_obj:
            return f"Error: No template found for {task_key}"

        # 🚀 THE MAGIC: Use Langfuse native compile
        # This replaces {{proposal_data}}, {{user_department}}, etc.
        compiled_task_str = task_prompt_obj.compile(**merged_facts)
        domain_enforcement = self._resolve_dharma_enforcement(agent_id)

        # --- LAYER 5: THE DEEP MERGE & LOCKED DHARMA ---
        return f"""
{base_dna_str}

{user_persona_str}

# MISSION TASK: {task_key.upper()}
{compiled_task_str}

# FINAL DHARMA ENFORCEMENT (LOCKED)
{domain_enforcement}
"""

    def _resolve_dharma_enforcement(self, agent_id: str) -> str:
        """Determines the 'Hard Law' based on the agent's domain."""
        if "lead_gen" in agent_id:
            return """
- NO SPAM: Do not generate outreach that lacks 3 specific technical points from the prospect's profile.
- CONSENT: Do not store personal data for individuals who have explicitly opted out.
- TRUTH: Do not hallucinate case studies or features YantramOps does not physically possess.
"""
        # Default/Finance Dharma
        return f"- PII MASKING: All PAN numbers matching {self.locked_regex} MUST be masked."

    def verify_enforcement(self, final_output: str) -> bool:
        """
        The Narasimha Post-Process Check.
        Final verification that no raw PAN numbers escaped the masking policy.
        """
        if re.search(self.locked_regex, final_output):
            # In a real scenario, this would trigger an alert in Langfuse
            print("🚨 CRITICAL BREACH: Raw PII (PAN) detected in agent output!")
            return False
        return True