#!/usr/bin/env python3
"""
XYFLOW Dual-Copy Chatbot with proper instruct template, sampling, and reliable temperature.
"""

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

# =============================================================================
# CONFIGURATION
# =============================================================================
MODEL_NAME = "HuggingFaceTB/SmolLM2-135M-Instruct"
MAX_NEW_TOKENS = 1024
SWITCH_INTERVAL = 6
THERMAL_THRESHOLD = 80.0   # °C – early switch if exceeded
CORE_A = 0
CORE_B = 2

# Generation settings to avoid repetition
TEMPERATURE = 0.7
TOP_P = 0.9

# =============================================================================
# RELIABLE TEMPERATURE READING
# =============================================================================

def get_cpu_temperature():
    """Try multiple methods to get CPU package temperature on Linux."""
    # Method 1: psutil (often works on Intel)
    try:
        temps = psutil.sensors_temperatures()
        if 'coretemp' in temps:
            return temps['coretemp'][0].current
        if 'cpu_thermal' in temps:
            return temps['cpu_thermal'][0].current
        # fallback to first available
        for sensor in temps.values():
            if sensor:
                return sensor[0].current
    except:
        pass

    # Method 2: read from sysfs (most reliable)
    try:
        base = "/sys/class/thermal/thermal_zone"
        for i in range(10):
            path = f"{base}{i}/temp"
            if os.path.exists(path):
                with open(path, 'r') as f:
                    val = int(f.read().strip()) / 1000.0
                    if val > 0:
                        return val
    except:
        pass

    # Method 3: run `sensors` command
    try:
        output = subprocess.check_output(['sensors'], text=True)
        for line in output.split('\n'):
            if 'Package' in line or 'CPU' in line:
                # e.g., "Package id 0:  +45.0°C"
                import re
                match = re.search(r'\+([0-9.]+)°C', line)
                if match:
                    return float(match.group(1))
    except:
        pass

    # Fallback: return 0.0 if all fail
    return 0.0

# =============================================================================
# THERMAL MONITOR (logs temperatures during generation)
# =============================================================================

class ThermalMonitor:
    def __init__(self, interval=1.0):
        self.interval = interval
        self.running = False
        self.temps = []
        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):
        return max((t for _, t in self.temps), default=0.0)

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

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

class DualCopyChatSession:
    def __init__(self):
        # Force single-threaded for core 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
        ).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
        ).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.messages = []          # list of dicts: [{"role": "user", "content": ...}, ...]

    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, input_ids, past_kv, attention_mask, max_new_tokens):
        """
        Generate up to max_new_tokens using the current active model.
        Returns: (list of new token ids, updated past_kv, updated attention_mask, eos_reached)
        """
        generated = []
        current_ids = input_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, :]
            # Apply temperature and top-p sampling
            if TEMPERATURE > 0:
                logits = logits / TEMPERATURE
                probs = torch.softmax(logits, dim=-1)
                # Top-p (nucleus) sampling
                sorted_probs, sorted_indices = torch.sort(probs, descending=True)
                cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
                sorted_indices_to_remove = cumulative_probs > TOP_P
                # Shift to keep the first token (which is always <= TOP_P)
                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                sorted_indices_to_remove[..., 0] = False
                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                probs[indices_to_remove] = 0.0
                next_token = torch.multinomial(probs, 1).squeeze(1)
            else:
                next_token = torch.argmax(logits, dim=-1)

            next_id = next_token.item()
            generated.append(next_id)
            current_past = outputs.past_key_values
            current_ids = next_token.unsqueeze(0)  # shape (1,1)

            # Extend attention mask
            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)

            # Stop if EOS token is generated
            if next_id == self.tokenizer.eos_token_id:
                break

        return generated, current_past, current_mask

    def respond(self, user_message, max_tokens=MAX_NEW_TOKENS):
        """
        Process user message, generate assistant reply, and update conversation.
        """
        # Append user message to conversation history
        self.messages.append({"role": "user", "content": user_message})

        # Apply chat template to get the input for the model
        # For instruct models, this formats with <|user|>, <|assistant|>, etc.
        prompt = self.tokenizer.apply_chat_template(
            self.messages,
            tokenize=False,
            add_generation_prompt=True   # adds the assistant start token
        )

        # Tokenize the full prompt (we'll use it for the first chunk)
        # But we need to handle caching across chunks.
        # For simplicity, we'll tokenize the entire prompt and use it as input_ids.
        # For subsequent chunks we'll feed only the last token.
        input_ids = self.tokenizer.encode(prompt, return_tensors="pt")
        attention_mask = torch.ones_like(input_ids)

        # We'll generate using the active model, but we need to manage past_kv.
        # We'll reset past_kv at the start of each turn because the prompt changed.
        # (We could try to reuse past from previous turns, but it's easier to just reset.)
        # However, we want the model to have the full conversation context.
        # Using the full prompt each turn is correct; we just start fresh cache.
        # The dual-copy switching will still distribute heat during generation of this reply.
        current_past = None
        current_mask = attention_mask

        print(f"\n[Assistant generating]")
        monitor = ThermalMonitor()
        monitor.start()

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

        # We'll start with the full input_ids for the first chunk
        start_ids = input_ids

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

                # Thermal emergency switch
                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 input for this chunk
                if current_past is None:
                    # First chunk: use the whole prompt (input_ids)
                    chunk_input = start_ids
                else:
                    # Subsequent chunks: use only the last generated token
                    if last_token_id is None:
                        # Should not happen
                        break
                    chunk_input = torch.tensor([[last_token_id]], dtype=torch.long)

                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_input, current_past, current_mask, chunk_size
                )

                if not chunk_tokens:
                    break  # no tokens generated

                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 ({len(chunk_tokens)/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

                # If we hit EOS, stop early
                if last_token_id == self.tokenizer.eos_token_id:
                    break

        finally:
            monitor.stop()

        # Decode assistant reply
        reply_text = self.tokenizer.decode(all_assistant_tokens, skip_special_tokens=True)

        # Append assistant reply to conversation history
        self.messages.append({"role": "assistant", "content": reply_text})

        return reply_text, monitor

# =============================================================================
# UTILITY: CPU AFFINITY
# =============================================================================

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

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

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

    session = DualCopyChatSession()

    print("\nChat started. 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

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

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

if __name__ == "__main__":
    main()