"""LLM-based entity, relationship, and claim extraction for GraphRAG."""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field

import httpx

from docsctx.config import DocsSettings

# ---------------------------------------------------------------------------
# Data types
# ---------------------------------------------------------------------------

ENTITY_TYPES = {"Person", "Organization", "Concept", "Location", "Event", "Technology", "Other"}
CLAIM_STATUSES = {"CONFIRMED", "REFUTED", "SPECULATIVE"}


@dataclass(slots=True)
class ExtractedEntity:
    name: str
    entity_type: str
    description: str
    score: int = 0


@dataclass(slots=True)
class ExtractedRelationship:
    source: str
    target: str
    predicate: str
    description: str
    weight: float


@dataclass(slots=True)
class ExtractionResult:
    entities: list[ExtractedEntity]
    relationships: list[ExtractedRelationship]
    claims: list[ExtractedClaim] = field(default_factory=list)


@dataclass(slots=True)
class ExtractedClaim:
    entity_name: str
    claim: str
    status: str  # CONFIRMED | REFUTED | SPECULATIVE


# ---------------------------------------------------------------------------
# Prompts (matching Go docscontext extractor)
# ---------------------------------------------------------------------------

ENTITY_PROMPT = """\
You are a knowledge-graph construction assistant.
Given the text below, extract entities and relationships.

Return **only** valid JSON in the following schema:
{
  "entities": [
    {"name": "...", "type": "...", "description": "..."}
  ],
  "relationships": [
    {"source": "...", "target": "...", "predicate": "...", "description": "...", "weight": 0.0}
  ]
}

Rules:
- Entity type must be one of: Person, Organization, Concept, Location, Event, Technology, Other.
- Extract 3-10 entities per chunk.
- Extract up to 15 relationships per chunk. Weight is 0.0-1.0.
- Be concise in descriptions (1-2 sentences max).

Text:
"""

CLAIM_PROMPT = """\
You are a factual analysis assistant.
Given the text below, extract specific verifiable factual claims.

Return **only** valid JSON in the following schema:
{
  "claims": [
    {"entity_name": "...", "claim": "...", "status": "..."}
  ]
}

Rules:
- Status must be one of: CONFIRMED, REFUTED, SPECULATIVE.
- Claims should be specific, verifiable statements about entities.
- Extract up to 10 claims.

Text:
"""

COMMUNITY_SUMMARY_PROMPT = """\
You are a knowledge graph analyst.
Summarize the following community of related entities and relationships.

Entities:
{entities}

Relationships:
{relationships}

Provide:
1. A short title (3-7 words) that captures the main theme.
2. A 2-4 sentence summary explaining the community and key connections.

Format your response as:
TITLE: <title>
SUMMARY: <summary>
"""


# ---------------------------------------------------------------------------
# LLM call helper
# ---------------------------------------------------------------------------

def _llm_call(settings: DocsSettings, prompt: str, max_tokens: int = 2048, temperature: float = 0.0) -> str:
    """Make a single LLM chat completion call. Returns the assistant text."""
    extraction_timeout = max(settings.chat.timeout_seconds, 300.0)
    provider = settings.chat.provider.lower().strip()
    if provider in {"openai", "azure"}:
        base_url = settings.chat.base_url.rstrip("/")
        if provider == "azure":
            base_url = f"{base_url}/openai/deployments/{settings.chat.model}"
        response = httpx.post(
            f"{base_url}/chat/completions",
            headers={"Authorization": f"Bearer {settings.chat.api_key}"},
            json={
                "model": settings.chat.model,
                "temperature": temperature,
                "max_tokens": max_tokens,
                "messages": [
                    {"role": "system", "content": "You are a structured data extraction assistant. Respond with valid JSON only."},
                    {"role": "user", "content": prompt},
                ],
                "response_format": {"type": "json_object"},
            },
            timeout=extraction_timeout,
        )
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]

    if provider == "ollama":
        headers = {}
        if settings.chat.api_key:
            headers["Authorization"] = f"Bearer {settings.chat.api_key}"
            
        response = httpx.post(
            f"{settings.chat.base_url.rstrip('/')}/api/chat",
            headers=headers,
            json={
                "model": settings.chat.model,
                "stream": False,
                "messages": [
                    {"role": "system", "content": "You are a structured data extraction assistant. Respond with valid JSON only."},
                    {"role": "user", "content": prompt},
                ],
                "format": "json",
                "options": {"temperature": temperature, "num_predict": max_tokens},
            },
            timeout=extraction_timeout,
        )
        response.raise_for_status()
        return str(response.json().get("message", {}).get("content", "{}"))

    if provider == "mock":
        # Return mock JSON response for testing
        return json.dumps({
            "entities": [],
            "relationships": [],
            "claims": []
        })

    raise ValueError(f"Unsupported LLM provider for extraction: {provider}")


def _parse_json(text: str) -> dict:
    """Try to parse JSON from LLM output, handling markdown fences."""
    text = text.strip()
    if not text:
        return {}
    if not text.startswith("{"):
        start = text.find("{")
        end = text.rfind("}")
        if start >= 0 and end > start:
            text = text[start : end + 1]
        else:
            return {}
    try:
        return json.loads(text)
    except Exception:
        return {}


# ---------------------------------------------------------------------------
# Extraction functions
# ---------------------------------------------------------------------------

def extract_entities(settings: DocsSettings, chunks: list[str]) -> ExtractionResult:
    """Extract entities and relationships from document chunks via LLM."""
    combined = "\n\n".join(chunks)[:8000]
    prompt = ENTITY_PROMPT + combined
    raw = _llm_call(settings, prompt, max_tokens=4096, temperature=0.0)
    data = _parse_json(raw)

    entities: list[ExtractedEntity] = []
    for item in data.get("entities", []):
        etype = str(item.get("type", "Other"))
        if etype not in ENTITY_TYPES:
            etype = "Other"
        entities.append(ExtractedEntity(
            name=str(item.get("name", "")),
            entity_type=etype,
            description=str(item.get("description", "")),
        ))

    relationships: list[ExtractedRelationship] = []
    for item in data.get("relationships", []):
        try:
            weight = float(item.get("weight", 0.5))
        except (TypeError, ValueError):
            weight = 0.5
        relationships.append(ExtractedRelationship(
            source=str(item.get("source", "")),
            target=str(item.get("target", "")),
            predicate=str(item.get("predicate", "")),
            description=str(item.get("description", "")),
            weight=max(0.0, min(1.0, weight)),
        ))

    return ExtractionResult(entities=entities, relationships=relationships)


def extract_claims(settings: DocsSettings, chunks: list[str]) -> list[ExtractedClaim]:
    """Extract factual claims from document chunks via LLM."""
    combined = "\n\n".join(chunks)[:6000]
    prompt = CLAIM_PROMPT + combined
    raw = _llm_call(settings, prompt, max_tokens=1024, temperature=0.0)
    data = _parse_json(raw)

    claims: list[ExtractedClaim] = []
    for item in data.get("claims", []):
        status = str(item.get("status", "SPECULATIVE")).upper()
        if status not in CLAIM_STATUSES:
            status = "SPECULATIVE"
        claims.append(ExtractedClaim(
            entity_name=str(item.get("entity_name", "")),
            claim=str(item.get("claim", "")),
            status=status,
        ))
    return claims


def summarize_community(
    settings: DocsSettings,
    entity_descriptions: list[str],
    relationship_descriptions: list[str],
) -> tuple[str, str]:
    """Generate a title and summary for a community via LLM. Returns (title, summary)."""
    prompt = COMMUNITY_SUMMARY_PROMPT.format(
        entities="\n".join(f"- {d}" for d in entity_descriptions[:20]),
        relationships="\n".join(f"- {d}" for d in relationship_descriptions[:20]),
    )
    provider = settings.chat.provider.lower().strip()
    if provider in {"openai", "azure"}:
        base_url = settings.chat.base_url.rstrip("/")
        if provider == "azure":
            base_url = f"{base_url}/openai/deployments/{settings.chat.model}"
        response = httpx.post(
            f"{base_url}/chat/completions",
            headers={"Authorization": f"Bearer {settings.chat.api_key}"},
            json={
                "model": settings.chat.model,
                "temperature": 0.3,
                "max_tokens": 512,
                "messages": [
                    {"role": "system", "content": "You are a knowledge graph analyst."},
                    {"role": "user", "content": prompt},
                ],
            },
            timeout=max(settings.chat.timeout_seconds, 120.0),
        )
        response.raise_for_status()
        text = response.json()["choices"][0]["message"]["content"]
    elif provider == "ollama":
        headers = {}
        if settings.chat.api_key:
            headers["Authorization"] = f"Bearer {settings.chat.api_key}"
            
        response = httpx.post(
            f"{settings.chat.base_url.rstrip('/')}/api/chat",
            headers=headers,
            json={
                "model": settings.chat.model,
                "stream": False,
                "messages": [
                    {"role": "system", "content": "You are a knowledge graph analyst."},
                    {"role": "user", "content": prompt},
                ],
                "options": {"temperature": 0.3, "num_predict": 512},
            },
            timeout=max(settings.chat.timeout_seconds, 120.0),
        )
        response.raise_for_status()
        text = str(response.json().get("message", {}).get("content", ""))
    else:
        return ("", "")

    return _parse_community_summary(text)


_TITLE_RE = re.compile(r"TITLE:\s*(.+)", re.IGNORECASE)
_SUMMARY_RE = re.compile(r"SUMMARY:\s*(.+)", re.IGNORECASE | re.DOTALL)


def _parse_community_summary(text: str) -> tuple[str, str]:
    title_match = _TITLE_RE.search(text)
    summary_match = _SUMMARY_RE.search(text)
    title = title_match.group(1).strip() if title_match else ""
    summary = summary_match.group(1).strip() if summary_match else text.strip()
    if title and summary.startswith(title):
        summary = summary[len(title):].strip().lstrip(".")
    return (title[:200], summary[:2000])
