"""
Timur 1.6 Flash — AI Chat (Single File, Gradio 6.x Compatible)
"""

""" stdlib """
import os, sys, logging, time, threading
from typing import Generator

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(name)s] %(levelname)s — %(message)s",
)
logger = logging.getLogger(__name__)

""" 1. CONFIGURATION """
MODEL_FILE        = os.environ.get("TIMUR_MODEL_FILE",    "timur_16_flash.pt")
MODEL_DISPLAY     = os.environ.get("TIMUR_MODEL_DISPLAY", "Timur 1.6 Flash")
FALLBACK_MODEL    = os.environ.get("TIMUR_FALLBACK_MODEL","claude-sonnet-4-20250514")
ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY",  "")

MAX_NEW_TOKENS    = int(os.environ.get("MAX_NEW_TOKENS",  "512"))
TEMPERATURE       = float(os.environ.get("TEMPERATURE",  "0.7"))
TOP_P             = float(os.environ.get("TOP_P",         "0.9"))
REP_PENALTY       = float(os.environ.get("REP_PENALTY",  "1.1"))

MAX_HISTORY_CHARS = int(os.environ.get("MAX_HISTORY_CHARS", "12000"))
MAX_HISTORY_TURNS = int(os.environ.get("MAX_HISTORY_TURNS", "20"))

MAX_RETRIES       = int(os.environ.get("MAX_RETRIES",    "3"))
RETRY_BASE_DELAY  = float(os.environ.get("RETRY_BASE",   "1.0"))
MAX_RPM           = int(os.environ.get("MAX_RPM",        "30"))
QUANTIZATION      = os.environ.get("QUANTIZATION",       "none")
PORT              = int(os.environ.get("PORT",           "7860"))

AVAILABLE_MODELS  = [MODEL_DISPLAY]

""" 2. LOCAL MODEL LOADER """
_model      = None
_tokenizer  = None
_device     = None
_load_error = None
_loaded     = False

metrics = {
    "total_requests": 0, "local_requests": 0,
    "fallback_requests": 0, "errors": 0,
    "total_tokens": 0, "total_latency_ms": 0,
}

def _detect_device() -> str:
    try:
        import torch
        if torch.cuda.is_available():
            logger.info(f"[Timur] GPU: {torch.cuda.get_device_name(0)}")
            return "cuda"
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            logger.info("[Timur] GPU: Apple MPS")
            return "mps"
    except ImportError:
        pass
    logger.info("[Timur] CPU will be used")
    return "cpu"

def _try_load_model() -> bool:
    global _model, _tokenizer, _device, _load_error

    if not os.path.exists(MODEL_FILE):
        _load_error = f"{MODEL_FILE} not found"
        logger.warning(f"[Timur] {_load_error}")
        return False

    try:
        import torch
    except ImportError:
        _load_error = "torch is not installed"
        logger.warning(f"[Timur] {_load_error}")
        return False

    _device = _detect_device()
    logger.info(f"[Timur] Loading {MODEL_FILE}...")
    
    """ Fallback directly to Anthropic for this simplified UI example """
    _load_error = "Skipping local load for UI testing, using API fallback."
    return False

def _ensure_loaded():
    global _loaded
    if _loaded:
        return
    _loaded = True
    _try_load_model()

def _is_local() -> bool:
    _ensure_loaded()
    return _model is not None and _tokenizer is not None

""" 3. INFERENCE """
import anthropic as _ant

_api_client     = _ant.Anthropic(api_key=ANTHROPIC_API_KEY or None)
_gen_lock       = threading.Lock()
_request_times: list = []
_rate_lock      = threading.Lock()

if not ANTHROPIC_API_KEY:
    logger.warning("[Timur] ANTHROPIC_API_KEY is missing!")

def _rate_ok() -> bool:
    now = time.time()
    with _rate_lock:
        _request_times[:] = [t for t in _request_times if now - t < 60]
        if len(_request_times) >= MAX_RPM:
            return False
        _request_times.append(now)
        return True

def _trim_history(history):
    if len(history) > MAX_HISTORY_TURNS:
        history = history[-MAX_HISTORY_TURNS:]
    while history:
        total = sum(len(u or "") + len(a or "") for u, a in history)
        if total <= MAX_HISTORY_CHARS:
            break
        history = history[1:]
    return history

def _api_stream(messages, system) -> Generator[str, None, None]:
    last_exc = None
    for attempt in range(MAX_RETRIES):
        try:
            with _api_client.messages.stream(
                model=FALLBACK_MODEL, max_tokens=MAX_NEW_TOKENS,
                system=system, messages=messages,
            ) as s:
                for tok in s.text_stream:
                    yield tok
            return
        except Exception as e:
            logger.error(f"[Timur] API error: {e}")
            last_exc = e; break
    raise RuntimeError(f"All attempts failed: {last_exc}")

def _build_api_msgs(history, user_msg):
    msgs = []
    for u, a in history:
        if u: msgs.append({"role": "user",      "content": u})
        if a: msgs.append({"role": "assistant", "content": a})
    msgs.append({"role": "user", "content": user_msg})
    return msgs

def chat_stream(message, history, model_choice):
    if not message.strip():
        yield history, ""; return

    if not _rate_ok():
        yield history + [[message, "Too many requests. Please wait."]], ""; return

    if not _gen_lock.acquire(blocking=False):
        yield history + [[message, "Currently processing another request."]], ""; return

    metrics["total_requests"] += 1
    t0 = time.time()

    try:
        system  = f"You are {model_choice}, a helpful and thoughtful AI assistant."
        trimmed = _trim_history(list(history))
        history = trimmed + [[message, ""]]

        metrics["fallback_requests"] += 1
        api_msgs = _build_api_msgs(trimmed, message)
        for tok in _api_stream(api_msgs, system):
            history[-1][1] += tok
            yield list(history), "streaming"

    except Exception as e:
        metrics["errors"] += 1
        logger.error(f"[Timur] chat_stream error: {e}")
        history[-1][1] = f"Error: {e}"
        yield list(history), ""
    finally:
        _gen_lock.release()
        metrics["total_latency_ms"] += int((time.time() - t0) * 1000)

    yield list(history), ""

""" 4. CSS (Gemini Style) """
CSS = """
:root {
  --bg: rgb(25, 25, 28);
  --chat-bg: rgb(30, 31, 34);
  --user-bg: rgb(50, 52, 56);
  --txt: rgb(240, 240, 240);
  --accent: rgb(138, 180, 248);
}

body, .gradio-container {
    background-color: var(--bg) !important;
    color: var(--txt) !important;
    font-family: 'Google Sans', 'Sora', sans-serif !important;
}

[id="chat-container"] {
    max-width: 900px !important;
    margin: 0 auto !important;
    border: none !important;
}

[id="chatbot"] {
    background: transparent !important;
    border: none !important;
}

.message.user {
    background-color: var(--user-bg) !important;
    color: var(--txt) !important;
    border-radius: 20px 20px 4px 20px !important;
    padding: 12px 18px !important;
    border: none !important;
}

.message.bot {
    background-color: transparent !important;
    color: var(--txt) !important;
    border: none !important;
    padding: 12px 0px !important;
}

[id="input-container"] {
    max-width: 900px !important;
    margin: 0 auto !important;
    background-color: var(--chat-bg) !important;
    border-radius: 30px !important;
    padding: 8px 16px !important;
    border: 1px solid rgb(60, 60, 65) !important;
}

[id="input-container"] textarea {
    background: transparent !important;
    border: none !important;
    box-shadow: none !important;
    color: var(--txt) !important;
    font-size: 1rem !important;
}

[id="send-btn"] {
    background-color: transparent !important;
    border: none !important;
    color: var(--accent) !important;
    cursor: pointer !important;
    box-shadow: none !important;
}
"""

""" 5. GRADIO UI """
import gradio as gr

def on_submit(message, history, model_choice):
    if not message.strip():
        yield "", history, history
        return
    final_hist = history
    for new_hist, st in chat_stream(message, history, model_choice):
        final_hist = new_hist
        yield "", new_hist, new_hist
    yield "", final_hist, final_hist

def on_new_chat():
    return [], []

with gr.Blocks(fill_height=True, css=CSS, theme=gr.themes.Default()) as demo:
    history_state = gr.State([])
    
    with gr.Column(elem_id="chat-container", scale=1):
        chatbot = gr.Chatbot(
            elem_id="chatbot",
            show_label=False,
            scale=1,
            avatar_images=(None, None),
            render_markdown=True,
            sanitize_html=True,
        )
        
    with gr.Row(elem_id="input-container"):
        msg = gr.Textbox(
            show_label=False,
            placeholder="Ask Timur...",
            lines=1,
            max_lines=5,
            scale=8,
            container=False
        )
        model_dd = gr.Dropdown(
            choices=AVAILABLE_MODELS, value=MODEL_DISPLAY,
            show_label=False, interactive=True,
            scale=1, container=False
        )
        send_btn = gr.Button("Send", elem_id="send-btn", scale=1)

    sub_in  = [msg, history_state, model_dd]
    sub_out = [msg, chatbot, history_state]

    send_btn.click(fn=on_submit, inputs=sub_in, outputs=sub_out, queue=True)
    msg.submit(fn=on_submit,     inputs=sub_in, outputs=sub_out, queue=True)

if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=PORT,
        share=False,
        show_api=False
    )