#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "requests",
# ]
# ///

"""
Ara NPU LLM Chatbot terminal
============================
"""

import json
import os
import subprocess
import sys
import time
import requests

API_URL = "http://127.0.0.1:8000/v1/chat/completions"
MODELS_URL = "http://127.0.0.1:8000/v1/models"
MODEL_NAME = "Qwen2.5-7B-Instruct"

CONFIG_PATH = "/usr/share/eiq/aaf-connector/server_config.json"
SERVICE_NAME = "eiq-aaf-connector.service"

def is_service_running(service_name: str) -> bool:
    """Checks if a systemd service is currently active and running."""
    try:
        result = subprocess.run(
            ["systemctl", "is-active", service_name],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        return result.stdout.strip() == "active"
    except Exception:
        return False

def ensure_model_enabled(target_model: str, config_path: str) -> bool:
    """Updates server_config.json so only target_model has enabled=True."""
    if not os.path.exists(config_path):
        print(f"[Warning] Config file not found at {config_path}.")
        return False

    with open(config_path, "r") as f:
        config = json.load(f)

    config_changed = False
    available_models = config.get("available_models", [])

    if isinstance(available_models, list):
        for model_info in available_models:
            model_id = model_info.get("name") or model_info.get("id")
            should_be_enabled = (model_id == target_model)
            if model_info.get("enabled") != should_be_enabled:
                model_info["enabled"] = should_be_enabled
                config_changed = True

    if config_changed:
        print(f"[Config] Enabling '{target_model}' in server_config.json...")
        with open(config_path, "w") as f:
            json.dump(config, f, indent=4)
        return True

    print(f"[Config] '{target_model}' is already set as enabled.")
    return False


def wait_for_model_ready(target_model: str, timeout: int = 300) -> bool:
    print(f"[Service] Waiting for '{target_model}' to finish loading on Ara NPU...", end="", flush=True)
    start_time = time.time()

    while time.time() - start_time < timeout:
        try:
            res = requests.get(MODELS_URL, timeout=2)
            if res.status_code == 200:
                data = res.json()
                loaded_models = [m.get("id") or m.get("name") for m in data.get("data", [])]
                
                # Check if the endpoint responds AND the target model is registered
                if target_model in loaded_models or len(loaded_models) > 0:
                    print("\n[Service] AAF Connector is ready and endpoint is active!")
                    return True
        except requests.exceptions.RequestException:
            pass

        print(".", end="", flush=True)
        time.sleep(2)

    print("\n[Error] Timed out waiting for AAF Connector to register model endpoint.")
    return False


def restart_aaf_service(service_name: str):
    """Restarts the AAF connector systemd service."""
    print(f"[Service] Restarting {service_name}... This could take several minutes")
    try:
        subprocess.run(["systemctl", "restart", service_name], check=True)
    except subprocess.CalledProcessError as e:
        print(f"[Error] Failed to restart service: {e}")
        sys.exit(1)


def chat():
    print(f"\n--- Gateworks AI LLM Session (Model: {MODEL_NAME}) ---")
    print("Type 'exit' to stop.\n")

    history = [{"role": "system", "content": "You are a helpful AI assistant."}]

    while True:
        print("You: ", end="", flush=True)
        user_input = input()
        if user_input.lower() in ['exit', 'quit']:
            break

        history.append({"role": "user", "content": user_input})
        payload = {
            "model": MODEL_NAME,
            "messages": history,
            "temperature": 0.7,
            "stream": True
        }

        print("AI: ", end="", flush=True)

        start_time = time.time()
        full_reply = ""
        token_count = 0

        try:
            response = requests.post(API_URL, json=payload, stream=True)
            response.raise_for_status()

            for line in response.iter_lines():
                if line:
                    decoded_line = line.decode('utf-8')
                    if decoded_line.startswith("data: "):
                        content = decoded_line[6:]
                        if content.strip() == "[DONE]":
                            break

                        chunk = json.loads(content)
                        if "choices" in chunk and chunk["choices"][0]["delta"].get("content"):
                            text = chunk["choices"][0]["delta"]["content"]
                            print(text, end="", flush=True)
                            full_reply += text
                            token_count += 1

            end_time = time.time()
            duration = end_time - start_time
            tps = token_count / duration if duration > 0 else 0

            print(f"\n\n--- Stats ---")
            print(f"Time taken: {duration:.2f} seconds")
            print(f"Throughput: {tps:.2f} tokens/sec")
            print(f"-------------\n")

            history.append({"role": "assistant", "content": full_reply})

        except Exception as e:
            print(f"\nError: {e}")

if __name__ == "__main__":
    # 1. Update config file to reflect desired model state
    config_updated = ensure_model_enabled(MODEL_NAME, CONFIG_PATH)

    # 2. Check if service is currently running
    service_active = is_service_running(SERVICE_NAME)

    # 3. Restart if config was modified OR if the service is stopped/failed
    if config_updated or not service_active:
        if not service_active:
            print(f"[Service] '{SERVICE_NAME}' is not active. Starting/Restarting now...")
        restart_aaf_service(SERVICE_NAME)

    # 4. Wait until the Ara NPU finishes compiling/loading the model and registers the endpoint
    if not wait_for_model_ready(MODEL_NAME, timeout=300):
        print("[Error] Could not connect to model service. Try checking logs with:")
        print(f"        journalctl -u {SERVICE_NAME} -n 50 --no-pager")
        sys.exit(1)

    # 5. Begin Chat
    chat()