import json
import os
import torch
from torch.utils.data import Dataset

class InstructionDataset(Dataset):
    """
    Dataset for Instruction Tuning (Q&A).
    Format: [INSTRUCTION] Question [RESPONSE] Answer <EOS>
    """
    def __init__(self, data_path, tokenizer, block_size):
        self.tokenizer = tokenizer
        self.block_size = block_size
        self.data = []
        
        try:
            import sys
            import os
            # Ensure the Rust .dll/.so is on path
            core_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'agent_core', 'target', 'release')
            if core_path not in sys.path:
                sys.path.append(core_path)

            import agent_core
            
            # Check if the attribute actually exists before calling to prevent unhandled AttributeError
            if hasattr(agent_core, 'parse_instruction_dataset_safe'):
                self.data = agent_core.parse_instruction_dataset_safe(data_path)
                print(f"Loaded {len(self.data)} instruction pairs via Rust Memory-Safe Core.")
            else:
                raise AttributeError(f"module 'agent_core' has no attribute 'parse_instruction_dataset_safe'")
                
        except (ImportError, AttributeError, Exception) as e:
             # The Rust extension handles all JSON loading, error checking, 
             # and parallel string formatting in memory-safe C-bindings natively
             print(f"Warning: Rust `agent_core` fallback triggered. Reason: {e}")
             print("Falling back to native Python parsing.")
             self._fallback_load(data_path)

    def _fallback_load(self, data_path):
        try:
            import json
            with open(data_path, 'r', encoding='utf-8') as f:
                raw_data = json.load(f)
                
            for item in raw_data:
                if 'instruction' in item and 'response' in item:
                    system_prompt = item.get('system', 'You are a helpful, smart AI assistant.')
                    thought = item.get('thought', '')
                    tools = item.get('tools', [])
                    
                    thought_str = f"[THOUGHT] {thought} " if thought else ""
                    tools_str = ""
                    for tool in tools:
                         tools_str += f"[TOOL_CALL] {tool.get('name')} [TOOL_ARG] {tool.get('args')} [TOOL_RESULT] {tool.get('result')} "
                    
                    text = f"[SYSTEM] {system_prompt} [USER] {item['instruction']} {thought_str}{tools_str}[RESPONSE] {item['response']} <EOS>"
                    self.data.append(text)
        except Exception as e:
            print(f"Error loading instruction data: {e}")
            self.data = []

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

    def __getitem__(self, idx):
        text = self.data[idx]
        
        # Tokenize
        # Note: We need a custom encode method that handles the special tags if they aren't in vocab
        # For now, we assume the tokenizer splits them or we add them as special tokens
        
        # Simple encoding for now, assuming tokenizer handles it or splits by space
        # Ideally, we should add [INSTRUCTION] and [RESPONSE] to tokenizer specials
        
        # For simplicity in this iteration, we'll just encode the text string
        # The tokenizer.encode needs to be robust
        
        # We will use the tokenizer's encode method
        # If tokenizer doesn't have the special tokens, it might split them. 
        # We should add them to the tokenizer in train.py
        
        ids = self.tokenizer.encode(text)
        
        # Pad or Truncate
        if len(ids) > self.block_size:
            ids = ids[:self.block_size]
        else:
            ids = ids + [self.tokenizer.word_to_id.get("<PAD>", 0)] * (self.block_size - len(ids))
            
        x = torch.tensor(ids[:-1], dtype=torch.long)
        y = torch.tensor(ids[1:], dtype=torch.long)
        
        return x, y