"""LLM-powered relationship extraction."""

from __future__ import annotations

import json
import logging
from typing import Any

from neurovault.cognitive.llm_client import OllamaClient

logger = logging.getLogger(__name__)

RELATION_SYSTEM_PROMPT = """Extract entity relationships from text.
Return a JSON array of objects with keys: subject, predicate, object, confidence.
Only extract relationships explicitly stated or strongly implied.
Use lowercase for predicates. Keep entity names as they appear.
Return [] if no clear relationships exist."""


class LLMRelationshipExtractor:
    """Extract typed relationship triples using local LLM."""

    def __init__(self, client: OllamaClient | None = None):
        self._client = client or OllamaClient()

    def extract(self, text: str) -> list[dict[str, Any]]:
        """Extract relationship triples from text."""
        if not self._client.is_available() or len(text) < 10:
            return []

        prompt = (
            "Extract all entity relationships from this text as JSON:\n\n"
            f"Text: \"{text[:500]}\"\n\n"
            "Return JSON array (example format):\n"
            "[{\"subject\": \"Sarah\", \"predicate\": \"manages\", \"object\": \"Project Alpha\", \"confidence\": 0.9}]\n\n"
            "Relationships:"
        )

        raw = self._client.generate(
            prompt=prompt,
            system=RELATION_SYSTEM_PROMPT,
            temperature=0.1,
            max_tokens=512,
        )
        return self._parse_response(raw)

    def _parse_response(self, raw: str) -> list[dict[str, Any]]:
        """Parse LLM JSON response into relationship triples."""
        if not raw:
            return []

        try:
            start = raw.find("[")
            end = raw.rfind("]") + 1
            if start < 0 or end <= start:
                return []
            parsed = json.loads(raw[start:end])
            if not isinstance(parsed, list):
                return []

            results: list[dict[str, Any]] = []
            for item in parsed:
                if not isinstance(item, dict):
                    continue
                if all(k in item for k in ("subject", "predicate", "object")):
                    results.append(
                        {
                            "subject": str(item["subject"]),
                            "predicate": str(item["predicate"]),
                            "object": str(item["object"]),
                            "confidence": float(item.get("confidence", 0.7)),
                        }
                    )
            return results
        except (json.JSONDecodeError, ValueError, TypeError) as exc:
            logger.debug("Failed to parse LLM relationship response: %s", exc)
            return []
