"""Local inference engine wrapper for Lean code generation."""

from __future__ import annotations

import os
import re
import time
from dataclasses import dataclass
from typing import Any

from rc.inference.catalog import LLAMA_CPP_RUNTIME, MLX_RUNTIME, get_model_spec
from rc.inference.flops import compute_credit_subunits, compute_flops, compute_mflops, compute_tflops
from rc.inference.prompts import build_task_prompt

LEAN_TASK_TOKEN_LIMITS = {
    "lean_proof_gen": 128,
    "proof_attempt": 96,
    "proof_repair": 96,
    "conjecture_formalization": 128,
}

GENERATION_STOP_MARKERS = (
    "\n###",
    "\n---",
    "\n```",
    "\nExplanation",
    "\nProof Sketch",
    "\nComplete Lean",
    "\nRequired template:",
    "\nTarget template:",
)


@dataclass
class InferenceResult:
    lean_code: str
    raw_output: str
    prompt_tokens: int
    generated_tokens: int
    model_params: int
    prefill_eval_ms: float
    decode_eval_ms: float
    total_eval_ms: float
    flops: int
    mflops: int
    tflops: float
    credit_subunits: int


class InferenceEngine:
    """Wrap a supported local runtime for Lean proof generation."""

    def __init__(
        self,
        model: str,
        runtime: str | None = None,
        model_params: int | None = None,
        temperature: float = 0.6,
        top_p: float = 0.95,
        max_tokens: int = 4096,
        n_ctx: int = 2048,
        n_threads: int | None = None,
        n_gpu_layers: int | None = None,
    ) -> None:
        self.model_name = model
        self.temperature = temperature
        self.top_p = top_p
        self.max_tokens = max_tokens
        spec = get_model_spec(model)
        self.runtime = runtime or (spec.runtime if spec else LLAMA_CPP_RUNTIME)
        self.model_params = model_params or (spec.param_count if spec else 7_000_000_000)

        if self.runtime == MLX_RUNTIME:
            self._backend = _MlxBackend(
                model=model,
                temperature=temperature,
                top_p=top_p,
                max_tokens=max_tokens,
                model_params=self.model_params,
            )
            self.n_threads = 0
            return

        self.n_threads = n_threads or suggest_thread_count()
        self._backend = _LlamaCppBackend(
            model=model,
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens,
            n_ctx=n_ctx,
            n_threads=self.n_threads,
            n_gpu_layers=n_gpu_layers,
            model_params=self.model_params,
        )

    def generate_lean(
        self,
        problem: str,
        task_type: str = "lean_proof_gen",
        context: dict[str, Any] | None = None,
    ) -> InferenceResult:
        """Generate Lean 4 code for a mathematical problem.

        Returns an InferenceResult with the generated code, actual token counts,
        runtime timings, and token-accounted compute usage.
        """
        prompt = build_task_prompt(task_type=task_type, problem=problem, context=context)
        return self._backend.generate(prompt, task_type=task_type)


class _LlamaCppBackend:
    def __init__(
        self,
        model: str,
        temperature: float,
        top_p: float,
        max_tokens: int,
        n_ctx: int,
        n_threads: int,
        n_gpu_layers: int | None,
        model_params: int,
    ) -> None:
        try:
            from llama_cpp import Llama, llama_perf_context, llama_perf_context_reset
        except ImportError:
            raise RuntimeError(
                "llama.cpp is not installed. Install the CLI dependencies again."
            )

        resolved_gpu_layers = (
            n_gpu_layers if n_gpu_layers is not None else _default_gpu_layers()
        )
        self.temperature = temperature
        self.top_p = top_p
        self.max_tokens = max_tokens
        self.model_params = model_params
        self._perf_context = llama_perf_context
        self._perf_context_reset = llama_perf_context_reset
        self.llm = Llama(
            model_path=model,
            n_ctx=n_ctx,
            n_threads=n_threads,
            n_gpu_layers=resolved_gpu_layers,
            verbose=False,
        )

    def generate(self, prompt: str, task_type: str = "lean_proof_gen") -> InferenceResult:
        self._perf_context_reset(self.llm._ctx.ctx)
        resolved_max_tokens = _resolve_max_tokens(task_type, self.max_tokens)
        response = self.llm.create_completion(
            prompt=prompt,
            temperature=self.temperature,
            top_p=self.top_p,
            max_tokens=resolved_max_tokens,
            stop=list(GENERATION_STOP_MARKERS),
        )
        perf = self._perf_context(self.llm._ctx.ctx)
        generated_text = _truncate_at_stop_markers(response["choices"][0]["text"])
        usage = response.get("usage", {})
        prompt_tokens = int(
            usage.get("prompt_tokens") or len(self.llm.tokenize(prompt.encode("utf-8")))
        )
        generated_tokens = int(
            usage.get("completion_tokens")
            or len(self.llm.tokenize(generated_text.encode("utf-8")))
        )
        model_params = int(getattr(self.llm._model, "n_params", lambda: self.model_params)())
        prefill_eval_ms = float(perf.t_p_eval_ms)
        decode_eval_ms = float(perf.t_eval_ms)
        total_eval_ms = prefill_eval_ms + decode_eval_ms
        flops = compute_flops(
            model_params=model_params,
            prompt_tokens=prompt_tokens,
            generated_tokens=generated_tokens,
        )
        mflops = compute_mflops(
            model_params=model_params,
            prompt_tokens=prompt_tokens,
            generated_tokens=generated_tokens,
        )
        tflops = compute_tflops(
            model_params=model_params,
            prompt_tokens=prompt_tokens,
            generated_tokens=generated_tokens,
        )
        credit_subunits = compute_credit_subunits(
            model_params=model_params,
            prompt_tokens=prompt_tokens,
            generated_tokens=generated_tokens,
        )
        return InferenceResult(
            lean_code=_extract_lean_code(generated_text),
            raw_output=generated_text,
            prompt_tokens=prompt_tokens,
            generated_tokens=generated_tokens,
            model_params=model_params,
            prefill_eval_ms=prefill_eval_ms,
            decode_eval_ms=decode_eval_ms,
            total_eval_ms=total_eval_ms,
            flops=flops,
            mflops=mflops,
            tflops=tflops,
            credit_subunits=credit_subunits,
        )


class _MlxBackend:
    def __init__(
        self,
        model: str,
        temperature: float,
        top_p: float,
        max_tokens: int,
        model_params: int,
    ) -> None:
        try:
            from mlx_lm import load, stream_generate
            from mlx_lm.sample_utils import make_sampler
        except ImportError:
            raise RuntimeError(
                "mlx-lm is not installed. Install the CLI dependencies again."
            )

        self.temperature = temperature
        self.top_p = top_p
        self.max_tokens = max_tokens
        self.model_params = model_params
        self._stream_generate = stream_generate
        self._sampler = make_sampler(temp=temperature, top_p=top_p)
        self.model, self.tokenizer = load(model, lazy=False)

    def generate(self, prompt: str, task_type: str = "lean_proof_gen") -> InferenceResult:
        formatted_prompt = _format_mlx_prompt(self.tokenizer, prompt)
        pieces: list[str] = []
        last_response: Any | None = None
        started_at = time.perf_counter()
        resolved_max_tokens = _resolve_max_tokens(task_type, self.max_tokens)
        for response in self._stream_generate(
            self.model,
            self.tokenizer,
            prompt=formatted_prompt,
            max_tokens=resolved_max_tokens,
            sampler=self._sampler,
        ):
            pieces.append(response.text)
            last_response = response
            if _contains_stop_marker("".join(pieces)):
                break
        total_eval_ms = (time.perf_counter() - started_at) * 1000.0

        generated_text = _truncate_at_stop_markers("".join(pieces))
        prompt_tokens = int(getattr(last_response, "prompt_tokens", 0) or 0)
        generated_tokens = int(getattr(last_response, "generation_tokens", 0) or 0)
        prompt_tps = float(getattr(last_response, "prompt_tps", 0.0) or 0.0)
        generation_tps = float(getattr(last_response, "generation_tps", 0.0) or 0.0)
        prefill_eval_ms = (prompt_tokens / prompt_tps * 1000.0) if prompt_tps > 0 else 0.0
        decode_eval_ms = (
            generated_tokens / generation_tps * 1000.0 if generation_tps > 0 else 0.0
        )
        if total_eval_ms <= 0:
            total_eval_ms = prefill_eval_ms + decode_eval_ms

        flops = compute_flops(
            model_params=self.model_params,
            prompt_tokens=prompt_tokens,
            generated_tokens=generated_tokens,
        )
        mflops = compute_mflops(
            model_params=self.model_params,
            prompt_tokens=prompt_tokens,
            generated_tokens=generated_tokens,
        )
        tflops = compute_tflops(
            model_params=self.model_params,
            prompt_tokens=prompt_tokens,
            generated_tokens=generated_tokens,
        )
        credit_subunits = compute_credit_subunits(
            model_params=self.model_params,
            prompt_tokens=prompt_tokens,
            generated_tokens=generated_tokens,
        )
        return InferenceResult(
            lean_code=_extract_lean_code(generated_text),
            raw_output=generated_text,
            prompt_tokens=prompt_tokens,
            generated_tokens=generated_tokens,
            model_params=self.model_params,
            prefill_eval_ms=prefill_eval_ms,
            decode_eval_ms=decode_eval_ms,
            total_eval_ms=total_eval_ms,
            flops=flops,
            mflops=mflops,
            tflops=tflops,
            credit_subunits=credit_subunits,
        )


def _extract_lean_code(text: str) -> str:
    """Extract Lean code from model output, stripping markdown fences if present."""
    stripped = _truncate_at_stop_markers(text.strip())

    if "```lean4" in stripped:
        parts = stripped.split("```lean4", 1)
        if len(parts) > 1:
            code = parts[1].split("```", 1)[0]
            return code.strip()

    if "```lean" in stripped:
        parts = stripped.split("```lean", 1)
        if len(parts) > 1:
            code = parts[1].split("```", 1)[0]
            return code.strip()

    if "```" in stripped:
        parts = stripped.split("```", 1)
        if len(parts) > 1:
            code = parts[1].split("```", 1)[0]
            return code.strip()

    first_code_offset = _find_first_code_offset(stripped)
    if first_code_offset is not None:
        return stripped[first_code_offset:].strip()
    if stripped.startswith("###"):
        return ""

    return stripped


def _resolve_max_tokens(task_type: str, configured_max_tokens: int) -> int:
    limit = LEAN_TASK_TOKEN_LIMITS.get(task_type, configured_max_tokens)
    return max(16, min(configured_max_tokens, limit))


def _contains_stop_marker(text: str) -> bool:
    return any(marker in text for marker in GENERATION_STOP_MARKERS)


def _truncate_at_stop_markers(text: str) -> str:
    truncated = text
    for marker in GENERATION_STOP_MARKERS:
        marker_index = truncated.find(marker)
        if marker_index != -1:
            truncated = truncated[:marker_index]
    return truncated.rstrip()


def _find_first_code_offset(text: str) -> int | None:
    patterns = (
        r"(?m)^import Mathlib\b",
        r"(?m)^theorem\s+[A-Za-z_][\w']*",
        r"(?m)^lemma\s+[A-Za-z_][\w']*",
        r"(?m)^example\b",
        r"(?m)^def\s+[A-Za-z_][\w']*",
    )
    offsets = [
        match.start()
        for pattern in patterns
        if (match := re.search(pattern, text)) is not None
    ]
    if not offsets:
        return None
    return min(offsets)


def _format_mlx_prompt(tokenizer: Any, prompt: str) -> str:
    """Use the model chat template when available for instruction-tuned MLX models."""
    apply_chat_template = getattr(tokenizer, "apply_chat_template", None)
    if callable(apply_chat_template):
        try:
            return apply_chat_template(
                [{"role": "user", "content": prompt}],
                add_generation_prompt=True,
                tokenize=False,
            )
        except TypeError:
            pass
    return prompt


def _default_gpu_layers() -> int:
    """Use GPU offload where available, otherwise stay on CPU."""
    if platform := os.environ.get("RC_N_GPU_LAYERS"):
        try:
            return int(platform)
        except ValueError:
            return 0
    return -1


def suggest_thread_count(
    active_provider_count: int | None = None,
    cpu_count: int | None = None,
) -> int:
    """Split local CPU capacity across concurrently running providers."""
    detected_cpu_count = max(cpu_count or _detect_cpu_count(), 1)
    detected_provider_count = max(
        active_provider_count or _detect_active_provider_count(),
        1,
    )
    usable_threads = max(1, detected_cpu_count - 1)
    return max(1, usable_threads // detected_provider_count)


def _detect_cpu_count() -> int:
    try:
        import psutil

        return psutil.cpu_count(logical=False) or psutil.cpu_count() or os.cpu_count() or 4
    except ImportError:
        return os.cpu_count() or 4


def _detect_active_provider_count() -> int:
    try:
        import psutil
    except ImportError:
        return 1

    provider_count = 0
    for proc in psutil.process_iter(["cmdline"]):
        try:
            cmdline = " ".join(proc.info.get("cmdline") or [])
        except (psutil.Error, OSError):
            continue
        lowered = cmdline.lower()
        if "rc serve" in lowered and "python" in lowered:
            provider_count += 1

    return provider_count or 1
