import os
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
from concurrent.futures import ThreadPoolExecutor
import logging
import time
import threading

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

MODEL_ID = "YoussefElsafi/PlayerAI-1.2B-v1.5"

# ── CPU threading (set BEFORE any torch ops) ──────────────────────────────────
PHYSICAL_CORES   = os.cpu_count() or 2
NUM_MODEL_COPIES = 2                          # run 2 model instances in parallel
MAX_CONCURRENT   = NUM_MODEL_COPIES * 2           # two requests per model at a time
THREADS_PER_GEN  = max(1, PHYSICAL_CORES // NUM_MODEL_COPIES)

torch.set_num_threads(THREADS_PER_GEN)
torch.set_num_interop_threads(2)

for var in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS",
            "VECLIB_MAXIMUM_THREADS", "NUMEXPR_NUM_THREADS"):
    os.environ[var] = str(THREADS_PER_GEN)

logger.info(f"CPU cores: {PHYSICAL_CORES} | "
            f"Model copies: {NUM_MODEL_COPIES} | "
            f"Concurrent gens: {MAX_CONCURRENT} | "
            f"Threads/gen: {THREADS_PER_GEN}")

# ── Model loading ─────────────────────────────────────────────────────────────
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

# Detect device
if torch.cuda.is_available():
    DEVICE = "cuda"
    DTYPE  = torch.float16
    logger.info("Device: CUDA — float16")
else:
    DEVICE = "cpu"
    # INT8 quantization REQUIRES float32 weights as input.
    DTYPE  = torch.float32
    logger.info("Device: CPU — float32 (INT8 quantization will halve this)")

def load_one_model(idx: int):
    logger.info(f"Loading model copy #{idx}...")
    m = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        dtype             = DTYPE,
        device_map        = "auto" if DEVICE == "cuda" else "cpu",
        trust_remote_code = True,
        low_cpu_mem_usage = True,
    )
    m.eval()

    # INT8 dynamic quantization (CPU only, before torch.compile)
    if DEVICE == "cpu":
        try:
            m = torch.quantization.quantize_dynamic(
                m, {torch.nn.Linear}, dtype=torch.qint8,
            )
            logger.info(f"Model #{idx}: INT8 quantization applied ✓")
        except Exception as e:
            logger.warning(f"Model #{idx}: quantization skipped: {e}")

    if hasattr(torch, "compile"):
        try:
            m = torch.compile(m, mode="default")
            logger.info(f"Model #{idx}: torch.compile applied ✓")
        except Exception as e:
            logger.warning(f"Model #{idx}: torch.compile skipped: {e}")

    return m

print(f"Loading {NUM_MODEL_COPIES} model copies...")
models = [load_one_model(i) for i in range(NUM_MODEL_COPIES)]
model_locks = [threading.Lock() for _ in range(NUM_MODEL_COPIES)]
print(f"All {NUM_MODEL_COPIES} model copies loaded!")

# ── Generation config ─────────────────────────────────────────────────────────
GEN_CONFIG = dict(
    max_new_tokens     = 80,
    do_sample          = True,
    temperature        = 0.85,
    top_p              = 0.9,
    repetition_penalty = 1.1,
    pad_token_id       = tokenizer.eos_token_id,
    use_cache          = True,
)

# ── Thread pool + concurrency guards ─────────────────────────────────────────
executor      = ThreadPoolExecutor(max_workers=MAX_CONCURRENT, thread_name_prefix="gen")
queue_counter = threading.Semaphore(8)   # max 8 requests waiting

def acquire_free_model(timeout: float = 60.0):
    """Try to grab the first free model lock. Returns (idx, model) or (None, None)."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        for i, lock in enumerate(model_locks):
            if lock.acquire(blocking=False):
                return i, models[i]
        time.sleep(0.05)
    return None, None

# ── Auth ──────────────────────────────────────────────────────────────────────
ALLOWED_TOKEN = os.environ.get("PLAYERAI_API_KEYS", "")

def validate_key(api_key: str) -> bool:
    if not ALLOWED_TOKEN:
        return True
    return api_key.strip() == ALLOWED_TOKEN.strip()

# ── System prompt ─────────────────────────────────────────────────────────────
SYSTEM_PROMPT = (
    "You are a human-like player in a multiplayer chat environment. "
    "Respond casually, with short informal messages and natural tone. "
    "Use lowercase, slang, and act like a real person chatting."
)

# ── Core generation (runs in thread pool) ─────────────────────────────────────
def _run_generation(model_instance, prompt: str, streamer: TextIteratorStreamer) -> None:
    try:
        inputs = tokenizer(
            prompt,
            return_tensors = "pt",
            truncation     = True,
            max_length     = 512,
        )
        with torch.inference_mode():
            model_instance.generate(
                input_ids      = inputs.input_ids,
                attention_mask = inputs.attention_mask,
                streamer       = streamer,
                **GEN_CONFIG,
            )
    except Exception as e:
        logger.error(f"Generation error: {e}")
        streamer.text_queue.put(streamer.stop_signal)
        raise

# ── Streaming response ────────────────────────────────────────────────────────
def respond(message: str, history: list, api_key: str):
    if not validate_key(api_key):
        yield "[error] unauthorized"
        return

    if not queue_counter.acquire(blocking=False):
        yield "[error] server busy — try again shortly"
        return

    try:
        # Trim history to keep context short (long context = slow CPU prefill)
        trimmed = history[-4:] if len(history) > 4 else history

        messages = [{"role": "system", "content": SYSTEM_PROMPT}]
        messages.extend(trimmed)
        messages.append({"role": "user", "content": message})

        try:
            prompt = tokenizer.apply_chat_template(
                messages,
                tokenize              = False,
                add_generation_prompt = True,
            )
        except Exception as e:
            yield f"[error] template: {e}"
            return

        # Get a free model instance
        idx, model_instance = acquire_free_model(timeout=60.0)
        if model_instance is None:
            yield "[error] all models busy — try again"
            return

        response = ""
        try:
            streamer = TextIteratorStreamer(
                tokenizer,
                skip_prompt         = True,
                skip_special_tokens = True,
                timeout             = 120.0,
            )

            future = executor.submit(_run_generation, model_instance, prompt, streamer)

            for token in streamer:
                response += token
                yield response

            # Raise any exception from the worker thread
            future.result(timeout=10)

        except Exception as e:
            logger.error(f"respond() error on model #{idx}: {e}")
            if not response:
                yield "[error] generation failed"
        finally:
            model_locks[idx].release()

    finally:
        queue_counter.release()


# ── Warmup (runs in background, errors are non-fatal) ────────────────────────
def _warmup():
    logger.info("Warmup starting for all model copies...")
    try:
        dummy_prompt = tokenizer.apply_chat_template(
            [{"role": "system", "content": SYSTEM_PROMPT},
             {"role": "user",   "content": "hi"}],
            tokenize              = False,
            add_generation_prompt = True,
        )
        inputs = tokenizer(
            dummy_prompt,
            return_tensors = "pt",
            truncation     = True,
            max_length     = 128,
        )
        warmup_cfg = {**GEN_CONFIG, "max_new_tokens": 10}

        for idx, m in enumerate(models):
            t0 = time.time()
            try:
                with torch.inference_mode():
                    m.generate(
                        input_ids      = inputs.input_ids,
                        attention_mask = inputs.attention_mask,
                        **warmup_cfg,
                    )
                logger.info(f"Model #{idx} warmup done in {time.time()-t0:.1f}s ✓")
            except Exception as e:
                logger.warning(f"Model #{idx} warmup failed (non-fatal): {e}")

        logger.info("All warmups complete ✓")

    except Exception as e:
        logger.warning(f"Warmup failed (non-fatal): {e}")

Thread(target=_warmup, daemon=True).start()


# ── CSS ───────────────────────────────────────────────────────────────────────
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Press+Start+2P&display=swap');
* { font-family: 'Press Start 2P', monospace !important; border-radius: 0 !important; }
html, body, .gradio-container, .main, .wrap, gradio-app {
    background-color: #000000 !important; color: #ffffff !important; }
.gradio-container { max-width: 800px !important; margin: 0 auto !important; padding: 10px !important; }
.chatbot, [data-testid="chatbot"], .message-wrap, .messages {
    background-color: #000000 !important; border: none !important; box-shadow: none !important; }
.message.user .message-bubble-border, div[data-testid="user"] .message-bubble-border, .message.user {
    background-color: #00d000 !important; color: #000000 !important; border: none !important;
    font-size: 12px !important; line-height: 1.6 !important; padding: 10px 14px !important; }
.message.bot .message-bubble-border, div[data-testid="bot"] .message-bubble-border, .message.bot {
    background-color: #ffffff !important; color: #000000 !important; border: none !important;
    font-size: 12px !important; line-height: 1.6 !important; padding: 10px 14px !important; }
.message.user *, .message.bot * { color: #000000 !important; background-color: transparent !important; }
.avatar-container, .avatar-container img, img.avatar-image {
    width: 36px !important; height: 36px !important; min-width: 36px !important;
    background: transparent !important; image-rendering: pixelated !important;
    border: none !important; border-radius: 0 !important; }
textarea, input[type="text"] {
    background-color: #111111 !important; color: #00ff00 !important;
    border: 2px solid #00d000 !important; font-size: 12px !important; padding: 12px !important; }
button { background-color: #00d000 !important; color: #000000 !important;
    border: 2px solid #00d000 !important; font-size: 10px !important; }
button:hover { background-color: #00ff00 !important; }
h1, h2, h3 { color: #00ff00 !important; text-align: center; font-size: 16px !important; }
#lock-screen {
    position: fixed !important; top: 0 !important; left: 0 !important;
    width: 100vw !important; height: 100vh !important;
    background-color: #000000 !important; z-index: 9999 !important;
    display: flex !important; flex-direction: column !important;
    align-items: center !important; justify-content: center !important; gap: 20px !important; }
footer { display: none !important; }
.block, .form, .gap { background-color: #000000 !important; border: none !important; }
"""

USER_AVATAR = "face.png"
BOT_AVATAR  = "question.png"

# ── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(css=custom_css, theme=gr.themes.Base()) as demo:

    with gr.Column(elem_id="lock-screen", visible=True) as lock_screen:
        gr.Markdown("# >> player_ai <<")
        gr.Markdown("### enter access key to continue")
        key_input  = gr.Textbox(placeholder="access key...", type="password",
                                show_label=False, max_lines=1)
        unlock_btn = gr.Button("UNLOCK")
        lock_msg   = gr.Markdown("")

    with gr.Column(visible=False) as chat_screen:
        gr.Markdown("# >> player_ai <<")
        stored_key = gr.Textbox(value="", visible=False)

        chatbot = gr.Chatbot(
            type              = "messages",
            avatar_images     = (USER_AVATAR, BOT_AVATAR),
            height            = 500,
            show_label        = False,
            bubble_full_width = False,
            show_copy_button  = False,
        )

        gr.ChatInterface(
            fn                          = respond,
            chatbot                     = chatbot,
            type                        = "messages",
            additional_inputs           = [stored_key],
            additional_inputs_accordion = gr.Accordion(visible=False),
        )

    def try_unlock(key):
        if validate_key(key):
            return gr.update(visible=False), gr.update(visible=True), key, ""
        return gr.update(visible=True), gr.update(visible=False), "", "❌ invalid key"

    for trigger in (unlock_btn.click, key_input.submit):
        trigger(fn=try_unlock, inputs=[key_input],
                outputs=[lock_screen, chat_screen, stored_key, lock_msg])

if __name__ == "__main__":
    demo.launch(
        max_threads = MAX_CONCURRENT * 4,
        show_error  = True,
    )