"""LLM client abstraction for HIEROMEM agent."""

from __future__ import annotations

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


@dataclass
class LLMConfig:
    """Configuration for LLM client."""

    model: str = "gpt-4o-mini"
    temperature: float = 0.2
    max_tokens: int = 512
    max_retries: int = 3
    retry_delay: float = 1.0


class LLMClient:
    """LLM client with retry logic and provider abstraction."""

    def __init__(
        self,
        config: LLMConfig | None = None,
        logger: logging.Logger | None = None,
    ) -> None:
        self.config = config or LLMConfig()
        self.log = logger or logging.getLogger(self.__class__.__name__)
        self._client: Any = None
        self._pipeline: Any = None
        self._tokenizer: Any = None

    @property
    def model(self) -> str:
        return self.config.model

    @property
    def temperature(self) -> float:
        return self.config.temperature

    @property
    def max_tokens(self) -> int:
        return self.config.max_tokens

    @property
    def provider(self) -> str:
        """Infer provider from environment or model name."""
        env_provider = os.environ.get("HIEROMEM_LLM_PROVIDER", "").lower()
        if env_provider == "local_huggingface":
            return "local_huggingface"

        model = self.config.model.lower()
        if model.startswith("gpt") or model.startswith("o1"):
            return "openai"
        elif model.startswith("claude"):
            return "anthropic"
        elif model.startswith("gemini"):
            return "google"
        elif "/" in model:  # HuggingFace model format: org/model
            return "local_huggingface"
        else:
            return "openai"  # Default

    @property
    def client(self) -> Any:
        """Lazily initialize the OpenAI client."""
        if self._client is None:
            try:
                from openai import OpenAI

                self._client = OpenAI()
            except ImportError:
                raise ImportError(
                    "openai package is required. Install with: pip install openai"
                )
        return self._client

    def _get_hf_pipeline(self) -> Any:
        """Lazily initialize the HuggingFace pipeline."""
        if self._pipeline is None:
            try:
                from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
                import torch

                device = "cuda" if torch.cuda.is_available() else "cpu"
                self.log.info(f"Loading HuggingFace model {self.config.model} on {device}")

                self._tokenizer = AutoTokenizer.from_pretrained(self.config.model)
                model = AutoModelForCausalLM.from_pretrained(
                    self.config.model,
                    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
                    device_map="auto" if device == "cuda" else None,
                    low_cpu_mem_usage=True,
                )
                if device == "cpu":
                    model = model.to(device)

                self._pipeline = pipeline(
                    "text-generation",
                    model=model,
                    tokenizer=self._tokenizer,
                    device=None if device == "cuda" else 0 if device == "cpu" else device,
                )
            except ImportError:
                raise ImportError(
                    "transformers and torch are required for local models. "
                    "Install with: pip install transformers torch"
                )
        return self._pipeline

    def complete(self, prompt: str, system: str | None = None) -> str:
        """Generate a completion for the given prompt."""
        messages: list[dict[str, str]] = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        return self._call_with_retry(messages)

    def chat(self, messages: list[dict[str, str]]) -> str:
        """Generate a response for a chat conversation."""
        return self._call_with_retry(messages)

    def _call_with_retry(self, messages: list[dict[str, str]]) -> str:
        """Call the LLM with retry logic."""
        last_error: Exception | None = None

        for attempt in range(self.config.max_retries):
            try:
                if self.provider == "local_huggingface":
                    return self._call_huggingface(messages)
                else:
                    return self._call_openai(messages)

            except Exception as e:
                last_error = e
                self.log.warning(
                    f"LLM call failed (attempt {attempt + 1}/{self.config.max_retries}): {e}"
                )
                if attempt < self.config.max_retries - 1:
                    sleep_time = self.config.retry_delay * (2**attempt)
                    time.sleep(sleep_time)

        self.log.error(f"LLM call failed after {self.config.max_retries} attempts")
        raise last_error or RuntimeError("LLM call failed")

    def _call_openai(self, messages: list[dict[str, str]]) -> str:
        """Call OpenAI API."""
        response = self.client.chat.completions.create(
            model=self.config.model,
            messages=messages,
            temperature=self.config.temperature,
            max_tokens=self.config.max_tokens,
        )
        content = response.choices[0].message.content
        return content.strip() if content else ""

    def _call_huggingface(self, messages: list[dict[str, str]]) -> str:
        """Call local HuggingFace model."""
        pipe = self._get_hf_pipeline()

        # Format messages into a prompt
        if self._tokenizer is not None and hasattr(self._tokenizer, "apply_chat_template"):
            prompt = self._tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
        else:
            # Fallback: simple concatenation
            prompt_parts = []
            for msg in messages:
                role = msg.get("role", "user")
                content = msg.get("content", "")
                if role == "system":
                    prompt_parts.append(f"System: {content}")
                elif role == "user":
                    prompt_parts.append(f"User: {content}")
                elif role == "assistant":
                    prompt_parts.append(f"Assistant: {content}")
            prompt_parts.append("Assistant:")
            prompt = "\n".join(prompt_parts)

        outputs = pipe(
            prompt,
            max_new_tokens=self.config.max_tokens,
            temperature=self.config.temperature if self.config.temperature > 0 else None,
            do_sample=self.config.temperature > 0,
            pad_token_id=pipe.tokenizer.eos_token_id,
        )

        generated = outputs[0]["generated_text"]
        # Extract only the new content after the prompt
        if generated.startswith(prompt):
            generated = generated[len(prompt):]
        return generated.strip()
