"""
LLM client — supports 'claude' and 'openai' providers.

Uses the model configured in the models table (fetched via write_model_id
or resolve_model_id config keys).
"""
from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from typing import Any, Optional

import asyncpg

log = logging.getLogger("festinger.llm")


@dataclass
class ModelConfig:
    provider: str       # 'claude', 'openai', or 'lm-studio'
    model_name: str
    api_key: str
    base_url: str = ""  # custom endpoint for lm-studio / openai-compatible


async def get_model_config(pool: asyncpg.Pool, model_id: str) -> Optional[ModelConfig]:
    if not model_id:
        return None
    try:
        mid = int(model_id)
    except ValueError:
        return None
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            "SELECT provider, model_name, api_key, base_url FROM models WHERE id=$1", mid
        )
    if not row:
        return None
    return ModelConfig(
        provider=row["provider"],
        model_name=row["model_name"],
        api_key=row["api_key"],
        base_url=row["base_url"] or "",
    )


async def call_llm(model: ModelConfig, prompt: str) -> str:
    """Call the configured LLM and return the text response."""
    if model.provider == "claude":
        return await _call_claude(model, prompt)
    elif model.provider in ("openai", "lm-studio"):
        return await _call_openai(model, prompt)
    else:
        raise ValueError(f"Unknown LLM provider: {model.provider!r}")


async def _call_claude(model: ModelConfig, prompt: str) -> str:
    import anthropic
    client = anthropic.AsyncAnthropic(api_key=model.api_key)
    message = await client.messages.create(
        model=model.model_name,
        max_tokens=1024,
        messages=[{"role": "user", "content": prompt}],
    )
    return message.content[0].text


async def _call_openai(model: ModelConfig, prompt: str) -> str:
    from openai import AsyncOpenAI
    kwargs: dict = {"api_key": model.api_key or "lm-studio"}
    if model.base_url:
        kwargs["base_url"] = model.base_url
    client = AsyncOpenAI(**kwargs)
    response = await client.chat.completions.create(
        model=model.model_name,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=1024,
    )
    return response.choices[0].message.content or ""


# ---------------------------------------------------------------------------
# Structured prompts
# ---------------------------------------------------------------------------

CONTEXT_DISCOVER_PROMPT_TEMPLATE = """You are building a semantic memory graph for an AI agent.

The following words appeared in the user message and are NOT in the standard English dictionary. They may be domain-specific terms worth remembering, or they may be typos/noise:
{candidates}

Read the user message and evaluate each candidate. For candidates that ARE real domain concepts (project names, system names, proper nouns, technical identifiers), generate a relationship triple ONLY IF the message provides clear evidence. For typos, noise, or common words written oddly, return nothing.

USER MESSAGE:
{context}

Return a JSON array. Each item must have:
  - "concept": one of the candidate words above (use the exact spelling, lowercase)
  - "parent": what it is or belongs to (e.g. "software-repository", "glitch-university")
  - "dimension": one of [{dimensions}] or a new specific one if none fit
  - "is_isa": true if concept IS A type of parent (classification), false if concept BELONGS TO parent (containment)
  - "confidence": 0.0–1.0 reflecting how clearly the message supports this

Only include candidates you are confident are real domain concepts. Return [] if none qualify.
Return ONLY the JSON array. No explanation.
Example:
[
  {{"concept": "gnommoweb", "parent": "software-repository", "dimension": "type", "is_isa": true, "confidence": 0.9}},
  {{"concept": "gnommoweb", "parent": "glitch-university", "dimension": "membership", "is_isa": false, "confidence": 0.85}}
]
"""


@dataclass
class ContextTriple:
    concept: str
    parent: str
    dimension: str
    is_isa: bool
    confidence: float


def parse_context_triples(response: str) -> list[ContextTriple]:
    """Parse JSON array of discovered triples from LLM response."""
    try:
        start = response.find("[")
        end = response.rfind("]") + 1
        if start == -1 or end == 0:
            return []
        data = json.loads(response[start:end])
        triples = []
        for item in data:
            if not isinstance(item, dict):
                continue
            triples.append(ContextTriple(
                concept=str(item.get("concept", "")).strip().lower(),
                parent=str(item.get("parent", "")).strip().lower(),
                dimension=str(item.get("dimension", "type")).strip().lower(),
                is_isa=bool(item.get("is_isa", True)),
                confidence=float(item.get("confidence", 0.7)),
            ))
        return triples
    except (json.JSONDecodeError, ValueError) as e:
        log.warning("failed to parse context triples: %s", e)
        return []


# Kept for backwards-compatibility with the resolution job prompts below.
@dataclass
class LLMTriple:
    parent: str
    dimension: str
    is_isa: bool
    confidence: float


def parse_llm_triples(response: str) -> list[LLMTriple]:
    """Parse JSON array of single-concept triples (used by resolution job only)."""
    try:
        start = response.find("[")
        end = response.rfind("]") + 1
        if start == -1 or end == 0:
            return []
        data = json.loads(response[start:end])
        triples = []
        for item in data:
            if not isinstance(item, dict):
                continue
            triples.append(LLMTriple(
                parent=str(item.get("parent", "")).strip().lower(),
                dimension=str(item.get("dimension", "type")).strip().lower(),
                is_isa=bool(item.get("is_isa", True)),
                confidence=float(item.get("confidence", 0.7)),
            ))
        return triples
    except (json.JSONDecodeError, ValueError) as e:
        log.warning("failed to parse LLM triples: %s", e)
        return []


RESOLVE_ISA_ISA_PROMPT = """You are a knowledge graph resolution assistant.

Concept: "{concept}"
Current fact: "{concept}" is a "{existing_parent}" in dimension "{dimension}"
Conflicting fact: "{concept}" is a "{incoming_parent}" in dimension "{dimension}"

Both facts appear to be simultaneously true. The dimension "{dimension}" is too coarse.
Propose two more specific dimension names to replace it:
  - existing_dimension: the dimension where "{concept}" as "{existing_parent}" is valid
  - new_dimension: the dimension where "{concept}" as "{incoming_parent}" is valid

Known dimensions for reference: {known_dimensions}

Return ONLY a JSON object:
{{"decision": "decompose", "existing_dimension": "...", "new_dimension": "...", "reasoning": "..."}}

If the conflicting fact is clearly wrong, return:
{{"decision": "dismiss", "reasoning": "..."}}
"""

RESOLVE_ISPART_ISPART_PROMPT = """You are a knowledge graph resolution assistant.

Concept: "{concept}"
Current fact: "{concept}" belongs to "{existing_parent}" in dimension "{dimension}"
Conflicting fact: "{concept}" belongs to "{incoming_parent}" in dimension "{dimension}"

These facts contradict each other — a thing can only belong to one place per dimension.

Return ONLY a JSON object with your decision:
{{"decision": "update", "reasoning": "..."}}  — if the incoming fact is more current/correct
{{"decision": "dismiss", "reasoning": "..."}}  — if the existing fact is still correct
"""

RESOLVE_MISCLASS_PROMPT = """You are a knowledge graph resolution assistant.

Concept: "{concept}"
A relationship was classified inconsistently:
  - Existing: "{concept}" ISA "{existing_parent}" in dimension "{dimension}"
  - Incoming: "{concept}" ISPART "{incoming_parent}" in dimension "{dimension}"

The incoming fact uses a different relation type. Suggest the correct dimension for the incoming fact.

Known dimensions: {known_dimensions}

Return ONLY a JSON object:
{{"decision": "reclassify", "correct_dimension": "...", "reasoning": "..."}}
"""


@dataclass
class ResolutionDecision:
    decision: str           # 'decompose', 'update', 'dismiss', 'reclassify'
    existing_dimension: str = ""
    new_dimension: str = ""
    correct_dimension: str = ""
    reasoning: str = ""


def parse_resolution_decision(response: str) -> Optional[ResolutionDecision]:
    try:
        start = response.find("{")
        end = response.rfind("}") + 1
        if start == -1 or end == 0:
            return None
        data = json.loads(response[start:end])
        return ResolutionDecision(
            decision=str(data.get("decision", "dismiss")).lower(),
            existing_dimension=str(data.get("existing_dimension", "")).lower(),
            new_dimension=str(data.get("new_dimension", "")).lower(),
            correct_dimension=str(data.get("correct_dimension", "")).lower(),
            reasoning=str(data.get("reasoning", "")),
        )
    except (json.JSONDecodeError, ValueError) as e:
        log.warning("failed to parse resolution decision: %s", e)
        return None
