"""Causal relationship extraction for TCMGM."""
import json
import logging
from datetime import datetime
from typing import Dict, List

from outhad_contextkit.memory.temporal.enums import CausalType
from outhad_contextkit.memory.temporal.types import CausalLink

logger = logging.getLogger(__name__)


CAUSAL_EXTRACTION_PROMPT = """You are an expert at identifying causal relationships between events.

Given a sequence of events, identify cause-effect relationships.

Events:
{events}

For each causal relationship, provide:
1. cause_event_id: ID of the causing event
2. effect_event_id: ID of the effect event
3. causal_type: Type of causality (caused_by, leads_to, enables, prevents, correlates_with)
4. confidence: Confidence score (0-1)
5. evidence: Brief explanation

Return JSON format:
{{
  "causal_links": [
    {{
      "cause_id": "event1",
      "effect_id": "event2",
      "causal_type": "caused_by",
      "confidence": 0.9,
      "evidence": "Event1 directly triggered Event2"
    }}
  ]
}}
"""


class CausalExtractor:
    """Extracts causal relationships between events using LLM or rule-based methods."""
    
    def __init__(self, llm=None):
        """
        Initialize the causal extractor.
        
        Args:
            llm: LLM instance for causal extraction (optional, required for LLM-based extraction)
        """
        self.llm = llm
    
    def extract_causal_links(
        self,
        events: List[Dict],
        use_llm: bool = True
    ) -> List[CausalLink]:
        """
        Extract causal links between events.
        
        Args:
            events: List of event dictionaries with keys: id, timestamp, content
            use_llm: Use LLM for extraction (vs rule-based). Requires self.llm to be set.
        
        Returns:
            List of CausalLink objects
        """
        if not events:
            logger.warning("No events provided for causal extraction")
            return []
        
        if use_llm:
            if not self.llm:
                logger.warning("LLM not available, falling back to rule-based extraction")
                return self._extract_with_rules(events)
            return self._extract_with_llm(events)
        else:
            return self._extract_with_rules(events)
    
    def _extract_with_llm(self, events: List[Dict]) -> List[CausalLink]:
        """
        Extract causal links using LLM.
        
        Args:
            events: List of event dictionaries
            
        Returns:
            List of CausalLink objects
        """
        # Format events for prompt
        events_text = "\n".join([
            f"ID: {e.get('id', 'unknown')}, Time: {e.get('timestamp', 'unknown')}, Content: {e.get('content', '')}"
            for e in events
        ])
        
        prompt = CAUSAL_EXTRACTION_PROMPT.format(events=events_text)
        
        try:
            response = self.llm.generate_response(
                messages=[{"role": "user", "content": prompt}],
                response_format={"type": "json_object"}
            )
            
            result = json.loads(response)
            
            causal_links = []
            for link_data in result.get("causal_links", []):
                # Validate causal_type
                causal_type = link_data.get("causal_type", "correlates_with")
                if causal_type not in [ct.value for ct in CausalType]:
                    logger.warning(f"Invalid causal type '{causal_type}', defaulting to 'correlates_with'")
                    causal_type = CausalType.CORRELATES_WITH.value
                
                link = CausalLink(
                    cause_id=link_data["cause_id"],
                    effect_id=link_data["effect_id"],
                    causal_type=causal_type,
                    confidence=link_data.get("confidence", 0.7),
                    evidence=link_data.get("evidence"),
                    timestamp=datetime.utcnow()
                )
                causal_links.append(link)
            
            logger.info(f"LLM extracted {len(causal_links)} causal links")
            return causal_links
        
        except Exception as e:
            logger.error(f"LLM causal extraction failed: {e}", exc_info=True)
            return []
    
    def _extract_with_rules(self, events: List[Dict]) -> List[CausalLink]:
        """
        Extract causal links using rule-based heuristics.
        
        This method uses keyword matching and temporal proximity to infer causal relationships.
        
        Args:
            events: List of event dictionaries
            
        Returns:
            List of CausalLink objects
        """
        causal_links = []
        
        # Rule 1: Temporal proximity + causal keywords suggest causality
        causal_keywords = {
            CausalType.CAUSED_BY.value: ["because", "due to", "caused by", "resulted from", "as a result of"],
            CausalType.LEADS_TO.value: ["led to", "resulted in", "caused", "triggered", "resulted in"],
            CausalType.ENABLES.value: ["enabled", "allowed", "made possible", "facilitated", "permitted"],
            CausalType.PREVENTS.value: ["prevented", "stopped", "blocked", "inhibited", "avoided"]
        }
        
        for i, event in enumerate(events):
            content_lower = event.get('content', '').lower()
            event_id = event.get('id', f'event_{i}')
            
            # Check for causal keywords referencing other events
            for j, other_event in enumerate(events):
                if i == j:
                    continue
                
                other_event_id = other_event.get('id', f'event_{j}')
                
                for causal_type, keywords in causal_keywords.items():
                    if any(keyword in content_lower for keyword in keywords):
                        # Determine direction based on causal type
                        if causal_type == CausalType.CAUSED_BY.value:
                            # Current event was caused by other event
                            cause_id = other_event_id
                            effect_id = event_id
                        else:
                            # Current event caused/enabled/prevented other event
                            cause_id = event_id
                            effect_id = other_event_id
                        
                        # Create causal link
                        link = CausalLink(
                            cause_id=cause_id,
                            effect_id=effect_id,
                            causal_type=causal_type,
                            confidence=0.6,  # Lower confidence for rule-based
                            evidence=f"Keyword-based: '{next(kw for kw in keywords if kw in content_lower)}'",
                            timestamp=datetime.utcnow()
                        )
                        causal_links.append(link)
        
        # Deduplicate links (same cause-effect pair)
        seen = set()
        unique_links = []
        for link in causal_links:
            key = (link.cause_id, link.effect_id)
            if key not in seen:
                seen.add(key)
                unique_links.append(link)
        
        logger.info(f"Rule-based extraction found {len(unique_links)} causal links")
        return unique_links

