import json
import logging
from typing import TYPE_CHECKING, List, Tuple

if TYPE_CHECKING:
    from ...mem_agent import MemAgent

logger = logging.getLogger(__name__)


class GraphExtractor:
    """
    Extracts knowledge graph triplets (Entity, Relation, Entity) from text using an LLM.
    """

    def __init__(self, agent: "MemAgent"):
        self.agent = agent

    def extract(self, text: str) -> List[Tuple[str, str, str]]:
        """
        Extract triplets from text.
        Returns list of (Source, Relation, Target).
        """
        prompt = f"""
        Extract knowledge graph triplets from the following conversation.

        GOAL: Identify entities (people, places, interests, facts) and their relations.
        CRITICAL: DO NOT extract conversation metadata like "User says hello"
        or "Assistant replies".
        Only extract factual knowledge about the user or the topics discussed.

        Format: Return ONLY a JSON array of arrays: [["Source", "Relation", "Target"]].
        Example: [["User", "lives_in", "Istanbul"], ["Python", "is_a", "Programming Language"]]

        Text to analyze:
        ---
        {text}
        ---

        JSON Output:
        """

        # Use direct LLM call instead of agent.chat to avoid recursion/infinite loops
        # and to keep extraction logic separate from conversation history
        messages = [{"role": "user", "content": prompt}]
        response = self.agent.llm.chat(messages=messages, temperature=0.1)

        import re

        try:
            # More robust parsing: find all patterns that look like ["Source", "Relation", "Target"]
            triplet_pattern = r'\[\s*"([^"]+)"\s*,\s*"([^"]+)"\s*,\s*"([^"]+)"\s*\]'
            matches = re.findall(triplet_pattern, response)

            if matches:
                return [list(m) for m in matches]

            # Fallback for standard JSON array
            cleaned_response = response.strip()
            if cleaned_response.startswith("```json"):
                cleaned_response = cleaned_response[7:-3].strip()
            elif cleaned_response.startswith("```"):
                cleaned_response = cleaned_response[3:-3].strip()

            triplets = json.loads(cleaned_response)
            valid_triplets = []
            if isinstance(triplets, list):
                for t in triplets:
                    if isinstance(t, list) and len(t) == 3:
                        valid_triplets.append((str(t[0]), str(t[1]), str(t[2])))
            return valid_triplets

        except Exception as e:
            logger.warning(f"Failed to parse graph extraction: {e}. Response was: {response}")
            return []
