"""Model access for the graph-memory engines (HippoRAG 2 and ReMem).

The engines call a language model through an object with `infer(messages, **kw)` returning
`(text, metadata, cache_hit)` and embed text through an object with `encode` / `batch_encode`.
`EngineChat` and `EngineEmbedder` provide those interfaces on top of the evaluation's chat client
and an OpenAI-compatible embeddings endpoint, so the engines use the configured backbone and the
token usage of every call is logged under the current phase.
"""

from __future__ import annotations

import json
import logging
import os
import sys
import time
import types
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np

from ..llm_client import ENDPOINTS, MAX_ATTEMPTS, SDK_MAX_RETRIES, USAGE, api_key_for, http_timeout

logger = logging.getLogger(__name__)

VENDOR = Path(__file__).resolve().parents[2] / "vendor"
_FENCE_OPEN, _FENCE_CLOSE = "```json\n", "```"


def add_engine_source(name: str) -> None:
    """Put `vendor/<name>/src` on the import path."""
    src = VENDOR / name / "src"
    if not src.exists():
        raise ImportError(f"{name} source folder not found at {src}")
    if str(src) not in sys.path:
        sys.path.insert(0, str(src))


def stub_optional_imports() -> None:
    """HippoRAG imports vllm at module level for its offline modes, which are not used here.
    When vllm is absent, register empty modules carrying the names those imports expect."""
    if "vllm" in sys.modules:
        return
    try:
        import vllm  # noqa: F401
        return
    except ImportError:
        pass
    names = {"vllm": ("SamplingParams", "LLM"),
             "vllm.model_executor": (),
             "vllm.model_executor.guided_decoding": (),
             "vllm.model_executor.guided_decoding.guided_fields": ("GuidedDecodingRequest",)}
    for module_name, classes in names.items():
        module = types.ModuleType(module_name)
        for cls in classes:
            setattr(module, cls, type(cls, (), {}))
        sys.modules[module_name] = module
    logger.info("vllm is not installed; empty placeholder modules registered")


def read_openie_docs(path: str) -> list:
    """The `docs` list of an engine's OpenIE cache file; empty when the file is missing or unreadable."""
    if not os.path.exists(path):
        return []
    try:
        with open(path, encoding="utf-8") as fh:
            return json.load(fh).get("docs", [])
    except Exception as e:
        logger.warning("could not read the OpenIE cache %s: %s", path, e)
        return []


class EngineChat:
    """`infer` and `batch_infer` over the evaluation's chat client.

    With `max_tokens` set, each call sends an output limit: the engine's own `max_completion_tokens`
    or `max_tokens` when it passes one, else `max_tokens`. With `max_tokens=None` no limit is sent
    and the client's default applies. `echo=True` adds the messages and the reply to the metadata.
    The metadata carries the reply's finish reason ("stop" when the client reports none), so an
    engine can repair JSON that was cut off at the output limit ("length").
    """

    def __init__(self, client, *, temperature: float, max_tokens: int | None = None, echo: bool = False):
        self._client = client
        self._temperature = temperature
        self._max_tokens = max_tokens
        self._echo = echo

    def infer(self, messages, **kw):
        request = {"temperature": kw.get("temperature", self._temperature)}
        if self._max_tokens is not None:
            request["max_tokens"] = kw.get("max_completion_tokens", kw.get("max_tokens", self._max_tokens))
        if "response_format" in kw:
            request["response_format"] = kw["response_format"]
        try:
            reply = self._client.chat(messages=messages, **request)
            text = reply.content
            wants_json = (kw.get("response_format") or {}).get("type") == "json_object"
            if wants_json and text.startswith(_FENCE_OPEN) and text.endswith(_FENCE_CLOSE):
                text = text[len(_FENCE_OPEN):-len(_FENCE_CLOSE)].strip()
        except Exception as e:
            logger.error("engine chat call failed: %s", e)
            raise
        meta = {"prompt": messages, "response": text} if self._echo else {}
        meta.update(prompt_tokens=reply.input_tokens, completion_tokens=reply.output_tokens,
                    finish_reason=getattr(reply, "finish_reason", None) or "stop")
        return text, meta, False

    def batch_infer(self, messages_list, max_workers: int = 10, **kw):
        """`infer` over several message lists in parallel; results keep the input order."""
        with ThreadPoolExecutor(max_workers=max_workers) as pool:
            pending = [pool.submit(self.infer, messages, **kw) for messages in messages_list]
        return [p.result() for p in pending]


class EngineEmbedder:
    """`encode` / `batch_encode` over an OpenAI-compatible embeddings endpoint.

    Blank texts are sent as "empty"; an `instruction` keyword is prefixed to every text. Requests
    hold at most `batch_size` texts and are retried up to MAX_ATTEMPTS times (a 400 response is
    not retried). Rows are scaled to unit length unless `norm=False` is passed.
    """

    def __init__(self, *, model: str | None, batch_size: int | None, dim: int | None = None,
                 api_key: str | None = None, base_url: str | None = None):
        self.embedding_model_name = model
        self.embedding_size = dim
        self._batch_size = batch_size
        self._api_key = api_key
        self._base_url = base_url
        self._endpoint = None

    def _embeddings(self):
        if self._endpoint is None:
            from openai import OpenAI
            base_url = self._base_url or ENDPOINTS["embedding"]
            self._endpoint = OpenAI(api_key=self._api_key or api_key_for(base_url), base_url=base_url,
                                    timeout=http_timeout(), max_retries=SDK_MAX_RETRIES).embeddings
            logger.info("embedding endpoint ready, model=%s", self.embedding_model_name)
        return self._endpoint

    def _request(self, texts: list) -> list:
        attempt = 0
        while True:
            attempt += 1
            try:
                response = self._embeddings().create(input=texts, model=self.embedding_model_name)
            except Exception as e:
                if attempt < MAX_ATTEMPTS and "Error code: 400" not in str(e):
                    time.sleep(1.5 * attempt)
                    continue
                raise RuntimeError(f"embedding request failed after {attempt} attempt(s): {e}") from e
            usage = getattr(response, "usage", None)
            if usage is not None:
                USAGE.log(self.embedding_model_name, "embed", int(getattr(usage, "prompt_tokens", 0) or 0), 0)
            return [row.embedding for row in response.data]

    def encode(self, texts, **kw) -> np.ndarray:
        texts = [texts] if isinstance(texts, str) else texts
        prefix = kw.get("instruction", "")
        texts = [f"{prefix}{t if t and t.strip() else 'empty'}" for t in texts]
        if not self._batch_size:
            raise RuntimeError("set embedding.batch_size in backbone.yaml")
        step = int(self._batch_size)
        vectors = []
        for start in range(0, len(texts), step):
            vectors.extend(self._request(texts[start:start + step]))
        matrix = np.array(vectors)
        if kw.get("norm", True):
            lengths = np.linalg.norm(matrix, axis=1, keepdims=True)
            matrix = matrix / np.where(lengths == 0, 1, lengths)
        return matrix

    batch_encode = encode