"""
Knowledge Graph Builder Utility for LangXchange.
Optimized for speed with parallel processing and batching.
"""

import json
import logging
import asyncio
import re
from typing import List, Dict, Any, Optional, Callable
from dataclasses import dataclass, asdict, field
from functools import lru_cache
import hashlib
from .exceptions import InsufficientQuotaError

logger = logging.getLogger(__name__)


@dataclass
class Triple:
    """Represents a (subject, predicate, object) relation with optional metadata."""
    subject: str
    predicate: str
    object: str
    subject_type: Optional[str] = None
    object_type: Optional[str] = None
    confidence: float = 1.0
    metadata: Dict[str, Any] = field(default_factory=dict)  # Fix mutable default

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)
    
    def __hash__(self):
        return hash((self.subject, self.predicate, self.object))
    
    def __eq__(self, other):
        if not isinstance(other, Triple):
            return False
        return (self.subject, self.predicate, self.object) == (other.subject, other.predicate, other.object)


class KnowledgeGraphBuilder:
    """
    Utility for building knowledge graphs from unstructured text using LLMs.
    Optimized with parallel processing, batching, and caching.
    """

    DEFAULT_EXTRACTION_PROMPT = """Extract entities and relationships as subject-predicate-object triples.
Rules:
1. Subjects/Objects = specific entities (people, places, concepts).
2. Predicates = verbs/phrases describing relationships.
3. Include types (Person, Organization, Concept, etc.).
4. Return ONLY a JSON array with keys: subject, predicate, object, subject_type, object_type, confidence.
5. Return [] if no relationships found.

Text:
{text}"""

    def __init__(
        self, 
        llm_helper: Any,
        max_concurrency: int = 5,
        batch_size: int = 3,
        enable_cache: bool = True,
        timeout: float = 120.0
    ):
        """
        Initialize with an LLM helper.
        
        Args:
            llm_helper: Object with .chat() method (sync) or .achat() method (async).
            max_concurrency: Max concurrent LLM calls.
            batch_size: Number of texts to batch in single prompt (reduces API calls).
            enable_cache: Cache results for identical text inputs.
            timeout: Timeout per LLM call in seconds.
        """
        self.llm = llm_helper
        self.semaphore = asyncio.Semaphore(max_concurrency)
        self.batch_size = batch_size
        self.enable_cache = enable_cache
        self.timeout = timeout
        self._cache: Dict[str, List[Triple]] = {}

    def _cache_key(self, text: str) -> str:
        return hashlib.md5(text.encode()).hexdigest()

    def _parse_json_response(self, content: str) -> List[Dict]:
        """Robust JSON extraction from LLM response."""
        content = content.strip()
        
        # Try direct parse first
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            pass
        
        # Extract JSON from markdown code blocks
        patterns = [
            r'```json\s*([\s\S]*?)\s*```',
            r'```\s*([\s\S]*?)\s*```',
            r'\[[\s\S]*\]'  # Find array directly
        ]
        
        for pattern in patterns:
            match = re.search(pattern, content)
            if match:
                try:
                    return json.loads(match.group(1) if '```' in pattern else match.group(0))
                except (json.JSONDecodeError, IndexError):
                    continue
        
        return []

    def _build_triples(self, data: List[Dict]) -> List[Triple]:
        """Convert raw dicts to Triple objects with validation."""
        triples = []
        for item in data:
            if not isinstance(item, dict):
                continue
            subject = str(item.get("subject", "")).strip()
            predicate = str(item.get("predicate", "")).strip()
            obj = str(item.get("object", "")).strip()
            
            if not all([subject, predicate, obj]):
                continue
                
            triples.append(Triple(
                subject=subject,
                predicate=predicate,
                object=obj,
                subject_type=item.get("subject_type"),
                object_type=item.get("object_type"),
                confidence=min(1.0, max(0.0, float(item.get("confidence", 1.0)))),
                metadata=item.get("metadata") or {}
            ))
        return triples

    async def _call_llm(self, messages: List[Dict]) -> str:
        """Call LLM with async support and timeout."""
        # Check if LLM has async method
        if hasattr(self.llm, 'achat'):
            return await asyncio.wait_for(
                self.llm.achat(messages=messages, temperature=0),
                timeout=self.timeout
            )
        else:
            # Run sync method in executor
            loop = asyncio.get_event_loop()
            return await asyncio.wait_for(
                loop.run_in_executor(None, lambda: self.llm.chat(messages=messages, temperature=0)),
                timeout=self.timeout
            )

    async def extract_triples(
        self, 
        text: str, 
        custom_prompt: Optional[str] = None
    ) -> List[Triple]:
        """Extract triples from a single piece of text with caching."""
        if not text or not text.strip():
            return []
        
        # Check cache
        cache_key = self._cache_key(text)
        if self.enable_cache and cache_key in self._cache:
            return self._cache[cache_key]

        prompt_template = custom_prompt or self.DEFAULT_EXTRACTION_PROMPT
        prompt = prompt_template.format(text=text)
        messages = [
            {"role": "system", "content": "Extract structured triplets for Knowledge Graphs. Return only valid JSON."},
            {"role": "user", "content": prompt}
        ]

        async with self.semaphore:
            try:
                response = await self._call_llm(messages)
                logger.info(f"LLM Raw Response: {response}")
                data = self._parse_json_response(response)
                triples = self._build_triples(data)
                
                if self.enable_cache:
                    self._cache[cache_key] = triples
                return triples
                
            except asyncio.TimeoutError:
                logger.warning(f"LLM call timed out after {self.timeout}s")
                return []
            except InsufficientQuotaError:
                raise
            except Exception as e:
                logger.error(f"Error extracting triples: {e}")
                return []

    async def _extract_batch(self, texts: List[str]) -> List[Triple]:
        """Extract triples from multiple texts in a single LLM call."""
        if not texts:
            return []
        
        # Build batched prompt
        batch_prompt = "Extract triples from each numbered text section below. Return a single JSON array combining all triples.\n\n"
        for i, text in enumerate(texts, 1):
            batch_prompt += f"[TEXT {i}]\n{text}\n\n"
        
        messages = [
            {"role": "system", "content": "Extract structured triplets for Knowledge Graphs. Return only valid JSON array."},
            {"role": "user", "content": batch_prompt}
        ]

        async with self.semaphore:
            try:
                response = await self._call_llm(messages)
                data = self._parse_json_response(response)
                return self._build_triples(data)
            except asyncio.TimeoutError:
                logger.error(f"Batch extraction timed out after {self.timeout}s")
                # Fallback to individual processing
                results = await asyncio.gather(
                    *[self.extract_triples(t) for t in texts],
                    return_exceptions=True
                )
                return [t for r in results if isinstance(r, list) for t in r]
            except InsufficientQuotaError:
                raise
            except Exception as e:
                logger.error(f"Batch extraction error: {e}")
                # Fallback to individual processing
                results = await asyncio.gather(
                    *[self.extract_triples(t) for t in texts],
                    return_exceptions=True
                )
                return [t for r in results if isinstance(r, list) for t in r]

    async def process_chunks(
        self, 
        chunks: List[str],
        progress_callback: Optional[Callable[[int, int], Any]] = None,
        deduplicate: bool = True,
        use_batching: bool = True
    ) -> List[Triple]:
        """
        Process multiple chunks with parallel execution.
        
        Args:
            chunks: List of text chunks.
            progress_callback: Called with (completed, total) counts.
            deduplicate: Remove duplicate triples.
            use_batching: Batch multiple chunks per LLM call (faster, slightly less accurate).
        """
        if not chunks:
            return []
        
        total = len(chunks)
        completed = 0
        all_triples: List[Triple] = []

        if use_batching and self.batch_size > 1:
            # Process in batches
            batches = [chunks[i:i + self.batch_size] for i in range(0, len(chunks), self.batch_size)]
            tasks = [self._extract_batch(batch) for batch in batches]
            
            for coro in asyncio.as_completed(tasks):
                try:
                    batch_triples = await coro
                    all_triples.extend(batch_triples)
                except Exception as e:
                    logger.error(f"Batch processing error: {e}")
                
                completed += self.batch_size
                if progress_callback:
                    progress_callback(min(completed, total), total)
        else:
            # Parallel individual processing
            tasks = [self.extract_triples(chunk) for chunk in chunks]
            
            for coro in asyncio.as_completed(tasks):
                try:
                    chunk_triples = await coro
                    all_triples.extend(chunk_triples)
                except Exception as e:
                    logger.error(f"Chunk processing error: {e}")
                
                completed += 1
                if progress_callback:
                    progress_callback(completed, total)

        if deduplicate:
            all_triples = list(set(all_triples))
        
        return all_triples

    def clear_cache(self):
        """Clear the result cache."""
        self._cache.clear()