"""LLM client abstraction — swap providers without touching business logic.

Provider is configured via LLM_PROVIDER env var (default: "gemini").
API keys are loaded internally via config.load_api_key().

Usage:
    llm = create_client()              # uses CLASSIFIER_PROVIDER from config
    llm = create_client("groq")        # explicit override
    response = await llm.call("prompt", "text", max_tokens=300)
"""

import logging
import os
import sys
from abc import ABC, abstractmethod

import httpx

from neo_cortex import config

logger = logging.getLogger(__name__)

_VERIFY_SSL = sys.platform != "win32"

from functools import partial
from neo_cortex._log import boot_log
_bl = partial(boot_log, "LLM")


class LLMClient(ABC):
    """Abstract LLM client with transport + task methods.

    Transport: call(prompt, text) -> str (abstract, implemented by each provider).
    Task methods: distill(), classify(), etc. (concrete in base, use _prompts dict).
    """

    def __init__(self):
        from neo_cortex.prompts import DEFAULT_PROMPTS
        self._prompts = dict(DEFAULT_PROMPTS)

    @abstractmethod
    async def call(
        self, prompt: str, text: str, max_tokens: int = 250, temperature: float = 0.1
    ) -> str: ...

    @abstractmethod
    async def close(self) -> None: ...

    @property
    @abstractmethod
    def name(self) -> str:
        """Provider name for logging."""
        ...

    # --- helpers ---

    def _strip_fences(self, text: str) -> str:
        """Strip markdown code fences from LLM output."""
        text = text.strip()
        if text.startswith("```"):
            text = "\n".join(text.split("\n")[1:])
        if text.endswith("```"):
            text = "\n".join(text.split("\n")[:-1])
        return text.strip()

    # --- task methods (concrete, override in subclass for provider-specific behavior) ---

    async def distill(self, text: str) -> str:
        raw = await self.call(self._prompts["distill"], text, max_tokens=16000, temperature=0.2)
        return self._strip_fences(raw)

    async def classify(self, text: str) -> str:
        raw = await self.call(self._prompts["classify"], text, max_tokens=500)
        return self._strip_fences(raw)

    async def check_noise(self, question: str) -> str:
        raw = await self.call(self._prompts["noise"], question, max_tokens=60, temperature=0.3)
        return self._strip_fences(raw)

    async def analyze_query(self, query: str) -> str:
        raw = await self.call(self._prompts["query_analyze"], query, max_tokens=200)
        return self._strip_fences(raw)

    async def resolve_conflict(self, new_text: str, existing_title: str,
                                existing_summary: str, existing_facts: str) -> str:
        prompt = (
            self._prompts["conflict_prefix"]
            + f"Title: {existing_title}\nSummary: {existing_summary}\nFacts: {existing_facts}\n\n"
            + "INCOMING memory:\n"
        )
        raw = await self.call(prompt, new_text, max_tokens=100)
        return self._strip_fences(raw)

    async def aggregate(self, text: str) -> str:
        return await self.call(self._prompts["aggregate"], text, max_tokens=800)

    async def merge_memories(self, text: str) -> str:
        raw = await self.call(self._prompts["merge"], text, max_tokens=1000)
        return self._strip_fences(raw)

    async def focal_points(self, text: str) -> str:
        return await self.call(self._prompts["focal_points"], text, max_tokens=400)

    async def generate_insight(self, focal_point: str, evidence_text: str) -> str:
        prompt = self._prompts["insight"].format(focal_point=focal_point)
        return await self.call(prompt, evidence_text, max_tokens=400)

    async def verify_truthfulness(self, claim: str, evidence_text: str) -> str:
        prompt = self._prompts["truthfulness"].format(claim=claim)
        return await self.call(prompt, evidence_text, max_tokens=200)

    async def identity_proposal(self, text: str) -> str:
        return await self.call(self._prompts["identity_proposal"], text, max_tokens=400)

    async def coherence_check(self, text: str) -> str:
        return await self.call(self._prompts["coherence_check"], text, max_tokens=200)


class GeminiClient(LLMClient):
    """Gemini — single content block API."""

    def __init__(self, api_key: str, model: str = ""):
        super().__init__()
        self._api_key = api_key
        self._model = model or config.GEMINI_MODEL
        self._client = httpx.AsyncClient(timeout=30.0, verify=_VERIFY_SSL)
        _bl(f"GeminiClient init: model={self._model}")

    @property
    def name(self) -> str:
        return f"gemini/{self._model}"

    async def call(
        self, prompt: str, text: str, max_tokens: int = 250, temperature: float = 0.1
    ) -> str:
        url = f"{config.GEMINI_URL}/models/{self._model}:generateContent"
        content = f"{prompt}\n{text}" if text else prompt
        resp = await self._client.post(
            url,
            headers={"x-goog-api-key": self._api_key},
            json={
                "contents": [{"parts": [{"text": content}]}],
                "generationConfig": {
                    "temperature": temperature,
                    "maxOutputTokens": max_tokens,
                    "thinkingConfig": {"thinkingBudget": 0},
                },
            },
        )
        resp.raise_for_status()
        data = resp.json()
        return data["candidates"][0]["content"]["parts"][0]["text"].strip()

    async def close(self) -> None:
        await self._client.aclose()


class GroqClient(LLMClient):
    """Groq — fast inference via LPU, OpenAI-compatible messages."""

    def __init__(self, api_key: str, model: str = ""):
        super().__init__()
        from neo_cortex.prompts.groq import GROQ_PROMPTS
        self._prompts.update(GROQ_PROMPTS)
        self._api_key = api_key
        self._model = model or config.GROQ_MODEL
        self._client = httpx.AsyncClient(timeout=15.0, verify=_VERIFY_SSL)
        _bl(f"GroqClient init: model={self._model}")

    @property
    def name(self) -> str:
        return f"groq/{self._model}"

    async def call(
        self, prompt: str, text: str, max_tokens: int = 250, temperature: float = 0.1
    ) -> str:
        if text:
            messages = [
                {"role": "system", "content": prompt},
                {"role": "user", "content": text},
            ]
        else:
            messages = [{"role": "user", "content": prompt}]
        resp = await self._client.post(
            config.GROQ_URL,
            headers={
                "Authorization": f"Bearer {self._api_key}",
                "Content-Type": "application/json",
            },
            json={
                "model": self._model,
                "messages": messages,
                "temperature": temperature,
                "max_tokens": max_tokens,
            },
        )
        resp.raise_for_status()
        data = resp.json()
        return data["choices"][0]["message"]["content"].strip()

    async def close(self) -> None:
        await self._client.aclose()


class OpenAICompatibleClient(LLMClient):
    """OpenAI-compatible API — covers Ollama, LM Studio, vLLM, OpenAI, etc."""

    def __init__(self, base_url: str, api_key: str, model: str,
                 prompt_overrides: dict[str, str] | None = None):
        super().__init__()
        if prompt_overrides:
            self._prompts.update(prompt_overrides)
        self._base_url = base_url.rstrip("/")
        self._api_key = api_key
        self._model = model
        self._client = httpx.AsyncClient(timeout=120.0, verify=_VERIFY_SSL)

    @property
    def name(self) -> str:
        return f"openai-compatible/{self._model}"

    async def call(self, prompt: str, text: str, max_tokens: int = 250, temperature: float = 0.1) -> str:
        headers = {"Content-Type": "application/json"}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"
        if text:
            messages = [
                {"role": "system", "content": prompt},
                {"role": "user", "content": text},
            ]
        else:
            messages = [{"role": "user", "content": prompt}]
        resp = await self._client.post(
            f"{self._base_url}/chat/completions",
            headers=headers,
            json={
                "model": self._model,
                "messages": messages,
                "temperature": temperature,
                "max_tokens": max_tokens,
            },
        )
        resp.raise_for_status()
        data = resp.json()
        return data["choices"][0]["message"]["content"].strip()

    async def close(self) -> None:
        await self._client.aclose()


def _extract_text(msg) -> str:
    """Extract text from a Claude SDK message (handles str and list[Block])."""
    if not hasattr(msg, "content"):
        return ""
    content = msg.content
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "".join(block.text for block in content if hasattr(block, "text"))
    return ""


class ClaudeCodeClient(LLMClient):
    """Uses Claude Code CLI via Agent SDK — persistent session.

    First call() connects to Claude Code CLI and keeps the process alive.
    Subsequent calls reuse the same process. close() disconnects.
    """

    def __init__(self, model: str = "haiku"):
        super().__init__()
        from neo_cortex.prompts.anthropic import ANTHROPIC_PROMPTS
        self._prompts.update(ANTHROPIC_PROMPTS)
        self._model = model
        self._client = None  # ClaudeSDKClient, lazy init
        _bl(f"ClaudeCodeClient init: model={model}")

    @property
    def name(self) -> str:
        return f"claude-code/{self._model}"

    async def _ensure_connected(self):
        if self._client is None:
            import tempfile
            from claude_code_sdk import ClaudeCodeOptions, ClaudeSDKClient, SystemMessage
            from claude_code_sdk._internal import message_parser

            # Patch SDK parser to skip unknown message types (rate_limit_event etc)
            _orig_parse = message_parser.parse_message

            def _tolerant_parse(data):
                try:
                    return _orig_parse(data)
                except Exception:
                    return SystemMessage(subtype="unknown", data=data)

            message_parser.parse_message = _tolerant_parse

            # Empty MCP config and cwd — no servers, no CLAUDE.md, no hooks
            self._mcp_tmp = tempfile.NamedTemporaryFile(
                mode="w", suffix=".json", delete=False,
            )
            self._mcp_tmp.write('{"mcpServers":{}}')
            self._mcp_tmp.flush()
            self._cwd_tmp = tempfile.mkdtemp(prefix="claude-llm-")

            self._client = ClaudeSDKClient(
                ClaudeCodeOptions(
                    model=self._model,
                    system_prompt="Follow instructions exactly. Return only what is asked.",
                    mcp_servers=self._mcp_tmp.name,
                    cwd=self._cwd_tmp,
                    env={"MAX_THINKING_TOKENS": "0"},
                    extra_args={
                        "no-session-persistence": None,
                        "strict-mcp-config": None,
                        "tools": "",
                        "setting-sources": "",
                    },
                ),
            )
            await self._client.connect()

    async def call(
        self, prompt: str, text: str, max_tokens: int = 250, temperature: float = 0.1
    ) -> str:
        import uuid
        from claude_code_sdk import ResultMessage
        await self._ensure_connected()
        full_prompt = f"{prompt}\n{text}" if text else prompt
        await self._client.query(full_prompt, session_id=uuid.uuid4().hex)
        result = ""
        async for msg in self._client.receive_response():
            if isinstance(msg, ResultMessage):
                break
            result += _extract_text(msg)
        return result.strip()

    async def close(self) -> None:
        if self._client is not None:
            await self._client.disconnect()
            self._client = None
        if hasattr(self, "_mcp_tmp") and self._mcp_tmp:
            import os
            os.unlink(self._mcp_tmp.name)
            self._mcp_tmp = None
        if hasattr(self, "_cwd_tmp") and self._cwd_tmp:
            import shutil
            shutil.rmtree(self._cwd_tmp, ignore_errors=True)
            self._cwd_tmp = None


try:
    from anthropic import AsyncAnthropic
except ImportError:
    AsyncAnthropic = None  # type: ignore[misc,assignment]


class AnthropicClient(LLMClient):
    """Anthropic Claude via API — direct SDK, no CLI wrapper.

    Uses system param for prompt, user message for text.
    Requires `pip install neo-cortex-mcp[anthropic]`.
    """

    def __init__(self, api_key: str, model: str = "claude-haiku-4-5-20251001"):
        super().__init__()
        from neo_cortex.prompts.anthropic import ANTHROPIC_PROMPTS
        self._prompts.update(ANTHROPIC_PROMPTS)
        if AsyncAnthropic is None:
            raise ImportError("Install anthropic: pip install neo-cortex-mcp[anthropic]")
        self._model = model
        self._client = AsyncAnthropic(api_key=api_key)
        _bl(f"AnthropicClient init: model={model}")

    @property
    def name(self) -> str:
        return f"anthropic/{self._model}"

    async def call(
        self, prompt: str, text: str, max_tokens: int = 250, temperature: float = 0.1
    ) -> str:
        kwargs: dict = {
            "model": self._model,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "messages": [{"role": "user", "content": text if text else prompt}],
        }
        if text:
            kwargs["system"] = prompt
        resp = await self._client.messages.create(**kwargs)
        return resp.content[0].text.strip()

    async def close(self) -> None:
        await self._client.close()


class ZhipuClient(LLMClient):
    """Zhipu AI (Z.AI) — uses the Anthropic-compatible proxy endpoint.

    The Coding Plan subscription only works via /api/anthropic (not /api/paas/v4).
    Available models: glm-4.5-air, glm-4.7, glm-5.
    """

    def __init__(self, api_key: str, model: str = "glm-4.7"):
        super().__init__()
        from neo_cortex.prompts.zhipu import ZHIPU_PROMPTS
        self._prompts.update(ZHIPU_PROMPTS)
        self._api_key = api_key
        self._model = model
        self._client = httpx.AsyncClient(timeout=120.0, verify=_VERIFY_SSL)
        _bl(f"ZhipuClient init: model={model}")

    @property
    def name(self) -> str:
        return f"z.ai/{self._model}"

    async def call(
        self, prompt: str, text: str, max_tokens: int = 250, temperature: float = 0.1
    ) -> str:
        kwargs: dict = {
            "model": self._model,
            "max_tokens": max_tokens,
            "messages": [{"role": "user", "content": text if text else prompt}],
        }
        if text:
            kwargs["system"] = prompt
        resp = await self._client.post(
            "https://api.z.ai/api/anthropic/v1/messages",
            headers={
                "x-api-key": self._api_key,
                "anthropic-version": "2023-06-01",
                "Content-Type": "application/json",
            },
            json=kwargs,
        )
        resp.raise_for_status()
        data = resp.json()
        return data["content"][0]["text"].strip()

    async def close(self) -> None:
        await self._client.aclose()


def _load_openai_prompts() -> dict[str, str]:
    from neo_cortex.prompts.openai import OPENAI_PROMPTS
    return OPENAI_PROMPTS


_PROVIDERS = {
    "gemini": lambda: GeminiClient(config.load_api_key("gemini")),
    "groq": lambda: GroqClient(config.load_api_key("groq")),
    "openai-compatible": lambda: OpenAICompatibleClient(
        config.LLM_BASE_URL,
        os.environ.get("OPENAI_API_KEY", ""),
        config.LLM_MODEL,
    ),
    "openai": lambda: OpenAICompatibleClient(
        "https://api.openai.com/v1",
        config.load_api_key("openai"),
        os.environ.get("OPENAI_MODEL", "gpt-4.1-mini"),
        prompt_overrides=_load_openai_prompts(),
    ),
    "anthropic": lambda: AnthropicClient(
        config.load_api_key("anthropic"),
        os.environ.get("ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"),
    ),
    "x.ai": lambda: OpenAICompatibleClient(
        "https://api.x.ai/v1",
        config.load_api_key("xai"),
        os.environ.get("XAI_MODEL", "grok-3-mini"),
    ),
    "z.ai": lambda: ZhipuClient(
        config.load_api_key("zai"),
        os.environ.get("ZAI_MODEL", "glm-4.7"),
    ),
    "claude-code": lambda: ClaudeCodeClient(
        model=config.LLM_MODEL or "haiku",
    ),
}


def create_client(provider: str | None = None) -> LLMClient:
    """Create an LLM client based on config or explicit provider name.

    Provider is resolved from:
    1. Explicit `provider` argument
    2. CLASSIFIER_PROVIDER env var
    3. Default: "gemini"
    """
    name = provider or config.LLM_PROVIDER
    factory = _PROVIDERS.get(name)
    if not factory:
        raise ValueError(f"Unknown LLM provider '{name}'. Available: {list(_PROVIDERS.keys())}")
    client = factory()
    _bl(f"create_client: provider={name} → {client.name}")
    return client
