import os
import re
from importlib import import_module, util
from threading import Lock
from typing import Any, cast
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# Prevent thread thrashing on standard/shared 2-vCPU hardware architectures
# NOTE: We removed torch entirely! It's no longer needed for CPU GGUF inference.

if util.find_spec("spaces"):
    spaces = import_module("spaces")
else:
    class _SpacesFallback:
        @staticmethod
        def GPU(duration=60):
            def _decorator(func):
                return func
            return _decorator
    spaces = _SpacesFallback()

HF_SECRET_TOKEN = os.environ.get("HF_TOKEN")
PRIVATE_MODEL_ID = "rotextechnologiesuk/nuclear-sqep"
GGUF_FILE = "rotex_nuclear_sqep_Q4_K_M.gguf"

GPU_DECORATOR_DURATION = int(os.environ.get("GPU_DECORATOR_DURATION", "60"))
MAX_CONTEXT_TURNS = int(os.environ.get("MAX_CONTEXT_TURNS", "8"))
MAX_INPUT_CHARS = int(os.environ.get("MAX_INPUT_CHARS", "512"))
DEFAULT_MAX_NEW_TOKENS = int(os.environ.get("DEFAULT_MAX_NEW_TOKENS", "768"))
MAX_NEW_TOKENS_LIMIT = int(os.environ.get("MAX_NEW_TOKENS_LIMIT", "2048"))

try:
    GRADIO_MAJOR_VERSION = int(gr.__version__.split(".")[0])
except (AttributeError, ValueError, IndexError):
    GRADIO_MAJOR_VERSION = 5

DEFAULT_SYSTEM_PROMPT = (
    "You are a Nuclear SQEP specialist assistant. "
    "Provide precise, practical, and safety-conscious answers. "
)

THINK_BLOCK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)
THINK_OPEN_RE = re.compile(r"<think>.*$", re.DOTALL | re.IGNORECASE)
DETAILS_BLOCK_RE = re.compile(r"<details>.*?</details>", re.DOTALL | re.IGNORECASE)
THINKING_PREFIX_RE = re.compile(
    r"^\s*(?:[^\w<\n]+\s*)?(?:thinking(?:\s+process)?\s*:?)",
    re.IGNORECASE,
)
ASSISTANT_PREFIX_RE = re.compile(r"^\s*assistant\s*:\s*", re.IGNORECASE)

print("Downloading GGUF file from Hugging Face Hub...")
# Safely download the specific target GGUF binary from the repository
model_path = hf_hub_download(
    repo_id=PRIVATE_MODEL_ID,
    filename=GGUF_FILE,
    token=HF_SECRET_TOKEN
)

print("Loading GGUF model via lightning-fast native llama_cpp core...")
# Initialize the high-speed C++ engine optimized to utilize exactly 2 CPU threads
model = Llama(
    model_path=model_path,
    n_ctx=4096,         # Set healthy context allocation window
    n_threads=2,        # Match hardware cores exactly to maximize processing throughput
    n_batch=512,        # Batch processing limits
    verbose=False       # Keeps terminal logs clean
)
MODEL_LOCK = Lock()

def _remove_hidden_blocks(text):
    cleaned = DETAILS_BLOCK_RE.sub("", text or "")
    cleaned = THINK_BLOCK_RE.sub("", cleaned)
    cleaned = THINK_OPEN_RE.sub("", cleaned)
    return cleaned.strip()

def _extract_visible_answer(raw_response):
    response_text = (raw_response or "").strip()
    hidden_parts = []
    for match in re.finditer(r"<think>(.*?)</think>", response_text, re.DOTALL | re.IGNORECASE):
        candidate = match.group(1).strip()
        if candidate:
            hidden_parts.append(candidate)
    cleaned = THINK_BLOCK_RE.sub("", response_text)
    cleaned = THINK_OPEN_RE.sub("", cleaned).strip()
    if THINKING_PREFIX_RE.match(cleaned):
        parts = re.split(r"\n\s*\n", cleaned, maxsplit=1)
        if len(parts) == 2:
            hidden_candidate = THINKING_PREFIX_RE.sub("", parts[0]).strip()
            if hidden_candidate:
                hidden_parts.append(hidden_candidate)
            cleaned = parts[1].strip()
        else:
            lines = cleaned.splitlines()
            if lines:
                hidden_candidate = THINKING_PREFIX_RE.sub("", lines[0]).strip()
                if hidden_candidate:
                    hidden_parts.append(hidden_candidate)
                cleaned = "\n".join(lines[1:]).strip()
    final_answer = ASSISTANT_PREFIX_RE.sub("", cleaned).strip() or "I am ready for your next request."
    hidden_trace = "\n\n".join(hidden_parts).strip()
    return final_answer, hidden_trace

def _build_messages(message, history, system_prompt):
    prompt = (system_prompt or "").strip() or DEFAULT_SYSTEM_PROMPT
    messages = [{"role": "system", "content": prompt}]
    history_slice = (history or [])[-MAX_CONTEXT_TURNS:]
    if history_slice and isinstance(history_slice[0], dict):
        for item in history_slice:
            role = str(item.get("role", "")).strip().lower()
            content = "" if item.get("content") is None else str(item.get("content")).strip()
            if not role or not content:
                continue
            if role == "assistant":
                content = _remove_hidden_blocks(content)
            if role in {"user", "assistant", "system"}:
                messages.append({"role": role, "content": content})
    else:
        for item in history_slice:
            if not isinstance(item, (list, tuple)) or len(item) < 2:
                continue
            user_content = "" if item[0] is None else str(item[0]).strip()
            assistant_content = "" if item[1] is None else str(item[1]).strip()
            if user_content:
                messages.append({"role": "user", "content": user_content})
            if assistant_content:
                messages.append(
                    {
                        "role": "assistant",
                        "content": _remove_hidden_blocks(assistant_content),
                    }
                )
    truncated_message = (message or "")[:MAX_INPUT_CHARS].strip()
    messages.append({"role": "user", "content": truncated_message})
    return messages

@spaces.GPU(duration=GPU_DECORATOR_DURATION)
def generate_response(message, history, system_prompt, temperature, top_p, max_new_tokens):
    user_message = (message or "")[:MAX_INPUT_CHARS].strip()
    if not user_message:
        return "", ""
    normalized_temperature = max(0.0, min(float(temperature), 1.0))
    normalized_top_p = max(0.1, min(float(top_p), 1.0))
    bounded_max_new_tokens = max(32, min(int(max_new_tokens), MAX_NEW_TOKENS_LIMIT))
    
    with MODEL_LOCK:
        try:
            messages = _build_messages(user_message, history, system_prompt)
            do_sample = normalized_temperature > 0.0
            
            # Execute ultra-fast, native C++ chat inference via llama_cpp
            response_obj = model.create_chat_completion(
                messages=messages,
                max_tokens=bounded_max_new_tokens,
                temperature=normalized_temperature if do_sample else 0.0,
                top_p=normalized_top_p,
                repeat_penalty=1.03,
            )
            
            raw_response = response_obj["choices"][0]["message"]["content"].strip()
            return _extract_visible_answer(raw_response)
        except Exception as exc:
            print(f"Inference Engine Exception encountered: {str(exc)}")
            return "An engine error occurred while processing this request. Please clear and try again.", ""

def respond(message, history, system_prompt, temperature, top_p, max_new_tokens):
    history = history or []
    user_message = (message or "")[:MAX_INPUT_CHARS].strip()
    if not user_message:
        return history, ""
    final_answer, _hidden_trace = generate_response(
        user_message,
        history,
        system_prompt,
        temperature,
        top_p,
        max_new_tokens,
    )
    normalized_history = []
    if history and isinstance(history[0], dict):
        for item in history:
            role = str(item.get("role", "")).strip().lower()
            content = "" if item.get("content") is None else str(item.get("content")).strip()
            if role in {"user", "assistant"} and content:
                normalized_history.append({"role": role, "content": content})
    else:
        for item in history:
            if not isinstance(item, (list, tuple)) or len(item) < 2:
                continue
            user_turn = "" if item[0] is None else str(item[0]).strip()
            assistant_turn = "" if item[1] is None else str(item[1]).strip()
            if user_turn:
                normalized_history.append({"role": "user", "content": user_turn})
            if assistant_turn:
                normalized_history.append({"role": "assistant", "content": assistant_turn})
    updated_history = normalized_history + [
        {"role": "user", "content": user_message},
        {"role": "assistant", "content": final_answer},
    ]
    return updated_history, ""

def clear_chat():
    return [], ""

APP_CSS = """:root {
    --lm-bg: #0c1117;
    --lm-panel: #121924;
    --lm-panel-2: #182232;
    --lm-border: #2a3548;
    --lm-text: #edf2fb;
    --lm-muted: #9fb0cb;
    --lm-accent: #58c5ff;
    --lm-accent-soft: #18384a;
}
.gradio-container {
    min-height: 100vh;
    color: var(--lm-text);
    background:
        radial-gradient(1200px 520px at 12% -20%, #1d334a 0%, transparent 55%),
        radial-gradient(980px 620px at 95% 120%, #21384e 0%, transparent 60%),
        var(--lm-bg);
    font-family: Arial, "Helvetica Neue", Helvetica, sans-serif;
}
#lm-shell-header {
    margin-bottom: 14px;
    padding: 18px 20px;
    border: 1px solid var(--lm-border);
    border-radius: 14px;
    background: linear-gradient(140deg, rgba(16, 24, 37, 0.92), rgba(20, 32, 48, 0.92));
    display: flex;
    align-items: center;
    justify-content: space-between;
    gap: 12px;
    backdrop-filter: blur(6px);
}
.lm-title {
    font-size: 1.05rem;
    font-weight: 700;
    letter-spacing: 0.04em;
    text-transform: uppercase;
}
.lm-subtitle {
    margin-top: 4px;
    color: var(--lm-muted);
    font-size: 0.9rem;
}
.lm-badge {
    border: 1px solid #3d4f6b;
    border-radius: 999px;
    padding: 8px 14px;
    background: rgba(19, 33, 50, 0.8);
    color: #c9d8ef;
    font-family: "IBM Plex Mono", "Consolas", monospace;
    font-size: 0.75rem;
}
#lm-shell-body {
    gap: 14px;
}
#lm-sidebar {
    border: 1px solid var(--lm-border);
    border-radius: 14px;
    padding: 16px;
    background: linear-gradient(180deg, rgba(18, 26, 38, 0.95), rgba(14, 20, 30, 0.95));
}
#lm-main {
    border: 1px solid var(--lm-border);
    border-radius: 14px;
    padding: 16px;
    background: linear-gradient(180deg, rgba(17, 23, 34, 0.95), rgba(14, 19, 28, 0.95));
}
#lm-chat-window {
    border: 1px solid var(--lm-border);
    border-radius: 12px;
    background: radial-gradient(circle at 50% -40%, rgba(32, 54, 74, 0.5), rgba(16, 22, 32, 0.96));
}
#lm-chat-window .message.user {
    background: #1f3046;
    border: 1px solid #3a4e68;}
#lm-chat-window .message.bot {
    background: #16202d;
    border: 1px solid #2c3b52;}
#lm-input-row {
    margin-top: 12px;
    align-items: stretch;
    gap: 10px;}
#lm-input-box textarea {
    border: 1px solid var(--lm-border);
    border-radius: 10px;
    background: rgba(15, 22, 32, 0.9);
    color: var(--lm-text);
    font-family: Arial, "Helvetica Neue", Helvetica, sans-serif;}
#lm-send-button button,.lm-primary-btn button {
    border: 1px solid #4a7b96;
    border-radius: 10px;
    background: linear-gradient(120deg, var(--lm-accent-soft), #21465f);
    color: #eaf6ff;
    font-weight: 600;}
#lm-send-button button:hover,.lm-primary-btn button:hover {
    border-color: var(--lm-accent);
    box-shadow: 0 0 0 1px rgba(88, 197, 255, 0.35) inset;}
#lm-sidebar .gr-markdown,#lm-sidebar .gradio-markdown {
    color: var(--lm-text);}
#lm-sidebar .gr-markdown p,#lm-sidebar .gradio-markdown p {
    color: var(--lm-muted);}
@keyframes rise-in {
    from {
        opacity: 0;
        transform: translateY(8px);
    }    to {
        opacity: 1;
        transform: translateY(0);
    }}#lm-shell-header,#lm-sidebar,#lm-main {
    animation: rise-in 0.45s ease both;}
#lm-main {
    animation-delay: 0.08s;}
@media (max-width: 960px) {
    #lm-shell-header {
        flex-direction: column;
        align-items: flex-start;
    }    .lm-badge {
        width: 100%;
        text-align: center;
    }}"""

blocks_kwargs = {
    "title": "Rotex LLM Nuclear SQEP Shell",
    "fill_height": True,
}
if GRADIO_MAJOR_VERSION < 6:
    blocks_kwargs["css"] = APP_CSS

with gr.Blocks(**blocks_kwargs) as demo:
    gr.HTML(
        """
        <div id="lm-shell-header">
            <div>
                <div class="lm-title">Rotex Nuclear SQEP Assistant</div>
                <div class="lm-subtitle">LLM model for Nuclear SQEP response.</div>
            </div>
        </div>
        """
    )
    with gr.Row(elem_id="lm-shell-body"):
        with gr.Column(scale=1, min_width=310, elem_id="lm-sidebar"):
            gr.Markdown("## Session Controls")
            clear_button = gr.Button("New Chat", elem_classes=["lm-primary-btn"])
            system_prompt = gr.Textbox(
                label="System prompt",
                value=DEFAULT_SYSTEM_PROMPT,
                lines=7,
            )
            gr.Markdown("## Generation Settings")
            temperature = gr.Slider(
                minimum=0.0,
                maximum=1.0, 
                value=0.6,
                step=0.05,
                label="Temperature",
            )
            top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.01,
                label="Top-p",
            )
            max_new_tokens = gr.Slider(
                minimum=64,
                maximum=MAX_NEW_TOKENS_LIMIT,
                value=DEFAULT_MAX_NEW_TOKENS,
                step=32,
                label="Max new tokens",
            )
        with gr.Column(scale=3, elem_id="lm-main"):
            chatbot = gr.Chatbot(
                label="",
                height=640,
                elem_id="lm-chat-window",
            )
            with gr.Row(elem_id="lm-input-row"):
                user_input = gr.Textbox(
                    show_label=False,
                    placeholder="Message the assistant...",
                    lines=3,
                    max_lines=8,
                    elem_id="lm-input-box",
                )
                send_button = gr.Button("Send", variant="primary", elem_id="lm-send-button", min_width=60, scale=0)

    send_inputs = [
        user_input,
        chatbot,
        system_prompt,
        temperature,
        top_p,
        max_new_tokens,
    ]

    send_button.click(
        fn=respond,
        inputs=send_inputs,
        outputs=[chatbot, user_input],
    )
    user_input.submit(
        fn=respond,
        inputs=send_inputs,
        outputs=[chatbot, user_input],
    )
    clear_button.click(
        fn=clear_chat,
        outputs=[chatbot, user_input],
        queue=False,
    )

if __name__ == "__main__":
    queued_demo = demo.queue(default_concurrency_limit=1)
    if GRADIO_MAJOR_VERSION >= 6:
        cast(Any, queued_demo).launch(css=APP_CSS)
    else:
        queued_demo.launch()