#!/usr/bin/env python3
"""
XYFLOW Dual-Copy Chatbot for HuggingFaceTB/SmolLM2-135M-Instruct
Interactive terminal chat with thermal-aware model switching.
"""

import os
import time
import threading
import psutil
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# =============================================================================
# CONFIGURATION
# =============================================================================
MODEL_NAME = "HuggingFaceTB/SmolLM2-135M-Instruct"
MAX_NEW_TOKENS = 60                # max tokens per assistant reply
SWITCH_INTERVAL = 6                # switch models every N generated tokens
THERMAL_THRESHOLD = 80.0           # °C – early switch if exceeded

# Physical core IDs (adjust to your CPU topology)
CORE_A = 0
CORE_B = 2

# System prompt (optional, adapt to the instruct model)
SYSTEM_PROMPT = "You are a helpful AI assistant."

# =============================================================================
# SYSTEM UTILITIES
# =============================================================================

def set_affinity(pid, core_id):
    try:
        os.sched_setaffinity(pid, {core_id})
        return True
    except (AttributeError, OSError):
        try:
            psutil.Process(pid).cpu_affinity([core_id])
            return True
        except AttributeError:
            print("Warning: CPU affinity not supported.")
            return False

def get_cpu_temperature():
    try:
        temps = psutil.sensors_temperatures()
        for sensor in temps.values():
            if sensor:
                return sensor[0].current
    except:
        pass
    return 0.0

# =============================================================================
# THERMAL MONITOR (background logger)
# =============================================================================

class ThermalMonitor:
    def __init__(self, interval=1.0):
        self.interval = interval
        self.running = False
        self.temps = []          # list of (timestamp, temp)
        self.thread = None

    def start(self):
        self.running = True
        self.temps = []
        self.thread = threading.Thread(target=self._run)
        self.thread.daemon = True
        self.thread.start()

    def stop(self):
        self.running = False
        if self.thread:
            self.thread.join()

    def _run(self):
        start = time.time()
        while self.running:
            temp = get_cpu_temperature()
            self.temps.append((time.time() - start, temp))
            time.sleep(self.interval)

    def get_peak(self):
        if not self.temps:
            return 0.0
        return max(t for _, t in self.temps)

    def get_latest(self):
        return self.temps[-1][1] if self.temps else 0.0

# =============================================================================
# CHAT SESSION WITH DUAL-COPY
# =============================================================================

class DualCopyChatSession:
    def __init__(self, system_prompt=SYSTEM_PROMPT):
        # Force single-threaded for affinity
        os.environ["OMP_NUM_THREADS"] = "1"
        os.environ["MKL_NUM_THREADS"] = "1"
        torch.set_num_threads(1)

        print("Loading Tokenizer...")
        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        print(f"Loading Model Copy A (pinned to Core {CORE_A})...")
        self.model_a = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME, torch_dtype=torch.float32, low_cpu_mem_usage=True
        )
        self.model_a.to('cpu')

        print(f"Loading Model Copy B (pinned to Core {CORE_B})...")
        self.model_b = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME, torch_dtype=torch.float32, low_cpu_mem_usage=True
        )
        self.model_b.to('cpu')

        self.active_model = self.model_a
        self.active_core = CORE_A
        set_affinity(0, self.active_core)

        # Conversation state
        self.past_kv = None
        self.attention_mask = None
        self.full_tokens = []   # list of token IDs for the entire conversation
        self.system_prompt = system_prompt

        # Tokenize system prompt once (will be used as initial context)
        self._init_system()

    def _init_system(self):
        """Initialize conversation with system prompt."""
        if self.system_prompt:
            sys_ids = self.tokenizer.encode(self.system_prompt, add_special_tokens=True)
            self.full_tokens = sys_ids
            # We don't generate yet; the first user message will be appended and generate.

    def _switch_model(self):
        if self.active_model is self.model_a:
            self.active_model = self.model_b
            self.active_core = CORE_B
        else:
            self.active_model = self.model_a
            self.active_core = CORE_A
        set_affinity(0, self.active_core)
        print(f"\n[SWITCH] -> {'B' if self.active_core == CORE_B else 'A'} on Core {self.active_core}")

    def _generate_chunk(self, start_ids, past_kv, attention_mask, max_new_tokens):
        """
        Generate up to max_new_tokens starting from start_ids (single token or sequence).
        Returns: (list of new token ids, updated past_kv, updated attention_mask)
        """
        generated = []
        current_ids = start_ids
        current_past = past_kv
        current_mask = attention_mask

        for _ in range(max_new_tokens):
            with torch.no_grad():
                outputs = self.active_model(
                    input_ids=current_ids,
                    past_key_values=current_past,
                    attention_mask=current_mask,
                    use_cache=True
                )
            logits = outputs.logits[:, -1, :]
            next_token = torch.argmax(logits, dim=-1, keepdim=True)
            next_id = next_token.item()
            generated.append(next_id)

            current_past = outputs.past_key_values
            current_ids = next_token
            if current_mask is not None:
                current_mask = torch.cat([current_mask, torch.ones((1, 1), dtype=torch.long)], dim=1)
            else:
                current_mask = torch.ones((1, 1), dtype=torch.long)

        return generated, current_past, current_mask

    def respond(self, user_message, max_tokens=MAX_NEW_TOKENS, switch_interval=SWITCH_INTERVAL):
        """
        Process a user message, generate an assistant reply using dual-copy switching.
        Returns the assistant reply text and the thermal monitor.
        """
        # Append user message to conversation
        user_ids = self.tokenizer.encode(user_message, add_special_tokens=False)
        self.full_tokens.extend(user_ids)

        # Initial input_ids for generation: the new user tokens (if no past_kv yet)
        # If past_kv is None, we need to provide the whole conversation so far.
        if self.past_kv is None:
            # First turn: use the entire conversation (system + user) as input
            start_ids = torch.tensor([self.full_tokens], dtype=torch.long)
            attention_mask = None
        else:
            # Subsequent turns: only the new tokens (the user message) are fed,
            # but we need to extend the attention_mask with ones for the new tokens.
            start_ids = torch.tensor([user_ids], dtype=torch.long)
            # Extend mask: existing mask has length L, we add len(user_ids) ones
            if self.attention_mask is not None:
                new_mask = torch.ones((1, len(user_ids)), dtype=torch.long)
                attention_mask = torch.cat([self.attention_mask, new_mask], dim=1)
            else:
                attention_mask = torch.ones((1, len(user_ids)), dtype=torch.long)

        # We'll need to generate assistant tokens, but the model already has context.
        # However, we must pass start_ids and the past_kv from the previous conversation.
        # If past_kv is None, the model will process the whole start_ids (first turn).
        # For subsequent turns, past_kv contains the previous conversation, and start_ids are the new tokens.

        # Start generation
        print(f"\n[Assistant generating] (switch every {switch_interval} tokens)")
        monitor = ThermalMonitor()
        monitor.start()

        all_assistant_tokens = []
        remaining = max_tokens
        token_count_since_switch = 0
        last_token_id = None

        # Use a local reference to active_model and core to update during generation
        # We'll update the class attributes as we switch.
        # We'll use a local variable for past_kv and attention_mask that we'll update
        current_past = self.past_kv
        current_mask = attention_mask if self.past_kv is not None else None

        try:
            while remaining > 0:
                chunk_size = min(switch_interval, remaining)

                # Thermal emergency switch (check before each chunk)
                temp = get_cpu_temperature()
                if temp >= THERMAL_THRESHOLD:
                    print(f"\n[THERMAL TRIGGER] {temp:.1f}°C >= {THERMAL_THRESHOLD}°C. Forcing early switch.")
                    self._switch_model()
                    token_count_since_switch = 0

                # Determine starting token(s) for this chunk
                if current_past is None:
                    # First chunk of the whole generation: start from the full context (start_ids)
                    chunk_start = start_ids
                else:
                    # Subsequent chunks: only the last generated token
                    chunk_start = torch.tensor([[last_token_id]], dtype=torch.long) if last_token_id is not None else start_ids

                print(f"  -> Generating {chunk_size} tokens on Core {self.active_core}...", end="", flush=True)
                start_time = time.time()

                chunk_tokens, current_past, current_mask = self._generate_chunk(
                    chunk_start, current_past, current_mask, chunk_size
                )

                all_assistant_tokens.extend(chunk_tokens)
                last_token_id = chunk_tokens[-1]
                remaining -= chunk_size
                token_count_since_switch += chunk_size

                elapsed = time.time() - start_time
                print(f" done in {elapsed:.2f}s ({chunk_size/elapsed:.2f} tok/s)")

                # Switch if interval reached and more tokens remain
                if remaining > 0 and token_count_since_switch >= switch_interval:
                    self._switch_model()
                    token_count_since_switch = 0

        finally:
            monitor.stop()

        # Update session state with the new assistant tokens
        self.full_tokens.extend(all_assistant_tokens)
        self.past_kv = current_past
        self.attention_mask = current_mask

        # Decode assistant reply (the newly generated tokens)
        reply_text = self.tokenizer.decode(all_assistant_tokens, skip_special_tokens=True)
        return reply_text, monitor

# =============================================================================
# MAIN CHAT LOOP
# =============================================================================

def main():
    print("="*60)
    print("XYFLOW Dual-Copy Chatbot")
    print(f"Model: {MODEL_NAME}")
    print(f"Switch interval: {SWITCH_INTERVAL} tokens")
    print(f"Thermal threshold: {THERMAL_THRESHOLD}°C")
    print("Type your messages below. Type 'exit' or 'quit' to stop.")
    print("="*60)

    session = DualCopyChatSession()

    print("\nChat session started. You can now talk to the assistant.\n")
    while True:
        user_input = input("You: ").strip()
        if user_input.lower() in ("exit", "quit", "bye"):
            print("Goodbye!")
            break
        if not user_input:
            continue

        # Get assistant response
        reply, monitor = session.respond(user_input)
        print(f"Assistant: {reply}")

        # Show thermal stats
        peak = monitor.get_peak()
        latest = monitor.get_latest()
        print(f"[Thermal] Peak: {peak:.1f}°C, Current: {latest:.1f}°C")

if __name__ == "__main__":
    main()