"""
LLM-powered entity extraction from text.

Extracts people, organizations, locations, and their relationships.
Supports OpenAI, Anthropic, and Ollama providers.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from typing import TYPE_CHECKING

import structlog
from openai import AsyncOpenAI

if TYPE_CHECKING:
    from remembra.config import Settings

log = structlog.get_logger()


# ============================================================================
# Entity Extraction Prompt
# ============================================================================

ENTITY_EXTRACTION_PROMPT = """You are an entity extraction engine. Extract entities and relationships from text.

ENTITY TYPES:
- PERSON: People's names (include titles, roles if mentioned)
- ORG: Companies, organizations, teams, groups
- LOCATION: Cities, countries, addresses, places
- DATE: Specific dates, time periods, deadlines
- PRODUCT: Products, services, projects, software
- MONEY: Monetary amounts, prices, budgets
- CONCEPT: Abstract concepts, topics, skills

RELATIONSHIP TYPES:
- WORKS_AT: Person works at organization
- MANAGES: Person manages another person
- SPOUSE_OF: Married/partner relationship
- LOCATED_IN: Entity is located somewhere
- OWNS: Person/org owns something
- CREATED: Person/org created something
- ROLE: Person has a role (e.g., CEO, Manager)

RULES:
1. Extract ALL named entities, even if mentioned once
2. Infer relationships from context (e.g., "John's wife Lisa" → SPOUSE_OF)
3. Include role/title as part of person description
4. Resolve pronouns when clear (e.g., "He is the CEO" after mentioning John)
5. Return empty arrays if no entities found

OUTPUT FORMAT (strict JSON):
{
  "entities": [
    {
      "name": "John Smith",
      "type": "PERSON",
      "description": "CEO of Acme Corp",
      "aliases": ["John", "Mr. Smith"]
    },
    {
      "name": "Acme Corp",
      "type": "ORG",
      "description": "Company where John works",
      "aliases": ["Acme", "Acme Corporation"]
    }
  ],
  "relationships": [
    {
      "subject": "John Smith",
      "predicate": "WORKS_AT",
      "object": "Acme Corp"
    },
    {
      "subject": "John Smith",
      "predicate": "ROLE",
      "object": "CEO"
    }
  ]
}

EXAMPLES:

Input: "Sarah mentioned that her husband Mike works at Google as a Senior Engineer."
Output: {
  "entities": [
    {"name": "Sarah", "type": "PERSON", "description": "Mentioned Mike", "aliases": []},
    {"name": "Mike", "type": "PERSON", "description": "Senior Engineer at Google, Sarah's husband", "aliases": []},
    {"name": "Google", "type": "ORG", "description": "Tech company", "aliases": []}
  ],
  "relationships": [
    {"subject": "Mike", "predicate": "SPOUSE_OF", "object": "Sarah"},
    {"subject": "Mike", "predicate": "WORKS_AT", "object": "Google"},
    {"subject": "Mike", "predicate": "ROLE", "object": "Senior Engineer"}
  ]
}

Input: "The meeting is scheduled for March 15th in the Denver office."
Output: {
  "entities": [
    {"name": "March 15th", "type": "DATE", "description": "Meeting date", "aliases": []},
    {"name": "Denver office", "type": "LOCATION", "description": "Meeting location", "aliases": ["Denver"]}
  ],
  "relationships": []
}
"""


# ============================================================================
# Data Classes
# ============================================================================

@dataclass
class ExtractedEntity:
    """An entity extracted from text."""
    name: str
    type: str
    description: str
    aliases: list[str]


@dataclass
class ExtractedRelationship:
    """A relationship extracted from text."""
    subject: str
    predicate: str
    object: str


@dataclass
class ExtractionResult:
    """Result of entity extraction."""
    entities: list[ExtractedEntity]
    relationships: list[ExtractedRelationship]


# ============================================================================
# Entity Extractor
# ============================================================================

class EntityExtractor:
    """
    Extracts entities and relationships from text using LLM.
    
    Usage:
        extractor = EntityExtractor()
        result = await extractor.extract("John is the CEO of Acme Corp")
        # result.entities = [Entity(name="John", type="PERSON", ...), ...]
        # result.relationships = [Relationship(subject="John", predicate="WORKS_AT", ...)]
    """
    
    def __init__(
        self,
        model: str = "gpt-4o-mini",
        api_key: str | None = None,
    ):
        self.model = model
        self.api_key = api_key
        self._client: AsyncOpenAI | None = None
    
    def _get_client(self) -> AsyncOpenAI:
        """Get or create OpenAI client."""
        if self._client is None:
            self._client = AsyncOpenAI(api_key=self.api_key)
            log.info("entity_extractor_initialized", model=self.model)
        return self._client
    
    async def extract(self, content: str) -> ExtractionResult:
        """
        Extract entities and relationships from content.
        
        Args:
            content: Text to extract entities from
            
        Returns:
            ExtractionResult with entities and relationships
        """
        if not content.strip() or len(content.strip()) < 10:
            return ExtractionResult(entities=[], relationships=[])
        
        try:
            client = self._get_client()
            
            log.debug("extracting_entities", content_length=len(content))
            
            response = await client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": ENTITY_EXTRACTION_PROMPT},
                    {"role": "user", "content": f"Extract entities and relationships from:\n\n{content}"},
                ],
                temperature=0.1,
                response_format={"type": "json_object"},
                timeout=30.0,
            )
            
            result_text = response.choices[0].message.content
            if not result_text:
                return ExtractionResult(entities=[], relationships=[])
            
            data = json.loads(result_text)
            
            # Parse entities
            entities = []
            for e in data.get("entities", []):
                if isinstance(e, dict) and e.get("name"):
                    entities.append(ExtractedEntity(
                        name=e.get("name", ""),
                        type=e.get("type", "CONCEPT"),
                        description=e.get("description", ""),
                        aliases=e.get("aliases", []),
                    ))
            
            # Parse relationships
            relationships = []
            for r in data.get("relationships", []):
                if isinstance(r, dict) and r.get("subject") and r.get("object"):
                    relationships.append(ExtractedRelationship(
                        subject=r.get("subject", ""),
                        predicate=r.get("predicate", "RELATED_TO"),
                        object=r.get("object", ""),
                    ))
            
            log.info(
                "entities_extracted",
                entity_count=len(entities),
                relationship_count=len(relationships),
            )
            
            return ExtractionResult(entities=entities, relationships=relationships)
            
        except json.JSONDecodeError as e:
            log.error("entity_extraction_json_error", error=str(e))
            return ExtractionResult(entities=[], relationships=[])
        except Exception as e:
            log.error("entity_extraction_error", error=str(e))
            return ExtractionResult(entities=[], relationships=[])


# ============================================================================
# Shared JSON Parsing Helper
# ============================================================================

def _parse_extraction_json(raw_text: str) -> ExtractionResult:
    """
    Parse an LLM response into an ExtractionResult.

    Handles both clean JSON and JSON wrapped in markdown fences.
    Returns an empty result on any parse failure.
    """
    # Strip markdown code fences (```json ... ```) if present
    cleaned = raw_text.strip()
    fence_match = re.search(r"```(?:json)?\s*\n?(.*?)```", cleaned, re.DOTALL)
    if fence_match:
        cleaned = fence_match.group(1).strip()

    data = json.loads(cleaned)

    # Parse entities
    entities: list[ExtractedEntity] = []
    for e in data.get("entities", []):
        if isinstance(e, dict) and e.get("name"):
            entities.append(ExtractedEntity(
                name=e.get("name", ""),
                type=e.get("type", "CONCEPT"),
                description=e.get("description", ""),
                aliases=e.get("aliases", []),
            ))

    # Parse relationships
    relationships: list[ExtractedRelationship] = []
    for r in data.get("relationships", []):
        if isinstance(r, dict) and r.get("subject") and r.get("object"):
            relationships.append(ExtractedRelationship(
                subject=r.get("subject", ""),
                predicate=r.get("predicate", "RELATED_TO"),
                object=r.get("object", ""),
            ))

    return ExtractionResult(entities=entities, relationships=relationships)


# ============================================================================
# Anthropic Entity Extractor
# ============================================================================

class AnthropicEntityExtractor:
    """Entity extraction using Anthropic Claude."""

    def __init__(self, model: str = "claude-sonnet-4-5", api_key: str | None = None):
        import anthropic

        self._client = anthropic.AsyncAnthropic(api_key=api_key)
        self.model = model
        log.info("anthropic_entity_extractor_initialized", model=self.model)

    async def extract(self, content: str) -> ExtractionResult:
        """
        Extract entities and relationships from content.

        Args:
            content: Text to extract entities from

        Returns:
            ExtractionResult with entities and relationships
        """
        if not content.strip() or len(content.strip()) < 10:
            return ExtractionResult(entities=[], relationships=[])

        try:
            log.debug("extracting_entities", provider="anthropic", content_length=len(content))

            response = await self._client.messages.create(
                model=self.model,
                max_tokens=4096,
                temperature=0.1,
                system=ENTITY_EXTRACTION_PROMPT,
                messages=[
                    {
                        "role": "user",
                        "content": f"Extract entities and relationships from:\n\n{content}",
                    },
                ],
            )

            # Claude returns content blocks; concatenate text blocks
            result_text = "".join(
                block.text for block in response.content if block.type == "text"
            )
            if not result_text:
                return ExtractionResult(entities=[], relationships=[])

            result = _parse_extraction_json(result_text)

            log.info(
                "entities_extracted",
                provider="anthropic",
                entity_count=len(result.entities),
                relationship_count=len(result.relationships),
            )
            return result

        except json.JSONDecodeError as e:
            log.error("entity_extraction_json_error", provider="anthropic", error=str(e))
            return ExtractionResult(entities=[], relationships=[])
        except Exception as e:
            log.error("entity_extraction_error", provider="anthropic", error=str(e))
            return ExtractionResult(entities=[], relationships=[])


# ============================================================================
# Ollama Entity Extractor
# ============================================================================

class OllamaEntityExtractor:
    """Entity extraction using local Ollama models."""

    def __init__(
        self,
        model: str = "llama3.1",
        base_url: str = "http://localhost:11434",
    ):
        import httpx

        self._base_url = base_url.rstrip("/")
        self.model = model
        self._client = httpx.AsyncClient(timeout=120.0)
        log.info(
            "ollama_entity_extractor_initialized",
            model=self.model,
            base_url=self._base_url,
        )

    async def extract(self, content: str) -> ExtractionResult:
        """
        Extract entities and relationships from content.

        Args:
            content: Text to extract entities from

        Returns:
            ExtractionResult with entities and relationships
        """
        if not content.strip() or len(content.strip()) < 10:
            return ExtractionResult(entities=[], relationships=[])

        try:
            log.debug("extracting_entities", provider="ollama", content_length=len(content))

            response = await self._client.post(
                f"{self._base_url}/api/chat",
                json={
                    "model": self.model,
                    "messages": [
                        {"role": "system", "content": ENTITY_EXTRACTION_PROMPT},
                        {
                            "role": "user",
                            "content": (
                                f"Extract entities and relationships from:\n\n{content}"
                            ),
                        },
                    ],
                    "format": "json",
                    "stream": False,
                    "options": {
                        "temperature": 0.1,
                    },
                },
            )
            response.raise_for_status()

            data = response.json()
            result_text = data.get("message", {}).get("content", "")
            if not result_text:
                return ExtractionResult(entities=[], relationships=[])

            result = _parse_extraction_json(result_text)

            log.info(
                "entities_extracted",
                provider="ollama",
                entity_count=len(result.entities),
                relationship_count=len(result.relationships),
            )
            return result

        except json.JSONDecodeError as e:
            log.error("entity_extraction_json_error", provider="ollama", error=str(e))
            return ExtractionResult(entities=[], relationships=[])
        except Exception as e:
            log.error("entity_extraction_error", provider="ollama", error=str(e))
            return ExtractionResult(entities=[], relationships=[])


# ============================================================================
# Factory
# ============================================================================

def create_entity_extractor(
    settings: Settings,
) -> EntityExtractor | AnthropicEntityExtractor | OllamaEntityExtractor:
    """Create the appropriate entity extractor based on config.

    Reads ``settings.llm_provider`` to decide which backend to use and
    passes through the relevant model / API-key / URL settings.
    """
    provider = getattr(settings, "llm_provider", "openai").lower()

    if provider == "anthropic":
        return AnthropicEntityExtractor(
            model=getattr(settings, "extraction_model", "claude-sonnet-4-5"),
            api_key=getattr(settings, "anthropic_api_key", None),
        )

    if provider == "ollama":
        return OllamaEntityExtractor(
            model=getattr(settings, "extraction_model", "llama3.1"),
            base_url=getattr(settings, "ollama_url", "http://localhost:11434"),
        )

    # Default: OpenAI
    return EntityExtractor(
        model=getattr(settings, "extraction_model", "gpt-4o-mini"),
        api_key=getattr(settings, "openai_api_key", None),
    )


# ============================================================================
# Convenience function
# ============================================================================

async def extract_entities(
    content: str,
    model: str = "gpt-4o-mini",
) -> ExtractionResult:
    """
    Extract entities and relationships from content.

    Convenience function for one-off extraction.
    """
    extractor = EntityExtractor(model=model)
    return await extractor.extract(content)
