import torch
import os
import torch.nn as nn
from typing import Optional, Dict, Any, Union, List, Generator
from transformers import AutoProcessor

from xorfice.optimize import LRUExpertCache, ManagedMoELayer, PagedKVCacheManager
from xorfice.models.xoron import XoronMultimodalModel
from xorfice.config import XoronConfig
from xorfice.models.llm.moe_llama import AuxLosslessMoELayer
from xorfice.memory import FlatFileMemoryManager
from xorfice.personalization import PersonalizationManager, VoiceManager
from xorfice.agent import AgentManager

import torchaudio
from PIL import Image
import numpy as np
import requests
import io
import tempfile
import cv2

class XoronEngine:
    """
    SOTA Inference Engine for Xoron-Dev.
    Fully optimized pipelines for Text, Image, Video, and Audio (S2S, I2I, V2V).
    """
    def __init__(
        self, 
        model_path: str, 
        max_vram_experts: int = 4,
        kv_cache_blocks: int = 256,
        block_size: int = 16,
        device: str = "cuda",
        **kwargs
    ):
        self.model_path = model_path
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        self.max_vram_experts = max_vram_experts
        self.kv_cache_blocks = kv_cache_blocks
        self.block_size = block_size
        self.expert_offload_threshold = kwargs.get('expert_offload_threshold', 0) # 0 means disabled
        self.max_offload_ratio = kwargs.get('max_offload_ratio', 0.5)
        
        # Initialize attributes
        self.config = None
        self.model = None
        self.processor = None
        self.expert_cache = None
        self.kv_manager = None
        
        # SOTA Personalization & Memory
        self.memory = FlatFileMemoryManager()
        self.personalization = PersonalizationManager()
        self.voice = VoiceManager()
        self.agent = AgentManager(workspace_dir=os.getcwd())
        
        # Ensure outputs directory exists for media files
        os.makedirs("outputs", exist_ok=True)
        
        self._load_model()
        
        if self.expert_offload_threshold > 0:
            self.enable_expert_offloading(
                cold_threshold=self.expert_offload_threshold, 
                max_offload_ratio=self.max_offload_ratio
            )
            
        self._apply_optimizations()

    def _load_model(self):
        print(f"[XoronEngine] Loading SOTA model from {self.model_path}...")
        
        # Determine if this is a local path or a HuggingFace repo ID
        _is_local_path = (
            self.model_path.startswith(".")
            or self.model_path.startswith("/")
            or self.model_path.startswith("~")
            or os.path.exists(self.model_path)
        )
        
        if _is_local_path:
            # Local path — it MUST exist
            resolved = os.path.expanduser(self.model_path)
            if not os.path.exists(resolved):
                raise FileNotFoundError(
                    f"Local model path does not exist: {resolved}\n"
                    f"  → Use --hf <repo_id> to download from HuggingFace Hub, or\n"
                    f"  → Use --weights <path> to point to local weights."
                )
            self.model_path = resolved
        elif not os.path.exists(self.model_path):
            # Looks like a HF repo ID — download it
            print(f"[XoronEngine] Downloading from HuggingFace Hub: {self.model_path}")
            from huggingface_hub import snapshot_download
            self.model_path = snapshot_download(repo_id=self.model_path)
            
        self.config = XoronConfig.from_pretrained(self.model_path)
        self.model = XoronMultimodalModel(self.config)
        from safetensors.torch import load_model as st_load
        
        # Handle sharded or single safetensors
        weights_path = f"{self.model_path}/model.safetensors"
        
        try: 
            if os.path.isdir(self.model_path):
                # Safetensors cannot load a directory directly using load_model
                from safetensors.torch import load_file
                import glob
                safetensor_files = glob.glob(f"{self.model_path}/*.safetensors")
                if safetensor_files:
                    state_dict = {}
                    for f in safetensor_files:
                        state_dict.update(load_file(f, device=str(self.device)))
                    self.model.load_state_dict(state_dict, strict=False)
                else:
                    raise FileNotFoundError("No .safetensors found in directory")
            else:
                st_load(self.model, weights_path, strict=False)
        except Exception as e: print(f"[XoronEngine] Note: {e}")
        
        self.processor = AutoProcessor.from_pretrained(self.model_path, trust_remote_code=True)

    def enable_expert_offloading(self, cold_threshold: int = 10, max_offload_ratio: float = 0.5):
        """
        Natively enables CPU/GPU expert offloading memory management for the
        underlying LLM Mixture of Experts layers to save VRAM on limited devices.
        """
        try:
            from xorfice.models.components.expert_offload import apply_expert_offloading
            print(f"[XoronEngine] Enabling native Expert Offloading (Cold Threshold: {cold_threshold}, Max Offload: {max_offload_ratio * 100}%)")
            manager = apply_expert_offloading(
                self.model,
                cold_threshold=cold_threshold,
                max_offload_ratio=max_offload_ratio
            )
            return manager is not None
        except ImportError:
            print("[XoronEngine] Could not load `expert_offload.py`, check package imports.")
            return False

    def _apply_optimizations(self):
        print("[XoronEngine] Applying SOTA multimodal optimizations...")
        
        # --- Hardware Auto-Optimization ---
        num_gpus = torch.cuda.device_count()
        
        # NUMA Awareness Check
        numa_nodes = 1
        try:
            if os.path.exists('/sys/devices/system/node'):
                numa_nodes = len([n for n in os.listdir('/sys/devices/system/node') if n.startswith('node')])
        except Exception: pass
        if numa_nodes > 1:
            print(f"[XoronEngine] 🧠 NUMA Architecture Detected ({numa_nodes} nodes). Optimizing memory affinity...")
            
        if num_gpus > 0:
            total_vram = sum(torch.cuda.get_device_properties(i).total_memory for i in range(num_gpus)) / (1024**3)
            print(f"[XoronEngine] 🖥️ Detected {num_gpus} GPU(s) with {total_vram:.1f}GB total VRAM.")
            
            # Dynamic tuning based on available VRAM & NUMA awareness
            vram_per_numa_factor = 1.0 if numa_nodes == 1 else 0.8 # Conservative allocation across NUMA boundaries
            self.max_vram_experts = max(2, int((total_vram / 10) * vram_per_numa_factor))
            self.kv_cache_blocks = max(128, int(total_vram * 8 * vram_per_numa_factor))
            print(f"[XoronEngine] ⚙️ Auto-tuned efficiency: max_experts={self.max_vram_experts}, kv_blocks={self.kv_cache_blocks}")
            
            if num_gpus > 1:
                print(f"[XoronEngine] 🔄 Multiple GPUs detected. Enabling Ring Attention and Model Parallelism...")
                device_map = {
                    'vision_encoder': 'cuda:0',
                    'video_encoder': 'cuda:0',
                    'audio_encoder': 'cuda:0',
                    'audio_decoder': 'cuda:0',
                    'waveform_decoder': 'cuda:0',
                    'projector': f'cuda:{min(1, num_gpus-1)}',
                    'audio_projector': f'cuda:{min(1, num_gpus-1)}',
                    'llm': f'cuda:{min(1, num_gpus-1)}',
                    'cross_attention': f'cuda:{min(num_gpus-1, 2)}',
                    'generator': f'cuda:{min(num_gpus-1, 2)}',
                    'video_generator': f'cuda:{min(num_gpus-1, 2)}',
                    'modality_markers': 'cuda:0',
                }
                # Apply Model Parallelism across NUMA nodes if possible
                self.model.apply_model_parallel(device_map)
                
        else:
            cpu_count = os.cpu_count() or 1
            # Auto-tune for CPU taking NUMA nodes into account
            optimal_threads = cpu_count // numa_nodes
            print(f"[XoronEngine] 🖥️ No GPUs detected. Auto-tuning for {cpu_count} CPU cores across {numa_nodes} NUMA node(s).")
            torch.set_num_threads(optimal_threads)
            self.device = torch.device("cpu")
        # ----------------------------------
        
        # 1. Automated Expert Management
        self.expert_cache = LRUExpertCache(capacity=self.max_vram_experts, device=str(self.device))
        self._wrap_moe_layers()
        
        # 2. Paged KV Cache
        self.kv_manager = PagedKVCacheManager(
            num_blocks=self.kv_cache_blocks,
            block_size=self.block_size,
            num_heads=self.config.num_heads,
            head_dim=self.config.hidden_size // self.config.num_heads,
            device=str(self.device)
        )
        
        # 3. Native optimizations (Triton patches applied directly to modules)
        
        # 4. Apply LoRA for Personalization
        if not self.model.lora_applied:
            self.model.apply_lora()
        
        self.model.to(self.device)
        print("[XoronEngine] 🚀 Elite SOTA Pipelines Ready: T2I, T2V, I2I, V2V, S2S.")

    def _wrap_moe_layers(self):
        from xorfice.models.llm.moe_llama import AuxLosslessMoELayer
        from xorfice.models.generators.video import TemporalMoELayer
        
        for name, module in self.model.named_modules():
            if "." in name:
                parent_name, child_name = name.rsplit(".", 1)
                parent = self.model.get_submodule(parent_name)
                child = getattr(parent, child_name)
                if isinstance(child, (AuxLosslessMoELayer, TemporalMoELayer)):
                    setattr(parent, child_name, ManagedMoELayer(child, name, self.expert_cache))

    @torch.no_grad()
    def generate(self, prompt: str, user_id: Optional[str] = None, images=None, videos=None, audios=None, stream=False, **kwargs) -> Union[Dict[str, Any], Generator[Any, None, None]]:
        """Unified SOTA multimodal entry point with per-user personalization."""
        
        # 1. Load User Personalization
        if user_id:
            self.personalization.load_user_weights(self.model, user_id)
            history = self.memory.get_context(user_id)
            if history:
                context_str = "\n".join([f"{h['role']}: {h['content']}" for h in history])
                prompt = f"<|system|>Previous Context:\n{context_str}<|/system|>\n{prompt}"

        # 2. Handle Media Inputs (load from path/URL if needed)
        def load_media(m, m_type):
            if isinstance(m, str):
                is_url = m.startswith("http")
                if not (is_url or os.path.exists(m)):
                    return m
                
                try:
                    # Download to temp if URL
                    if is_url:
                        response = requests.get(m, timeout=10)
                        response.raise_for_status()
                        
                        if m_type == "image":
                            return Image.open(io.BytesIO(response.content)).convert("RGB")
                        
                        # Audio/Video require Temp Files for standard libs
                        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{m_type}") as tmp:
                            tmp.write(response.content)
                            tmp_path = tmp.name
                    else:
                        tmp_path = m
                        if m_type == "image":
                            return Image.open(m).convert("RGB")

                    if m_type == "audio":
                        waveform, _ = torchaudio.load(tmp_path)
                        if is_url: os.remove(tmp_path)
                        return waveform[0]
                        
                    if m_type == "video":
                        # Real Video Extraction using OpenCV
                        cap = cv2.VideoCapture(tmp_path)
                        frames = []
                        target_frames = 16 
                        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
                        step = max(1, total_frames // target_frames)
                        
                        curr_frame = 0
                        while len(frames) < target_frames and cap.isOpened():
                            cap.set(cv2.CAP_PROP_POS_FRAMES, curr_frame)
                            ret, frame = cap.read()
                            if not ret: break
                            # Convert BGR to RGB and resize to 224x224 SOTA standard
                            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                            frame = cv2.resize(frame, (224, 224))
                            frames.append(frame)
                            curr_frame += step
                            
                        cap.release()
                        if is_url: os.remove(tmp_path)
                        
                        # Pad if video is too short
                        while len(frames) < target_frames and frames:
                            frames.append(frames[-1])
                            
                        if not frames:
                            return torch.zeros(3, target_frames, 224, 224)
                            
                        # Convert to (C, F, H, W) tensor
                        video_tensor = torch.from_numpy(np.array(frames)).permute(3, 0, 1, 2).float() / 255.0
                        return video_tensor
                        
                except Exception as e:
                    print(f"[XoronEngine] Failed to load {m_type} from {m}: {e}")
            return m

        images = load_media(images, "image")
        audios = load_media(audios, "audio")
        videos = load_media(videos, "video")

        inputs = self.processor(text=prompt, images=images, videos=videos, audios=audios, return_tensors="pt").to(self.device)
        
        # 3. Direct Multimodal Paths
        if audios is not None and kwargs.get("mode") == "s2s":
            return self.speech_to_speech(audios, prompt, user_id=user_id)

        if "<|generate_image|>" in prompt: return {"image_url": self.generate_image(inputs)}
        if "<|generate_video|>" in prompt: return {"video_url": self.generate_video(inputs)}
        
        if "<|edit_image|>" in prompt and images is not None:
            clean_instruction = prompt.replace("<|edit_image|>", "").strip()
            return {"image_url": self.edit_image(images, clean_instruction)}
            
        if "<|edit_video|>" in prompt and videos is not None:
            clean_instruction = prompt.replace("<|edit_video|>", "").strip()
            return {"video_url": self.edit_video(videos, clean_instruction)}
        
        if "<|speak|>" in prompt:
            text_to_speak = prompt.replace("<|speak|>", "").strip()
            return self.text_to_speech(text_to_speak, user_id=user_id)

        # Standard Text/Stream Path
        if stream: return self._generate_stream(inputs, **kwargs)
        
        max_tokens = kwargs.get("max_new_tokens", 1024)
        temperature = kwargs.get("temperature", 0.7)
        
        # Fallback to eager model.generate
        output_ids = self.model.generate(**inputs, max_new_tokens=max_tokens, temperature=temperature)
        return {"text": self.processor.batch_decode(output_ids, skip_special_tokens=True)[0]}

    def speech_to_speech(self, audio_waveform: torch.Tensor, instruction: str, user_id: Optional[str] = None) -> Dict[str, Any]:
        """SOTA Optimized S2S: Audio In -> Audio Out with Native Agentic MCP Tooling."""
        print("[XoronEngine] S2S: Direct Audio Agent pipeline active...")
        assert self.model is not None, "Model not loaded"
        
        # 1. Extract speaker embedding for zero-shot cloning
        speaker_embedding = self.voice.extract_speaker_embedding(self.model, audio_waveform)
        if user_id:
            self.voice.set_user_voice(user_id, speaker_embedding)
        
        # 2. Extract text response using ASR + LLM Agent logic
        # For direct S2S, we use listen_and_respond with the full tool loop
        tools = self.agent.get_all_tool_definitions()
        
        output = self.model.listen_and_respond(
            audio_waveform=audio_waveform.to(self.device),
            tokenizer=self.processor.tokenizer if hasattr(self.processor, 'tokenizer') else self.processor,
            speaker_embedding=speaker_embedding.to(self.device) if speaker_embedding is not None else None,
            tool_executor=self.agent.execute_tool,
            available_tools=tools,
            system_prompt="You are Xoron-Dev, an elite AI. You can write code, run terminal commands, and use tools to fulfill the user's voice requests.",
            max_tool_calls=5
        )
        
        path = f"outputs/s2s_{user_id or 'anon'}.wav"
        
        # Real SOTA flow: Save waveform physically and return proper decoded text
        if "waveform" in output and output["waveform"] is not None:
            import torchaudio
            waveform_cpu = output["waveform"].cpu()
            # Ensure 2D tensor for torchaudio [channels, time]
            if waveform_cpu.dim() == 1:
                waveform_cpu = waveform_cpu.unsqueeze(0)
            elif waveform_cpu.dim() == 3:
                waveform_cpu = waveform_cpu.squeeze(0)
                
            torchaudio.save(path, waveform_cpu, sample_rate=16000)
            
        transcription = output.get("text", "Unknown")
        speaking_text = output.get("speaking_text", transcription)
        
        return {
            "audio_url": f"/{path}", 
            "text": transcription,
            "speaking_text": speaking_text,
            "tool_calls": output.get("tool_calls", [])
        }

    def speech_to_speech_stream(self, audio_waveform: torch.Tensor, instruction: str, user_id: Optional[str] = None):
        """SOTA True Real-Time S2S Stream: Audio In -> Streaming Text & Audio Chunks Out."""
        print("[XoronEngine] S2S: Real-Time WebSocket Streaming pipeline active...")
        assert self.model is not None, "Model not loaded"
        
        # 1. Zero-shot cloning lookup or extraction
        speaker_embedding = None
        if user_id:
            try: speaker_embedding = self.voice.load_user_voice(user_id)
            except Exception: pass
            if speaker_embedding is None:
                speaker_embedding = self.voice.extract_speaker_embedding(self.model, audio_waveform)
                self.voice.set_user_voice(user_id, speaker_embedding)
        
        # 2. Generator looping
        tools = self.agent.get_all_tool_definitions()
        generator = self.model.stream_listen_and_respond(
            audio_waveform=audio_waveform.to(self.device),
            tokenizer=self.processor.tokenizer if hasattr(self.processor, 'tokenizer') else self.processor,
            speaker_embedding=speaker_embedding.to(self.device) if speaker_embedding is not None else None,
            tool_executor=self.agent.execute_tool,
            available_tools=tools,
            system_prompt="You are Xoron-Dev, an elite AI. You can write code, run terminal commands, and use tools to fulfill the user's voice requests.",
            max_tool_calls=5
        )
        
        for item in generator:
            yield item

    def text_to_speech(self, text: str, user_id: Optional[str] = None, voice_id: Optional[str] = None) -> Dict[str, Any]:
        """SOTA Zero-Shot TTS: Text -> Audio."""
        print(f"[XoronEngine] TTS: Generating audio for text with voice {voice_id or user_id or 'default'}")
        assert self.model is not None, "Model not loaded"
        
        # 1. Get speaker embedding
        speaker_embedding = None
        if voice_id:
            speaker_embedding = self.voice.load_user_voice(voice_id)
        elif user_id:
            speaker_embedding = self.voice.load_user_voice(user_id)
            
        if speaker_embedding is not None:
            speaker_embedding = speaker_embedding.to(self.device)
            
        # 2. Generate audio waveform from text
        text_inputs = self.processor(text=text, return_tensors="pt").to(self.device)
        waveform = self.model.speak(
            text_inputs.input_ids,
            text_inputs.get("attention_mask"),
            speaker_embedding=speaker_embedding
        )
        
        path = f"outputs/tts_{user_id or 'anon'}.wav"
        # Save waveform to path
        return {"audio_url": f"/{path}"}

    def generate_image(self, inputs: Dict[str, Any]) -> str:
        """SOTA T2I with Fused Kernels."""
        print("[XoronEngine] T2I: Generating image...")
        assert self.model is not None, "Model not loaded"
        
        # Real SOTA flow: extract embeddings, run diffusion, save image
        path = "outputs/generated_image.png"
        try:
            # Assumes model.generate_image returns a (1, 3, H, W) tensor in [-1, 1] or [0, 1]
            output_tensor = self.model.generate_image(inputs["input_ids"], inputs.get("attention_mask"))
            if output_tensor is not None:
                # Post-process to [0, 255] RGB Image
                if output_tensor.min() < 0: output_tensor = (output_tensor + 1.0) / 2.0
                img_np = (output_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
                Image.fromarray(img_np).save(path)
        except Exception as e:
            raise RuntimeError(f"T2I implementation failed: {str(e)}")
        return f"/{path}"

    def edit_image(self, image: torch.Tensor, instruction: str) -> str:
        """SOTA I2I: Image editing path."""
        print(f"[XoronEngine] I2I: Editing image with instruction: {instruction}")
        assert self.model is not None, "Model not loaded"
        path = "outputs/edited_image.png"
        try:
            inputs = self.processor(text=instruction, images=image, return_tensors="pt").to(self.device)
            context = self.model.get_text_embeddings(inputs["input_ids"], inputs.get("attention_mask"))
            mask = torch.ones_like(inputs["pixel_values"])
            output_tensor = self.model.generator.edit_image(inputs["pixel_values"], context, mask)
            if output_tensor is not None:
                if output_tensor.min() < 0: output_tensor = (output_tensor + 1.0) / 2.0
                img_np = (output_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
                Image.fromarray(img_np).save(path)
        except Exception as e:
            raise RuntimeError(f"I2I implementation failed: {str(e)}")
        
        return f"/{path}"

    def generate_video(self, inputs: Dict[str, Any]) -> str:
        """SOTA T2V with 3D-RoPE and Ring Attention."""
        print("[XoronEngine] T2V: Generating video...")
        assert self.model is not None, "Model not loaded"
        path = "outputs/generated_video.mp4"
        try:
            output_tensor = self.model.generate_video(inputs["input_ids"], inputs.get("attention_mask"))
            if output_tensor is not None:
                import torchvision.io
                frames = output_tensor.squeeze(0).permute(1, 2, 3, 0) # F, H, W, C
                if frames.min() < 0: frames = (frames + 1.0) / 2.0
                frames = (frames * 255).to(torch.uint8).cpu()
                torchvision.io.write_video(path, frames, fps=8)
        except Exception as e:
            raise RuntimeError(f"T2V implementation failed: {str(e)}")
            
        return f"/{path}"

    def edit_video(self, video: torch.Tensor, instruction: str) -> str:
        """SOTA V2V: Video editing using I2V paths."""
        print(f"[XoronEngine] V2V: Editing video with temporal guidance...")
        assert self.model is not None, "Model not loaded"
        path = "outputs/edited_video.mp4"
        try:
            inputs = self.processor(text=instruction, videos=video, return_tensors="pt").to(self.device)
            first_frame = inputs["video_frames"][:, :, 0] 
            output_tensor = self.model.generate_video(inputs["input_ids"], inputs.get("attention_mask"), first_frame=first_frame)
            if output_tensor is not None:
                import torchvision.io
                frames = output_tensor.squeeze(0).permute(1, 2, 3, 0)
                if frames.min() < 0: frames = (frames + 1.0) / 2.0
                frames = (frames * 255).to(torch.uint8).cpu()
                torchvision.io.write_video(path, frames, fps=8)
        except Exception as e:
            raise RuntimeError(f"V2V implementation failed: {str(e)}")
        return f"/{path}"

    @torch.no_grad()
    def _generate_stream(self, inputs: Dict[str, Any], **kwargs) -> Generator[Any, None, None]:
        """Real SOTA Streaming: Yields tokens one-by-one using the actual model with Agentic behavior."""
        max_new_tokens = kwargs.get("max_new_tokens", 512)
        temperature = kwargs.get("temperature", 0.7)
        agent_loops = 0
        max_agent_loops = 5
        
        generated_ids = inputs["input_ids"]
        attention_mask = inputs.get("attention_mask")
        
        # Enforce explicitly mapping images -> pixel_values and videos -> video_frames into generation kwargs
        # so XoronMultimodalModel.forward doesn't drop the multimodal embeddings midway.
        pixel_values = inputs.get("pixel_values", kwargs.get("images"))
        video_frames = inputs.get("video_frames", kwargs.get("videos"))
        audio_features = inputs.get("audio_features", kwargs.get("audios"))
        
        prompt_tokens = int(generated_ids.shape[1])
        completion_tokens = 0
        current_response_text = ""
        
        while agent_loops < max_agent_loops:
            for _ in range(max_new_tokens):
                # Pass all active modalities natively down the multi-modal fusion bridge
                outputs = self.model.forward(
                    input_ids=generated_ids[:, -1:], 
                    attention_mask=attention_mask,
                    pixel_values=pixel_values,
                    video_frames=video_frames,
                    audio_features=audio_features
                )
                next_token_logits = outputs.logits[:, -1, :]
                
                if temperature != 1.0 and temperature > 0:
                    next_token_logits = next_token_logits / temperature
                    probs = torch.softmax(next_token_logits, dim=-1)
                    new_token_id = torch.multinomial(probs, num_samples=1)
                else:
                    new_token_id = torch.argmax(next_token_logits, dim=-1, keepdim=True)
                
                new_token_text = self.processor.batch_decode(new_token_id, skip_special_tokens=True)[0]
                
                if not new_token_text: break
                    
                completion_tokens += 1
                yield new_token_text
                current_response_text += new_token_text
                
                generated_ids = torch.cat([generated_ids, new_token_id], dim=-1)
                if attention_mask is not None:
                    attention_mask = torch.cat([attention_mask, torch.ones((attention_mask.shape[0], 1), device=self.device)], dim=-1)
                
                if "</s>" in new_token_text or "<|endoftext|>" in new_token_text:
                    break

            # Agentic Check: Did the model call a tool?
            tool_calls = self.agent.parse_tool_calls(current_response_text)
            if not tool_calls:
                break
                
            agent_loops += 1
            for call in tool_calls:
                print(f"[XoronEngine] 🛠️ Model calling tool: {call['name']}({call.get('args', {})})")
                observation = self.agent.execute_tool(call['name'], **call.get('args', {}))
                response_str = self.agent.format_tool_response(observation)
                yield response_str
                
                # Feed observation back into model
                obs_ids = self.processor.tokenizer(response_str, return_tensors="pt").input_ids.to(self.device)
                prompt_tokens += int(obs_ids.shape[1])
                generated_ids = torch.cat([generated_ids, obs_ids], dim=-1)
                if attention_mask is not None:
                    attention_mask = torch.cat([attention_mask, torch.ones((1, obs_ids.shape[1]), device=self.device)], dim=-1)
                
                current_response_text = "" # Reset for next turn
                
        # Final Payload: Usage Statistics
        yield {
            "usage": {
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens
            }
        }

    def train_user_on_interaction(self, user_id: str, prompt: str, response: str):
        """Triggers an on-the-fly training step for a specific user interaction."""
        print(f"[XoronEngine] Triggering SOTA On-the-fly Personalization for user {user_id}...")
        
        # Tokenize content
        prompt_ids = self.processor.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
        response_ids = self.processor.tokenizer(response, return_tensors="pt").input_ids.to(self.device)
        
        # Perform training step
        self.personalization.train_on_fly(self.model, user_id, prompt_ids, response_ids)
        
        # Save interaction to persistent memory
        self.memory.add_interaction(user_id, prompt, response)

    def get_model_info(self) -> Dict[str, Any]:
        return {"id": self.config.model_name if self.config else "Xoron-Dev", "object": "model", "created": 1234567890, "owned_by": "xoron-dev"}
