import torch
from transformers import PreTrainedTokenizerFast
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders, processors
import json
import os


def train_tokenizer(corpus_files, vocab_size=32768, output_path="fsi_edge_tokenizer"):
    """
    Train a BPE tokenizer optimized for code + NLP.
    Uses byte-level BPE with special tokens for code structure.
    """
    tokenizer = Tokenizer(models.BPE())
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tokenizer.decoder = decoders.ByteLevel()
    tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)
    
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=[
            "<unk>", "<s>", "</s>", "<pad>", "<|begin_of_text|>",
            "<|end_of_text|>", "<|code|>", "<|nl|>", "<|py|>", "<|js|>",
            "<|java|>", "<|cpp|>", "<|go|>", "<|rust|>", "<|sql|>",
            "<|explain|>", "<|debug|>", "<|solve|>", "<|test|>",
            "<|exec|>", "<|trace|>", "<|ast|>", "<|scope_start|>", "<|scope_end|>",
            "<|thought|>", "<|answer|>"
        ],
        min_frequency=2,
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
    )
    
    tokenizer.train(corpus_files, trainer)
    tokenizer.save(f"{output_path}/tokenizer.json")
    
    # Wrap as HuggingFace tokenizer
    hf_tokenizer = PreTrainedTokenizerFast(
        tokenizer_object=tokenizer,
        unk_token="<unk>",
        bos_token="<s>",
        eos_token="</s>",
        pad_token="<pad>",
    )
    hf_tokenizer.save_pretrained(output_path)
    return hf_tokenizer


class CodeDataset(torch.utils.data.Dataset):
    """
    Dataset for code + NLP training with AST annotations.
    Generates on-the-fly structural features for the model.
    """
    
    STRUCTURAL_TOKENS = {
        'import': 1, 'function_def': 2, 'class_def': 3, 'if': 4, 'elif': 5,
        'else': 6, 'for': 7, 'while': 8, 'try': 9, 'except': 10,
        'return': 11, 'assignment': 12, 'call': 13, 'comment': 14,
        'string': 15, 'number': 16, 'operator': 17, 'delimiter': 18,
        'indent': 19, 'dedent': 20, 'decorator': 21, 'lambda': 22,
        'with': 23, 'async': 24, 'await': 25, 'break': 26, 'continue': 27,
        'raise': 28, 'yield': 29, 'assert': 30, 'global': 31, 'nonlocal': 32,
    }
    
    def __init__(self, data_path, tokenizer_path, max_length=8192, split='train'):
        self.max_length = max_length
        self.samples = []
        
        if isinstance(tokenizer_path, str):
            from transformers import PreTrainedTokenizerFast
            self.tokenizer = PreTrainedTokenizerFast.from_pretrained(tokenizer_path)
        else:
            self.tokenizer = tokenizer_path
        
        self._load_data(data_path)
    
    def _load_data(self, data_path):
        """Load code samples from directory or single file."""
        if os.path.isfile(data_path):
            with open(data_path, 'r') as f:
                for line in f:
                    if line.strip():
                        self.samples.append(json.loads(line))
        elif os.path.isdir(data_path):
            for root, _, files in os.walk(data_path):
                for fname in files:
                    if fname.endswith('.jsonl'):
                        fpath = os.path.join(root, fname)
                        with open(fpath, 'r') as f:
                            for line in f:
                                if line.strip():
                                    self.samples.append(json.loads(line))
    
    def _compute_ast_depth(self, line):
        """Estimate AST depth from indentation."""
        depth = 0
        depths = []
        for ch in line:
            if ch in '({[':
                depth += 1
            elif ch in ')}]':
                depth = max(0, depth - 1)
            depths.append(min(depth, 31))
        if not depths:
            depths = [0]
        return depths
    
    def _detect_structural_tokens(self, tokens):
        """Detect code structure boundaries."""
        ast_types = []
        for tok in tokens:
            tok_lower = tok.lower().strip()
            found = False
            for keyword, idx in self.STRUCTURAL_TOKENS.items():
                if tok_lower == keyword:
                    ast_types.append(idx)
                    found = True
                    break
            if not found:
                ast_types.append(0)
        return ast_types
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        sample = self.samples[idx]
        code = sample.get('code', sample.get('text', ''))
        
        encoded = self.tokenizer(
            code,
            truncation=True,
            max_length=self.max_length,
            padding=False,
            return_tensors=None,
        )
        
        input_ids = encoded['input_ids']
        if len(input_ids) > self.max_length:
            input_ids = input_ids[:self.max_length]
        
        # Decode tokens for structural analysis
        tokens = [self.tokenizer.decode([tid]) for tid in input_ids]
        ast_types = self._detect_structural_tokens(tokens)
        ast_depths = self._compute_ast_depth(code[:len(input_ids)])
        
        # Pad to max_length
        pad_len = self.max_length - len(input_ids)
        if pad_len > 0:
            input_ids = input_ids + [self.tokenizer.pad_token_id] * pad_len
            ast_types = ast_types + [0] * pad_len
            ast_depths = ast_depths + [0] * pad_len
        else:
            input_ids = input_ids[:self.max_length]
            ast_types = ast_types[:self.max_length]
            ast_depths = ast_depths[:self.max_length]
        
        return {
            'input_ids': torch.tensor(input_ids, dtype=torch.long),
            'labels': torch.tensor(input_ids, dtype=torch.long),
            'ast_types': torch.tensor(ast_types[:self.max_length], dtype=torch.long),
            'ast_depths': torch.tensor(ast_depths[:self.max_length], dtype=torch.long),
            'attention_mask': torch.tensor(
                [1 if tid != self.tokenizer.pad_token_id else 0 for tid in input_ids],
                dtype=torch.long),
        }


def collate_fn(batch):
    """Collate batch, stacking tensors."""
    keys = batch[0].keys()
    result = {}
    for k in keys:
        result[k] = torch.stack([b[k] for b in batch], dim=0)
    return result