import os
import json
import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.utils.data import Dataset, DataLoader
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import Whitespace
from pathlib import Path
import argparse

class LightweightGPT(nn.Module):
    def __init__(self, vocab_size, block_size, n_embd, n_head, n_layer):
        super().__init__()
        self.block_size = block_size
        self.token_embedding = nn.Embedding(vocab_size, n_embd)
        self.position_embedding = nn.Embedding(block_size, n_embd)
        
        self.blocks = nn.ModuleList([
            nn.TransformerDecoderLayer(
                d_model=n_embd,
                nhead=n_head,
                dim_feedforward=4 * n_embd,
                dropout=0.1,
                activation='gelu',
                batch_first=True,
                norm_first=True
            )
            for _ in range(n_layer)
        ])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        device = idx.device
        causal_mask = torch.triu(torch.ones(T, T, device=device, dtype=torch.bool), diagonal=1)

        token_emb = self.token_embedding(idx)
        pos = torch.arange(0, T, dtype=torch.long, device=device)
        pos_emb = self.position_embedding(pos)
        
        x = token_emb + pos_emb

        for block in self.blocks:
            x = block(x, x, tgt_mask=causal_mask)

        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-1
            )
        return logits, loss

    def generate(self, idx, max_new_tokens, temperature=0.8, top_k=50, stop_token=None):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]
            logits = logits / temperature

            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')

            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)

            if stop_token is not None and idx_next.item() == stop_token:
                break

            idx = torch.cat((idx, idx_next), dim=1)
            
        return idx

class ConversationDataset(Dataset):
    def __init__(self, tokens, block_size, end_token_id):
        self.end_token = end_token_id
        self.block_size = block_size
        self.segments = []
        current_start = 0
        for i, token in enumerate(tokens):
            if token == end_token_id:
                segment = tokens[current_start:i+1]
                if len(segment) < block_size + 1:
                    padding = [end_token_id] * (block_size + 1 - len(segment))
                    segment.extend(padding)
                self.segments.append(segment)
                current_start = i + 1
        print(f"Created {len(self.segments)} conversation segments.")

    def __len__(self):
        return len(self.segments)

    def __getitem__(self, idx):
        segment = self.segments[idx]
        start_pos = torch.randint(0, max(1, len(segment) - self.block_size), (1,)).item()
        chunk = segment[start_pos:start_pos + self.block_size + 1]
        
        x = torch.tensor(chunk[:-1], dtype=torch.long)
        y = torch.tensor(chunk[1:], dtype=torch.long)
        return x, y

class AIBuilder:
    def __init__(self, model_name: str):
        self.model_name = model_name
        self.output_folder = model_name.replace(" ", "_").lower()
        self.device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Using device: {self.device}")
        
        self.model_config = {
            "block_size": 128,
            "n_embd": 128,
            "n_head": 4,
            "n_layer": 4,
            "vocab_size": 8000,
            "batch_size": 8,
            "grad_accum": 4,
            "max_epochs": 3,
        }

    def _build_tokenizer(self, training_data: str):
        tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
        tokenizer.pre_tokenizer = Whitespace()
        trainer = BpeTrainer(
            special_tokens=["[UNK]", "[PAD]", "user:", "ai:", "<|endoftext|>"],
            vocab_size=self.model_config["vocab_size"]
        )
        tokenizer.train_from_iterator(self._get_text_iterator(training_data), trainer)
        return tokenizer

    def _get_text_iterator(self, text, chunk_size=1000):
        for i in range(0, len(text), chunk_size):
            yield text[i:i + chunk_size]

    def _prepare_dataloader(self, tokenizer, text):
        tokens = tokenizer.encode(text).ids
        end_token_id = tokenizer.token_to_id("<|endoftext|>")
        dataset = ConversationDataset(tokens, self.model_config["block_size"], end_token_id)
        
        def collate_fn(batch):
            xs, ys = zip(*batch)
            return torch.stack(xs), torch.stack(ys)

        return DataLoader(dataset, batch_size=self.model_config["batch_size"], shuffle=True, collate_fn=collate_fn)

    def train(self, training_data: str):
        os.makedirs(self.output_folder, exist_ok=True)
        
        print("Building and saving tokenizer...")
        tokenizer = self._build_tokenizer(training_data)
        tokenizer.save(os.path.join(self.output_folder, "tokenizer.json"))
        
        print("Saving configuration file...")
        self._save_config(tokenizer) # MOVED HERE
        
        print("Preparing data for training...")
        dataloader = self._prepare_dataloader(tokenizer, training_data)
        
        model = LightweightGPT(
            vocab_size=tokenizer.get_vocab_size(),
            block_size=self.model_config["block_size"],
            n_embd=self.model_config["n_embd"],
            n_head=self.model_config["n_head"],
            n_layer=self.model_config["n_layer"]
        ).to(self.device)
        
        optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
        model_path = os.path.join(self.output_folder, "model.pt")
        
        print("\n--- Starting Model Training ---")
        model.train()
        best_loss = float('inf')
        
        for epoch in range(self.model_config["max_epochs"]):
            optimizer.zero_grad()
            for batch_idx, (x, y) in enumerate(dataloader):
                x, y = x.to(self.device), y.to(self.device)
                _, loss = model(x, y)
                
                loss = loss / self.model_config["grad_accum"]
                loss.backward()
                
                if (batch_idx + 1) % self.model_config["grad_accum"] == 0:
                    optimizer.step()
                    optimizer.zero_grad()
                
                current_loss = loss.detach().item() * self.model_config["grad_accum"]
                
                if batch_idx % 50 == 0:
                    print(f"Epoch {epoch+1} | Batch {batch_idx} | Loss: {current_loss:.4f}")
                
                if current_loss < best_loss:
                    best_loss = current_loss
                    torch.save(model.state_dict(), model_path)
                    print(f"🎉 New best model saved with loss: {best_loss:.4f}")
        
        print(f"✅ Training complete. Final best loss: {best_loss:.4f}")

    def _save_config(self, tokenizer):
        config = {
            "model_name": self.model_name,
            **self.model_config,
            "vocab_size": tokenizer.get_vocab_size(),
            "end_token_id": tokenizer.token_to_id("<|endoftext|>")
        }
        with open(os.path.join(self.output_folder, "config.json"), "w") as f:
            json.dump(config, f, indent=2)
        print(f"Configuration saved to {os.path.join(self.output_folder, 'config.json')}")

class ChatInterface:
    def __init__(self, model_dir="aglm"):
        self.model_dir = Path(model_dir)
        self.device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
        self.load_model()
        
    def load_model(self):
        with open(self.model_dir / "config.json", "r") as f:
            self.config = json.load(f)
        
        self.tokenizer = Tokenizer.from_file(str(self.model_dir / "tokenizer.json"))
        self.end_token_id = self.config.get("end_token_id")
        
        self.model = LightweightGPT(
            vocab_size=self.config["vocab_size"],
            block_size=self.config["block_size"],
            n_embd=self.config["n_embd"],
            n_head=self.config["n_head"],
            n_layer=self.config["n_layer"]
        ).to(self.device)
        
        self.model.load_state_dict(torch.load(self.model_dir / "model.pt", map_location=self.device))
        self.model.eval()
        print("✅ Model loaded successfully!")
    
    def chat(self):
        print("\n===== AI Assistant Ready =====")
        print("Type 'quit' or 'exit' to end the chat.\n")
        
        while True:
            user_input = input("user: ")
            if user_input.lower() in ["quit", "exit"]:
                break
                
            prompt = f"user: {user_input}\nai:"
            input_ids = self.tokenizer.encode(prompt).ids
            input_tensor = torch.tensor([input_ids], dtype=torch.long, device=self.device)
            
            with torch.no_grad():
                output_ids = self.model.generate(
                    input_tensor,
                    max_new_tokens=150,
                    temperature=0.7,
                    top_k=40,
                    stop_token=self.end_token_id
                )
            
            response_ids = output_ids[0, len(input_ids):].tolist()
            response = self.tokenizer.decode(response_ids)
            response = response.replace("<|endoftext|>", "").strip()
            
            print(f"ai: {response}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train or chat with an AgLM model.")
    parser.add_argument('action', choices=['train', 'chat'], nargs='?', default='train', help="Choose 'train' (default) or 'chat'.")
    args = parser.parse_args()

    model_folder = "aglm"

    if args.action == 'train':
        print("--- Starting Setup for AgLM ---")
        builder = AIBuilder("AgLM")
        try:
            with open("train.txt", "r", encoding="utf-8") as f:
                data = f.read()
            builder.train(data)
            print("\n✅ Training finished. You can now run with the 'chat' argument.")
            print(f"To chat, run: python {os.path.basename(__file__)} chat")
        except FileNotFoundError:
            print("\nERROR: train.txt not found. Please create train.txt with your conversational data to train the model.")
    
    elif args.action == 'chat':
        print("--- Starting Chat Interface for AgLM ---")
        if os.path.exists(model_folder) and os.path.exists(os.path.join(model_folder, "model.pt")):
            chat_bot = ChatInterface(model_dir=model_folder)
            chat_bot.chat()
        else:
            print(f"\nERROR: Model directory '{model_folder}' not found. Please run training first.")