from __future__ import annotations

import json
import math
from pathlib import Path

import torch
from datasets import Dataset
from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers
from transformers import GPT2Config, GPT2LMHeadModel, PreTrainedTokenizerFast

PROJECT_DIR = Path(__file__).resolve().parent
DATA_DIR = PROJECT_DIR / "data"
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "snip-0.4m-base"
TOKENIZER_DIR = ARTIFACT_DIR / "tokenizer"
CONTEXT_LENGTH = 128
VOCAB_SIZE = 512


def read_texts(path: Path) -> list[str]:
    with path.open("r", encoding="utf-8") as handle:
        return [json.loads(line)["text"] for line in handle if line.strip()]


def train_tokenizer(texts: list[str]) -> PreTrainedTokenizerFast:
    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tokenizer.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(
        vocab_size=VOCAB_SIZE,
        min_frequency=2,
        special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"],
        show_progress=True,
    )
    tokenizer.train_from_iterator(texts, trainer=trainer)
    wrapped = PreTrainedTokenizerFast(
        tokenizer_object=tokenizer,
        bos_token="<bos>",
        eos_token="<eos>",
        unk_token="<unk>",
        pad_token="<pad>",
        model_max_length=CONTEXT_LENGTH,
    )
    TOKENIZER_DIR.mkdir(parents=True, exist_ok=True)
    wrapped.save_pretrained(TOKENIZER_DIR)
    return wrapped


def build_tokenizer(train_texts: list[str] | None = None) -> PreTrainedTokenizerFast:
    tokenizer_file = TOKENIZER_DIR / "tokenizer.json"
    if tokenizer_file.exists():
        return PreTrainedTokenizerFast.from_pretrained(TOKENIZER_DIR)
    if not train_texts:
        raise FileNotFoundError("Tokenizer is missing and no training text was provided.")
    return train_tokenizer(train_texts)


def make_model(tokenizer: PreTrainedTokenizerFast) -> GPT2LMHeadModel:
    config = GPT2Config(
        vocab_size=len(tokenizer),
        n_positions=CONTEXT_LENGTH,
        n_ctx=CONTEXT_LENGTH,
        n_embd=96,
        n_layer=3,
        n_head=3,
        n_inner=384,
        activation_function="gelu_new",
        resid_pdrop=0.05,
        embd_pdrop=0.05,
        attn_pdrop=0.05,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
        tie_word_embeddings=True,
    )
    return GPT2LMHeadModel(config)


def texts_to_blocks(
    texts: list[str],
    tokenizer: PreTrainedTokenizerFast,
) -> Dataset:
    token_ids: list[int] = []
    eos = tokenizer.eos_token_id
    for text in texts:
        token_ids.extend(
            tokenizer.backend_tokenizer.encode(
                text,
                add_special_tokens=False,
            ).ids
        )
        token_ids.append(eos)
    usable = (len(token_ids) // CONTEXT_LENGTH) * CONTEXT_LENGTH
    blocks = [
        token_ids[index : index + CONTEXT_LENGTH]
        for index in range(0, usable, CONTEXT_LENGTH)
    ]
    return Dataset.from_dict(
        {
            "input_ids": blocks,
            "attention_mask": [[1] * CONTEXT_LENGTH for _ in blocks],
            "labels": [block.copy() for block in blocks],
        }
    )


def parameter_count(model: torch.nn.Module) -> int:
    return sum(parameter.numel() for parameter in model.parameters())


def perplexity(loss: float) -> float:
    return float(math.exp(min(loss, 20)))