import json
import random
from datetime import datetime
from math import ceil
from pathlib import Path
from typing import Any, Optional
from uuid import uuid4

import trio
from openai import OpenAI
from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    PrivateAttr,
)

from openiziai.schemas import DataDict
from openiziai.task import Task
from openiziai.utils import exponential_backoff


class TrainDataTool(BaseModel):
    """Cria dados preparados para fine tuning.

    Cria os dados para construir um fine tuning de um modelo LLM centrado
    em uma task contendo: backstory, role e o goal.
    """

    client: OpenAI = Field(default=None, description='Client da OpenAI.')
    data: DataDict = Field(
        description='Dados utilizados para construir o dado de treino.',
    )
    task: Task = Field(
        description='Descrição da task que o modelo deverá executar.',
    )
    model: str = Field(
        default='gpt-3.5-turbo-125',
        description='Modelo GPT que criará o dado de fine tuning.',
    )

    _n_examples: int = PrivateAttr(default=None)
    _n_batch: int = PrivateAttr(default=None)
    _root: Path = PrivateAttr(default_factory=Path.cwd)
    _train_data_dir: Path
    _template: str
    _hash: str = PrivateAttr(default_factory=uuid4)

    model_config = ConfigDict(arbitrary_types_allowed=True)

    def __init__(self, **data: Any) -> None:
        """Cria um novo modelo analisando e validando o input.

        Args:
            client (OpenAI): Client da OpenAI.
            data (DataDict): Dados utilizados para construir o dado de treino.
            task (Task): Descrição da task que o modelo treinado irá executar.
            model (str): Modelo GPT usado para criar os dados de treino.
                Padrão gpt-3.5-turbo-125.
        """
        super().__init__(**data)
        self._template = """You are generating data which will be used to train a machine learning model.
        To describe how the model will be trained, you will be given a high-level description in a dict with:
        ```{{'backstory': $backstory_goes_here, 'role': $role_goes_here, 'goal': $goal_goes_here, 'data': $data_goes_here}}
        From that, you will generate data samples, each with a prompt/response pair. You must return the sample in jsonl format with the exact informations:
        ```{{'prompt': $prompt_goes_here, 'response': $response_goes_here}}```
        Only one prompt/response pair should be generated by turn. For each turn, make the example slightly more complex than the last, while ensuring diversity.
        Here is the high-level description to generate the samples:
        {description}
        """  # noqa

        self._train_data_dir = self._create_dir_if_not_exist(
            self._root / 'data' / 'train'
        )

    @property
    def id(self) -> str:
        """Hash da classe."""
        return self._hash

    @property
    def train_data_dir(self) -> str:
        """Diretório que os dados de treino serão salvos."""
        return str(self._train_data_dir)

    @staticmethod
    def _create_dir_if_not_exist(dir: Path) -> Path:
        dir.mkdir(parents=True, exist_ok=True)

        return dir

    @exponential_backoff()
    async def create_examples(  # noqa
        self,
        n_examples: int,
        temperature: float,
        max_tokens: int,
        max_context_length: int,
        sender: trio.abc.SendChannel[list[dict[str, str]]],
    ) -> None:
        """Cria exemplos com o par: prompt/response.

        Utilizando um modelo GPT, cria exemplos com o par prompt/response no
            formato de json.

        Args:
            n_examples (int): Número de exemplos que serão criados.
            temperature (float): Temperatura que controla a criatividade ao
                criar os exemplos.
            max_tokens (int): Máximo de tokens que deve conter nas respostas.
                Valores maiores trarão prompts e responses maiores porém
                terá maior custo.
            max_context_length (int): Quantidade de exemplos que devem ser
                usados como contexto ao criar o próximo.
            sender (SendChannel): Canal de Pub par aplicar o Pub/Sub.
        """
        description = dict(
            backstory=self.task.backstory,
            role=self.task.role,
            goal=self.task.goal,
            **self.data,
        )
        prompt = self._template.format(description=description)

        short_system = {'role': 'system', 'content': self.task.short_backstory}

        examples = []

        for _ in range(n_examples):
            context = (
                [
                    {'role': 'assistant', 'content': example}
                    for example in examples
                ]
                if len(examples) <= max_context_length
                else [
                    {'role': 'assistant', 'content': example}
                    for example in random.sample(examples, 8)
                ]
            )

            messages = [
                {
                    'role': 'system',
                    'content': prompt,
                }
            ] + context

            result = self.client.chat.completions.create(
                model=self.model,
                messages=messages,  # pyright: ignore
                temperature=temperature,
                max_tokens=max_tokens,
            )

            _response = result.choices[0].message.content
            if not _response:
                print('Algo de errado aconteceu ao criar o prompt.')
                continue
            response = json.loads(_response)
            example = {
                'messages': [
                    short_system,
                    {'role': 'user', 'content': response.get('prompt')},
                    {
                        'role': 'assistant',
                        'content': response.get('response'),
                    },
                ]
            }

            examples.append(example)
        async with sender:
            await sender.send(examples)

    async def create_train_file(self, receiver: trio.abc.ReceiveChannel):
        """Salva os exemplos gerados em um arquivo jsonl.

        Método Sub do sistema Pub/Sub.

        Args:
            receiver (ReceiveChannel): Canal de receive do sistema de Pub/Sub.
        """
        self._file = (
            self._train_data_dir
            / f'train_{self.id}_{datetime.now().strftime("%Y%m%d")}.jsonl'
        )

        if self._file.exists():
            self._file.unlink()

        async with receiver:
            async for results in receiver:
                async with await trio.open_file(
                    self._file, 'a', encoding='utf-8'
                ) as file:
                    for result in results:
                        await file.write(json.dumps(result) + '\n')

    async def create_train_data(  # noqa
        self,
        n_examples: int,
        n_batch: int,
        temperature: float = 0.5,
        max_tokens: int = 1000,
        max_context_length: int = 8,
    ) -> str:
        """Cria os dados de treino.

        Aplica o sistema de Pub/Sub para criar os exemplos e salvar em um jsonl
            a medida que os batches de treino ficam prontos.

        Args:
            n_examples (int): Número de exemplos que devem ser criados.
            n_batch (int): Número de batches para serem executados de forma
                concorrente.
            temperature (float): Nível de criatividade do modelo ao criar os
                exemplos. Padrão 0.5.
            max_tokens (int): Máximo de tokens que deve conter nas respostas.
                Valores maiores trarão prompts e responses maiores porém
                terá maior custo.
            max_context_length (int): Quantidade de exemplos que devem ser
                usados como contexto ao criar o próximo.

        Returns:
            str: Nome do arquivo JSONL criado.
        """
        sender, receiver = trio.open_memory_channel(0)
        self._n_batch = n_batch
        _n_examples = ceil(n_examples / self._n_batch)
        self._n_examples = _n_examples * self._n_batch

        async with trio.open_nursery() as nursery:
            async with sender, receiver:
                for _ in range(self._n_batch):
                    nursery.start_soon(
                        self.create_examples,
                        _n_examples,
                        temperature,
                        max_tokens,
                        max_context_length,
                        sender.clone(),
                    )
                    nursery.start_soon(
                        self.create_train_file,
                        receiver.clone(),
                    )

        return str(self._file)

    @property
    def n_examples(self) -> Optional[int]:
        """Número de exemplos criados."""
        if not self._n_examples:
            print(
                'Você deve criar um dataset de treino com `create_train_data`.'
            )
            return None
        return self._n_examples

    @property
    def n_batch(self) -> Optional[int]:
        """Número de batches utilizados."""
        if not getattr(self, '_n_batch'):
            print(
                'Você deve criar um dataset de treino com `create_train_data`.'
            )
            return None
        return self._n_batch

    def execute(  # noqa
        self,
        n_examples: int,
        n_batch: int,
        temperature: float = 0.5,
        max_tokens: int = 1000,
        max_context_length: int = 8,
    ) -> str:
        """Executa todo o pipeline para criar os dados de treino.
        Args:
            n_examples (int): Número de exemplos que devem ser criados.
            n_batch (int): Número de batches para serem executados de forma
                concorrente.
            temperature (float): Nível de criatividade do modelo ao criar os
                exemplos. Padrão 0.5.
            max_tokens (int): Máximo de tokens que deve conter nas respostas.
                Valores maiores trarão prompts e responses maiores porém
                terá maior custo.
            max_context_length (int): Quantidade de exemplos que devem ser
                usados como contexto ao criar o próximo.

        Returns:
            str: Nome do arquivo JSONL criado.
        """  # noqa
        file = trio.run(
            self.create_train_data,
            n_examples,
            n_batch,
            temperature,
            max_tokens,
            max_context_length,
        )

        return file

    @property
    def file(self) -> Optional[str]:
        """Nome do arquivo com os dados de treino."""
        if not hasattr(self, '_file'):
            print('Nenhum dado de treino foi criado com esta instância.')
            return None
        return str(self._file)

    def __repr__(self) -> str:
        return (
            'TrainDataTool('
            f'task={self.task}, '
            f'n_examples={self.n_examples}, '
            f'n_batch={self.n_batch}, '
            f'file={self.file}'
            ')'
        )
