"""Llama generation module."""

from __future__ import annotations

import warnings

from copy import copy
from typing import Any

import instructor
import openai
import torch

from pydantic import BaseModel
from typeguard import typechecked

from rago._optional import require_dependency
from rago.generation.base import GenerationBase
from rago.generation.openai import OpenAIGen


@typechecked
class LlamaGen(GenerationBase):
    """Llama Generation class."""

    default_model_name: str = 'meta-llama/Llama-3.2-1B'
    default_temperature: float = 0.5
    default_output_max_length: int = 500
    default_api_params = {
        'top_p': 1.0,
        'num_return_sequences': 1,
    }

    def _load_optional_modules(self) -> None:
        self._langdetect = require_dependency(
            'langdetect',
            extra='langdetect',
            context='LangDetect',
        )
        self._transformers = require_dependency(
            'transformers',
            extra='transformers',
            context='transformers',
        )

        self._detect = self._langdetect.detect
        self._AutoModelForCausalLM = self._transformers.AutoModelForCausalLM
        self._AutoTokenizer = self._transformers.AutoTokenizer
        self._pipeline = self._transformers.pipeline

    def _validate(self) -> None:
        """Raise an error if the initial parameters are not valid."""
        if not self.model_name.startswith('meta-llama/'):
            raise Exception(
                f'The given model name {self.model_name} is not provided '
                'by meta.'
            )

        if self.structured_output:
            warnings.warn(
                'Structured output is not supported yet in '
                f'{self.__class__.__name__}.'
            )

    def _setup(self) -> None:
        """Set up the object with the initial parameters."""
        self.tokenizer = self._AutoTokenizer.from_pretrained(
            self.model_name, token=self.api_key
        )

        self.model = self._AutoModelForCausalLM.from_pretrained(
            self.model_name,
            token=self.api_key,
            torch_dtype=torch.float16
            if self.device_name == 'cuda'
            else torch.float32,
        )

        self.generator = self._pipeline(
            'text-generation',
            model=self.model,
            tokenizer=self.tokenizer,
            device=0 if self.device_name == 'cuda' else -1,
        )

    def generate(self, query: str, data: list[str]) -> str:
        """Generate text using Llama model with language support."""
        input_text = self._format_prompt(query, data)

        # Detect and set the language code for multilingual models (optional)
        language = str(self._detect(query)) or 'en'
        self.tokenizer.lang_code = language

        api_params = (
            self.api_params if self.api_params else self.default_api_params
        )

        # Generate the response with adjusted parameters

        model_params = dict(
            text_inputs=input_text,
            max_new_tokens=self.output_max_length,
            do_sample=True,
            temperature=self.temperature,
            eos_token_id=self.tokenizer.eos_token_id,
            **api_params,
        )
        response = self.generator(**model_params)

        # self.logs['model_params'] = model_params

        # Extract and return the answer only
        answer = str(response[0].get('generated_text', ''))
        # Strip off any redundant text after the answer itself
        return answer.split('Answer:')[-1].strip()


@typechecked
class OllamaGen(GenerationBase):
    """Ollama Generation class for local inference via ollama-python."""

    default_model_name = 'llama3.2:1b'
    default_temperature: float = 0.5
    default_output_max_length: int = 500
    default_api_params: dict[str, Any] = {
        'base_url': 'http://localhost:11434/'
    }

    def _load_optional_modules(self) -> None:
        self._ollama = require_dependency(
            'ollama',
            extra='ollama',
            context='Ollama',
        )
        self._Ollama = self._ollama.Client

    def _setup(self) -> None:
        """Instantiate the Ollama client."""
        self.api_params = copy(
            self.api_params if self.api_params else self.default_api_params
        )
        base_url = self.api_params.pop('base_url')

        self.model = self._Ollama(
            host=base_url, headers={'x-some-header': 'some-value'}
        )

    def generate(self, query: str, data: list[str]) -> str | BaseModel:
        """
        Generate text by sending a prompt to the local Ollama model.

        Parameters
        ----------
        query : str
            The user query.
        data : list[str]
            Augmented context strings.

        Returns
        -------
        str
            The generated response text.
        """
        input_text = self._format_prompt(query, data)

        messages = []
        if self.system_message:
            messages.append({'role': 'system', 'content': self.system_message})
        messages.append({'role': 'user', 'content': input_text})

        request_params = copy(self.api_params or {})
        options = copy(request_params.pop('options', {}))
        options.setdefault('temperature', self.temperature)
        options.setdefault('num_predict', self.output_max_length)

        params = {
            'model': self.model_name,
            'messages': messages,
            'options': options,
            **request_params,
        }
        response = self.model.chat(**params)
        return str(response.message.content).strip()


@typechecked
class OllamaOpenAIGen(OpenAIGen):
    """OllamaGen via the Ollama Python client."""

    default_model_name = 'llama3.2:1b'
    default_api_params: dict[str, Any] = {
        'base_url': 'http://localhost:11434/v1'
    }

    def _setup(self) -> None:
        self.api_params = copy(
            self.api_params if self.api_params else self.default_api_params
        )
        base_url = self.api_params.pop('base_url')

        model = openai.OpenAI(
            api_key='nokey',
            base_url=base_url,
        )

        self.model = (
            instructor.from_openai(model, mode=instructor.Mode.JSON)
            if self.structured_output
            else model
        )
