"""
llm_extractor.py
----------------
LLM-powered entity and relation extraction for HIEROMEM.

Automatically extracts knowledge graph structure from text using LLMs,
enabling HIEROMEM to work with unstructured data.

Author: Antigravity
"""

from __future__ import annotations

import json
import hashlib
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional, Tuple
import networkx as nx


@dataclass
class ExtractionResult:
    """Result of LLM extraction."""
    entities: List[str]
    relations: List[Tuple[str, str, str]]  # (subject, relation, object)
    graph: nx.DiGraph
    raw_response: str
    cached: bool = False


class LLMProvider(ABC):
    """Base class for LLM providers."""

    @abstractmethod
    def complete(self, prompt: str, **kwargs) -> str:
        """Generate completion for prompt."""
        pass

    @property
    @abstractmethod
    def name(self) -> str:
        """Provider name."""
        pass


class OpenAIProvider(LLMProvider):
    """OpenAI API provider."""

    def __init__(
        self,
        api_key: Optional[str] = None,
        model: str = "gpt-3.5-turbo",
        temperature: float = 0.0,
    ):
        self.model = model
        self.temperature = temperature
        self._client = None
        self._api_key = api_key

    @property
    def client(self):
        if self._client is None:
            try:
                from openai import OpenAI
                import os
                api_key = self._api_key or os.environ.get("OPENAI_API_KEY")
                if not api_key:
                    raise ValueError("OPENAI_API_KEY not set")
                self._client = OpenAI(api_key=api_key)
            except ImportError:
                raise ImportError("openai package not installed. Run: pip install openai")
        return self._client

    @property
    def name(self) -> str:
        return f"openai/{self.model}"

    def complete(self, prompt: str, **kwargs) -> str:
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=self.temperature,
            **kwargs
        )
        return response.choices[0].message.content


class AnthropicProvider(LLMProvider):
    """Anthropic Claude API provider."""

    def __init__(
        self,
        api_key: Optional[str] = None,
        model: str = "claude-3-haiku-20240307",
        temperature: float = 0.0,
    ):
        self.model = model
        self.temperature = temperature
        self._client = None
        self._api_key = api_key

    @property
    def client(self):
        if self._client is None:
            try:
                from anthropic import Anthropic
                import os
                api_key = self._api_key or os.environ.get("ANTHROPIC_API_KEY")
                if not api_key:
                    raise ValueError("ANTHROPIC_API_KEY not set")
                self._client = Anthropic(api_key=api_key)
            except ImportError:
                raise ImportError("anthropic package not installed. Run: pip install anthropic")
        return self._client

    @property
    def name(self) -> str:
        return f"anthropic/{self.model}"

    def complete(self, prompt: str, **kwargs) -> str:
        response = self.client.messages.create(
            model=self.model,
            max_tokens=1024,
            messages=[{"role": "user", "content": prompt}],
            temperature=self.temperature,
            **kwargs
        )
        return response.content[0].text


class StubProvider(LLMProvider):
    """Stub provider for testing (no API calls)."""

    @property
    def name(self) -> str:
        return "stub"

    def complete(self, prompt: str, **kwargs) -> str:
        # Extract simple entities from prompt using regex-like logic
        import re
        words = re.findall(r'\b[A-Z][a-z]+\b', prompt)
        entities = list(set(words))[:5]
        return json.dumps({
            "entities": entities,
            "relations": [
                [entities[i], "related_to", entities[i+1]]
                for i in range(len(entities)-1)
            ]
        })


class CallableProvider(LLMProvider):
    """
    Model-agnostic provider using any callable.

    Allows using any LLM by passing a function that takes a prompt
    and returns a string response.

    Example:
        # With LiteLLM
        from litellm import completion
        def my_llm(prompt):
            response = completion(model="gpt-4", messages=[{"role": "user", "content": prompt}])
            return response.choices[0].message.content

        extractor = LLMExtractor(llm_fn=my_llm)

        # With Ollama
        import ollama
        def local_llm(prompt):
            return ollama.chat(model='llama2', messages=[{'role': 'user', 'content': prompt}])['message']['content']

        extractor = LLMExtractor(llm_fn=local_llm)

        # With HuggingFace
        def hf_llm(prompt):
            return my_pipeline(prompt)[0]['generated_text']

        extractor = LLMExtractor(llm_fn=hf_llm)
    """

    def __init__(self, llm_fn: Callable[[str], str], name: str = "custom"):
        """
        Initialize with a callable.

        Args:
            llm_fn: Callable that takes (prompt: str) and returns str
            name: Optional name for this provider
        """
        if not callable(llm_fn):
            raise ValueError("llm_fn must be callable")
        self._llm_fn = llm_fn
        self._name = name

    @property
    def name(self) -> str:
        return self._name

    def complete(self, prompt: str, **kwargs) -> str:
        return self._llm_fn(prompt)


EXTRACTION_PROMPT = """Extract knowledge graph triplets from the following text.
Identify key entities and the relationships between them.

Text: {text}

Rules:
1. Extract as many meaningful relations as possible.
2. Relations must be between extracted entities.
3. Use concise predicates (e.g. "directed_by", "located_in", "is_a").

Return ONLY a JSON object with this structure:
{{
  "entities": ["list", "of", "entities"],
  "relations": [["Subject", "relation_predicate", "Object"], ...]
}}

Example:
Text: "Apple was founded by Steve Jobs in California."
Output: {{"entities": ["Apple", "Steve Jobs", "California"],
"relations": [["Apple", "founded_by", "Steve Jobs"], ["Apple", "located_in", "California"]]}}"""


QUERY_EXPANSION_PROMPT = """Given this question, identify all entities and concepts that would be needed to answer it.
Include both explicit and implicit references.

Question: {question}

Return a JSON object with:
- "entities": explicit entities mentioned in the question
- "implicit": entities/concepts implied but not directly mentioned
- "related_concepts": broader concepts that might help

Example for "Who founded the company that acquired Instagram?":
{{"entities": ["Instagram"], "implicit": ["Facebook", "Meta", "Mark Zuckerberg"],
"related_concepts": ["tech acquisition", "social media"]}}

Only return valid JSON, no other text."""


class LLMExtractor:
    """
    LLM-powered knowledge graph extractor.

    Example:
        # Using built-in providers
        extractor = LLMExtractor(provider="openai")

        # Model-agnostic: pass any LLM function
        extractor = LLMExtractor(llm_fn=my_custom_llm)

        result = extractor.extract("Alice joined Google in 2020")
        print(result.entities)  # ["Alice", "Google", "2020"]
        print(result.relations) # [("Alice", "works_at", "Google"), ...]
    """

    def __init__(
        self,
        provider: Optional[str] = None,
        api_key: Optional[str] = None,
        model: Optional[str] = None,
        cache_extractions: bool = True,
        llm_fn: Optional[callable] = None,
    ):
        """
        Initialize LLM extractor.

        Args:
            provider: One of "openai", "anthropic", "stub" (ignored if llm_fn provided)
            api_key: API key for provider
            model: Model name override
            cache_extractions: Whether to cache extraction results
            llm_fn: Custom LLM function (prompt: str) -> str for model-agnostic usage
        """
        self.cache_extractions = cache_extractions
        self._cache: Dict[str, ExtractionResult] = {}

        # Model-agnostic: use llm_fn if provided
        if llm_fn is not None:
            self._provider = CallableProvider(llm_fn)
        elif provider == "openai" or provider is None:
            self._provider = OpenAIProvider(
                api_key=api_key,
                model=model or "gpt-4o-mini"
            )
        elif provider == "anthropic":
            self._provider = AnthropicProvider(
                api_key=api_key,
                model=model or "claude-3-haiku-20240307"
            )
        elif provider == "stub":
            self._provider = StubProvider()
        else:
            raise ValueError(f"Unknown provider: {provider}")

        # Always log which model is being used for transparency
        print(f"[LLMExtractor] Using model: {self._provider.name}")

    def _cache_key(self, text: str) -> str:
        """Generate cache key for text."""
        return hashlib.md5(text.encode(), usedforsecurity=False).hexdigest()

    def _parse_response(self, response: str) -> Tuple[List[str], List[Tuple[str, str, str]]]:
        """Parse JSON response from LLM."""
        try:
            # Robust JSON extraction
            response = response.strip()

            # Find JSON object directly
            import re
            json_match = re.search(r'(\{.*\})', response, re.DOTALL)
            if json_match:
                response = json_match.group(1)
            else:
                # Try stripping markdown code blocks if regex failed (fallback)
                if response.startswith("```json"):
                    response = response[7:]
                if response.startswith("```"):
                    response = response[3:]
                if response.endswith("```"):
                    response = response[:-3]

            data = json.loads(response.strip())
            entities = data.get("entities", [])
            # Enforce 3-tuple format: (subject, relation, object)
            raw_relations = data.get("relations", [])
            relations = []
            for r in raw_relations:
                if isinstance(r, (list, tuple)) and len(r) >= 3:
                    relations.append((str(r[0]), str(r[1]), str(r[2])))
            return entities, relations
        except json.JSONDecodeError:
            print("[LLMExtractor] JSON Parsing Failed. Attempting Regex Fallback.")
            # Fallback: extract entities and relations via regex
            import re

            # Entities: anything in "quotations" that isn't a key
            # This is sloppy, so we just try to find list items if possible, or just all strings
            entities = re.findall(r'"([^"]+)"', response)
            entities = [e for e in entities if e not in ("entities", "relations")]

            # Relations: ["Subj", "rel", "Obj"]
            # Pattern: [ "String", "String", "String" ] allows for whitespace
            rel_pattern = r'\[\s*"([^"]+)"\s*,\s*"([^"]+)"\s*,\s*"([^"]+)"\s*\]'
            found_rels = re.findall(rel_pattern, response)
            relations = [(a, b, c) for a, b, c in found_rels]

            return entities[:20], relations

    def extract(self, text: str) -> ExtractionResult:
        """
        Extract entities and relations from text.

        Args:
            text: Text to extract from

        Returns:
            ExtractionResult with entities, relations, and graph
        """
        # Check cache
        cache_key = self._cache_key(text)
        if self.cache_extractions and cache_key in self._cache:
            result = self._cache[cache_key]
            result.cached = True
            return result

        # Call LLM
        prompt = EXTRACTION_PROMPT.format(text=text)
        response = self._provider.complete(prompt)

        # Parse response
        entities, relations = self._parse_response(response)

        # Build graph
        G = nx.DiGraph()
        for entity in entities:
            G.add_node(entity.lower())
        for subj, rel, obj in relations:
            G.add_edge(subj.lower(), obj.lower(), relation=rel)

        result = ExtractionResult(
            entities=entities,
            relations=relations,
            graph=G,
            raw_response=response,
            cached=False
        )

        # Cache result
        if self.cache_extractions:
            self._cache[cache_key] = result

        return result

    def expand_query(self, question: str) -> Dict[str, List[str]]:
        """
        Expand a query to include implicit entities and concepts.

        Args:
            question: The question to expand

        Returns:
            Dict with 'entities', 'implicit', 'related_concepts'
        """
        prompt = QUERY_EXPANSION_PROMPT.format(question=question)
        response = self._provider.complete(prompt)

        try:
            response = response.strip()
            if response.startswith("```json"):
                response = response[7:]
            if response.startswith("```"):
                response = response[3:]
            if response.endswith("```"):
                response = response[:-3]

            data = json.loads(response.strip())
            return {
                "entities": data.get("entities", []),
                "implicit": data.get("implicit", []),
                "related_concepts": data.get("related_concepts", [])
            }
        except json.JSONDecodeError:
            return {"entities": [], "implicit": [], "related_concepts": []}

    def build_query_graph(self, question: str) -> nx.DiGraph:
        """
        Build a query graph by expanding the question.

        Args:
            question: The query/question text

        Returns:
            DiGraph with extracted and expanded entities
        """
        expansion = self.expand_query(question)

        G = nx.DiGraph()
        all_entities = (
            expansion.get("entities", []) +
            expansion.get("implicit", [])
        )

        for entity in all_entities:
            G.add_node(entity.lower())

        # Connect entities
        from itertools import combinations
        for e1, e2 in combinations([e.lower() for e in all_entities], 2):
            G.add_edge(e1, e2, relation="query_related")

        return G

    def clear_cache(self) -> int:
        """Clear extraction cache. Returns number of cleared entries."""
        count = len(self._cache)
        self._cache.clear()
        return count
