"""Provider-backed extraction helpers for docsctx."""

from __future__ import annotations

import json
import re
import time
from dataclasses import dataclass

import httpx

from docsctx.config import DocsSettings
from docsctx.extractor import ExtractedClaim, ExtractedEntity, ExtractedRelationship, ExtractionResult


@dataclass(slots=True)
class ProviderAnalysis:
    summary: str
    headings: list[str]
    entities: list[tuple[str, int]]
    relationships: list[tuple[str, str, int]]



@dataclass(slots=True)
class ProviderAttemptResult:
    analysis: ProviderAnalysis | None
    attempts: int
    error: str = ""

    @property
    def succeeded(self) -> bool:
        return self.analysis is not None


class BaseProvider:
    def analyze(self, title: str, content: str) -> ProviderAnalysis:
        raise NotImplementedError

    def extract_graph(self, title: str, content: str) -> ExtractionResult:
        """Extract entities, relationships, and claims via LLM."""
        raise NotImplementedError

    def synthesize_global_answer(self, query: str, community_summaries: str) -> str:
        """Synthesize an answer from community summaries. Override in subclasses."""
        return ""


class MockProvider(BaseProvider):
    WORD_RE = re.compile(r"[A-Za-z][A-Za-z0-9_-]{2,}")

    def analyze(self, title: str, content: str) -> ProviderAnalysis:
        words: dict[str, int] = {}
        for match in self.WORD_RE.finditer(content.lower()):
            word = match.group(0)
            words[word] = words.get(word, 0) + 1
        top = sorted(words.items(), key=lambda item: (-item[1], item[0]))[:8]
        entities = [(word, count) for word, count in top]
        relationships: list[tuple[str, str, int]] = []
        for index, (left, left_count) in enumerate(entities[:4]):
            for right, right_count in entities[index + 1 : 4]:
                relationships.append((left, right, min(left_count, right_count)))
        headings = [line.lstrip("#").strip() for line in content.splitlines() if line.strip().startswith("#")][:10]
        summary = " ".join(content.splitlines()[:3]).strip()[:400]
        return ProviderAnalysis(summary=summary, headings=headings, entities=entities, relationships=relationships)

    def extract_graph(self, title: str, content: str) -> ExtractionResult:
        """Mock extraction returning empty results."""
        return ExtractionResult(entities=[], relationships=[], claims=[])

    def synthesize_global_answer(self, query: str, community_summaries: str) -> str:
        return f"Based on indexed documents: {community_summaries[:200]}"


from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import SystemMessage, HumanMessage


class LangChainProvider(BaseProvider):
    def __init__(self, chat_model: BaseChatModel):
        self.chat_model = chat_model

    def analyze(self, title: str, content: str) -> ProviderAnalysis:
        prompt = self._prompt(title, content)
        messages = [
            SystemMessage(content="Extract structured document metadata as strict JSON."),
            HumanMessage(content=prompt),
        ]
        response = self.chat_model.invoke(messages)
        return _parse_analysis_json(str(response.content))

    def extract_graph(self, title: str, content: str) -> ExtractionResult:
        """Extract entities, relationships, and claims via LLM."""
        prompt = self._extraction_prompt(title, content)
        messages = [
            SystemMessage(content="You are a knowledge graph extraction expert. Extract entities, relationships, and claims as strict JSON."),
            HumanMessage(content=prompt),
        ]
        response = self.chat_model.invoke(messages)
        return _parse_extraction_json(str(response.content))

    @staticmethod
    def _prompt(title: str, content: str) -> str:
        clipped = content[:12000]
        return (
            "Return only JSON with fields: summary (string), headings (array of strings), "
            "entities (array of {name,score}), relationships (array of {source,target,weight}). "
            f"Title: {title}\nContent:\n{clipped}"
        )

    @staticmethod
    def _extraction_prompt(title: str, content: str) -> str:
        clipped = content[:12000]
        return (
            "Extract a knowledge graph from the document. Return JSON with:\n"
            '- "entities": array of {name, type (Person/Org/Concept/Location/Event/Technology), description, score}\n'
            '- "relationships": array of {source, target, predicate (describes/relates_to/part_of/caused_by), description, weight}\n'
            '- "claims": array of {entity, claim, status (CONFIRMED/REFUTED/SPECULATIVE)}\n'
            f"Title: {title}\nContent:\n{clipped}"
        )

    def synthesize_global_answer(self, query: str, community_summaries: str) -> str:
        clipped = community_summaries[:6000]
        messages = [
            SystemMessage(content="You are a helpful assistant that synthesizes document knowledge to answer questions concisely."),
            HumanMessage(content=(
                f"Community summaries from indexed documents:\n{clipped}\n\n"
                f"Answer this question using only the provided summaries: {query}"
            )),
        ]
        response = self.chat_model.invoke(messages)
        return str(response.content)


def create_provider(settings: DocsSettings) -> BaseProvider | None:
    provider = settings.chat.provider.lower().strip()
    model = settings.chat.model
    if provider in {"", "none"}:
        return None
    if provider == "mock":
        return MockProvider()

    chat_model = None
    if provider == "ollama":
        from langchain_ollama import ChatOllama
        client_kwargs = {"timeout": settings.chat.timeout_seconds}
        if settings.chat.api_key:
            client_kwargs["headers"] = {"Authorization": f"Bearer {settings.chat.api_key}"}
            
        chat_model = ChatOllama(
            base_url=settings.chat.base_url,
            model=model,
            format="json",
            temperature=0,
            client_kwargs=client_kwargs,
        )
    elif provider in {"openai", "azure"}:
        if not settings.chat.api_key:
            raise ValueError("DOCSCTX_CHAT__API_KEY is required for openai/azure providers")
        if provider == "azure":
            from langchain_openai import AzureChatOpenAI
            chat_model = AzureChatOpenAI(
                azure_endpoint=settings.chat.base_url,
                api_key=settings.chat.api_key,
                azure_deployment=model,
                api_version="2024-02-15-preview",
                temperature=0,
                model_kwargs={"response_format": {"type": "json_object"}},
                timeout=settings.chat.timeout_seconds,
            )
        else:
            from langchain_openai import ChatOpenAI
            chat_model = ChatOpenAI(
                base_url=settings.chat.base_url,
                api_key=settings.chat.api_key,
                model=model,
                temperature=0,
                model_kwargs={"response_format": {"type": "json_object"}},
                timeout=settings.chat.timeout_seconds,
            )

    if chat_model is None:
        raise ValueError(f"Unsupported docs LLM provider: {provider}")

    return LangChainProvider(chat_model)



def validate_provider_access(settings: DocsSettings, provider: BaseProvider | None = None) -> BaseProvider:
    """Ensure provider configuration exists and the backend is reachable."""
    requested_provider = settings.chat.provider.lower().strip()
    if requested_provider in {"", "none"}:
        raise ValueError("DOCSCTX_CHAT__PROVIDER must be set to a real LLM provider (or 'mock' for testing)")

    # The provider might be cached or passed in
    p = provider or create_provider(settings)
    if p is None:
        raise ValueError("Failed to initialize a provider: configuration is incomplete")

    base_url = (settings.chat.base_url or "").rstrip("/")
    if base_url and requested_provider != "mock":
        try:
            import httpx
            if requested_provider == "ollama":
                headers = {}
                if settings.chat.api_key:
                    headers["Authorization"] = f"Bearer {settings.chat.api_key}"
                response = httpx.get(f"{base_url}/api/tags", headers=headers, timeout=5.0)
                response.raise_for_status()
            elif requested_provider in {"openai", "azure"}:
                headers = {"Authorization": f"Bearer {settings.chat.api_key}"}
                if requested_provider == "azure":
                    headers["api-key"] = settings.chat.api_key
                    url = f"{base_url}/models?api-version=2024-02-15-preview"
                else:
                    url = f"{base_url}/models"
                response = httpx.get(url, headers=headers, timeout=5.0)
                response.raise_for_status()
            else:
                response = httpx.get(base_url, timeout=5.0)
        except Exception as exc:
            raise RuntimeError(f"LLM provider host {base_url} is not accessible: {exc}")

    return p


def analyze_with_retries(
    provider: BaseProvider,
    title: str,
    content: str,
    max_retries: int = 0,
    backoff_seconds: float = 0.0,
) -> ProviderAttemptResult:
    attempts = 0
    last_error = ""
    total_attempts = max(1, max_retries + 1)
    for attempt in range(total_attempts):
        attempts = attempt + 1
        try:
            return ProviderAttemptResult(analysis=provider.analyze(title, content), attempts=attempts)
        except Exception as exc:
            last_error = f"{type(exc).__name__}: {exc}"
            if attempt + 1 >= total_attempts:
                break
            if backoff_seconds > 0:
                time.sleep(backoff_seconds * attempts)
    return ProviderAttemptResult(analysis=None, attempts=attempts, error=last_error)


def _parse_analysis_json(raw_text: str) -> ProviderAnalysis:
    text = raw_text.strip()
    if not text:
        return ProviderAnalysis(summary="", headings=[], entities=[], relationships=[])
    if not text.startswith("{"):
        start = text.find("{")
        end = text.rfind("}")
        if start >= 0 and end > start:
            text = text[start : end + 1]
    payload = json.loads(text)
    summary = str(payload.get("summary") or "")
    headings = [str(item) for item in payload.get("headings", []) if isinstance(item, str)]
    entities = []
    for item in payload.get("entities", []):
        if not isinstance(item, dict):
            continue
        name = str(item.get("name") or "").strip()
        score = int(item.get("score") or 0)
        if name:
            entities.append((name, score))
    relationships = []
    for item in payload.get("relationships", []):
        if not isinstance(item, dict):
            continue
        source = str(item.get("source") or "").strip()
        target = str(item.get("target") or "").strip()
        weight = int(item.get("weight") or 0)
        if source and target:
            relationships.append((source, target, weight))
    return ProviderAnalysis(summary=summary, headings=headings, entities=entities, relationships=relationships)


def _parse_extraction_json(raw_text: str) -> ExtractionResult:
    """Parse knowledge graph JSON from LLM."""
    text = raw_text.strip()
    if not text:
        return ExtractionResult(entities=[], relationships=[], claims=[])
    if not text.startswith("{"):
        start = text.find("{")
        end = text.rfind("}")
        if start >= 0 and end > start:
            text = text[start : end + 1]
    try:
        payload = json.loads(text)
    except (json.JSONDecodeError, ValueError):
        return ExtractionResult(entities=[], relationships=[], claims=[])

    # Parse entities
    entities: list[ExtractedEntity] = []
    for item in payload.get("entities", []):
        if not isinstance(item, dict):
            continue
        name = str(item.get("name") or "").strip()
        entity_type = str(item.get("type") or "Concept").strip()
        description = str(item.get("description") or "").strip()
        score = int(item.get("score") or 0)
        if name:
            entities.append(ExtractedEntity(name=name, entity_type=entity_type, description=description, score=score))

    # Parse relationships
    relationships: list[ExtractedRelationship] = []
    for item in payload.get("relationships", []):
        if not isinstance(item, dict):
            continue
        source = str(item.get("source") or "").strip()
        target = str(item.get("target") or "").strip()
        predicate = str(item.get("predicate") or "").strip()
        description = str(item.get("description") or "").strip()
        weight = float(item.get("weight") or 1.0)
        if source and target:
            relationships.append(ExtractedRelationship(source=source, target=target, predicate=predicate, description=description, weight=weight))

    # Parse claims
    claims: list[ExtractedClaim] = []
    for item in payload.get("claims", []):
        if not isinstance(item, dict):
            continue
        entity = str(item.get("entity") or "").strip()
        claim = str(item.get("claim") or "").strip()
        status = str(item.get("status") or "SPECULATIVE").strip()
        if entity and claim:
            claims.append(ExtractedClaim(entity_name=entity, claim=claim, status=status))

    return ExtractionResult(entities=entities, relationships=relationships, claims=claims)
