import torch
import torch.nn.functional as F
import os
from model.config import ModelConfig
from model.transformer import GPT
from model.tokenizer import AdvancedTokenizer

class InferenceEngine:
    def __init__(self, model_path='sail.pt'):
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.draft_model = None # Initialized if available
        
        if not os.path.exists(model_path):
            print(f"Warning: {model_path} not found. Agent will be uninitialized.")
            self.model = None
            return
            
        print(f"Loading model from {model_path}...")
        try:
            checkpoint = torch.load(model_path, map_location=self.device, weights_only=False)
            
            if isinstance(checkpoint, dict) and 'config' in checkpoint:
                self.config = checkpoint['config']
                self.config.device = self.device
                self.tokenizer = AdvancedTokenizer(vocab_size=self.config.vocab_size)
                if 'vocab' in checkpoint:
                    self.tokenizer.word_to_id = checkpoint['vocab']
                    self.tokenizer.id_to_word = {v: k for k, v in self.tokenizer.word_to_id.items()}
                state_dict = checkpoint['model_state_dict']
                # Strip _orig_mod. prefix if model was saved compiled
                state_dict = {k.replace('_orig_mod.', ''): v for k, v in state_dict.items()}
                self.model = GPT(self.config).to(self.device)
                self.model.load_state_dict(state_dict)
            else:
                import pickle
                with open('tokenizer.pkl', 'rb') as f:
                    self.tokenizer = pickle.load(f)
                self.config = ModelConfig()
                self.model = GPT(self.config).to(self.device)
                self.model.load_state_dict(checkpoint)

            self.model.eval()
            print("Model loaded successfully.")
        except Exception as e:
            print(f"Error loading model: {e}")
            self.model = None

    def speculative_generate(self, idx, max_new_tokens, K=4):
        """
        Speculative Decoding: Generate multiple tokens using a draft model, 
        then verify them in one pass with the big model.
        """
        if self.draft_model is None:
            return self.model.generate(idx, max_new_tokens)
            
        generated = 0
        while generated < max_new_tokens:
            T = idx.shape[1]
            draft_idx = idx.clone()
            for _ in range(min(K, max_new_tokens - generated)):
                logits, _ = self.draft_model(draft_idx)
                next_token = torch.multinomial(F.softmax(logits[:, -1, :], dim=-1), 1)
                draft_idx = torch.cat([draft_idx, next_token], dim=1)
            
            # Big model verifies in one pass
            full_logits, _ = self.model(draft_idx)
            
            # Simple acceptance logic (simplified for agentic use)
            idx = draft_idx
            generated += K
            if generated >= max_new_tokens: break
            
        return idx

    def chat(self, prompt, system_prompt="You are a smart, agentic AI assistant.", max_steps=3):
        if not self.model: return "Error: Model not loaded."
        
        from agent.tool_executor import parse_and_execute_tools
        current_text = f"[SYSTEM] {system_prompt} [USER] {prompt} [THOUGHT]"
        
        for step in range(max_steps):
            token_ids = self.tokenizer.encode(current_text)
            idx = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0).to(self.device)
            
            with torch.no_grad():
                out_idx = self.speculative_generate(idx, max_new_tokens=256)
            
            generated_text = self.tokenizer.decode(out_idx[0].tolist())
            
            # Tool Use
            tool_result = parse_and_execute_tools(generated_text)
            if tool_result:
                current_text = generated_text + " " + tool_result
                continue
            else:
                # Self-Critique
                if "[THOUGHT]" in generated_text and "error" in generated_text.lower():
                    current_text = generated_text + " [THOUGHT] I detected a potential error. I must correct it."
                    continue
                return generated_text
        return generated_text

    def generate(self, prompt, max_new_tokens=200, **kwargs):
        if not self.model: return "Model not loaded."
        token_ids = self.tokenizer.encode(prompt)
        idx = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0).to(self.device)
        with torch.no_grad():
            out_idx = self.model.generate(idx, max_new_tokens, **kwargs)
        return self.tokenizer.decode(out_idx[0].tolist())