import queue
import threading
import time
import logging

from bot.telegram_api import (
    send_message,
    edit_message,
    delete_message,
    md_to_telegram_html,
    split_for_telegram,
)
from model.qwen import generate_stream
from db import models as dbm

log = logging.getLogger("streaming")

EDIT_THROTTLE_SEC = 1.0
CURSOR = " ▌"
PREVIEW_LIMIT = 3500
QUEUE_TIMEOUT = 300
FALLBACK_PLAIN_LIMIT = 3900


async def stream_reply(chat_id, user_id, user_text):
    try:
        window = int(dbm.get_setting("conversation_window", "6"))
    except Exception:
        window = 6
    try:
        max_new = int(dbm.get_setting("max_new_tokens", "1024"))
    except Exception:
        max_new = 1024
    try:
        temperature = float(dbm.get_setting("temperature", "0.7"))
    except Exception:
        temperature = 0.7
    try:
        top_p = float(dbm.get_setting("top_p", "0.9"))
    except Exception:
        top_p = 0.9
    system_prompt = dbm.get_setting("system_prompt", "You are a helpful assistant.")

    history = dbm.get_history(user_id, window)
    messages = [{"role": "system", "content": system_prompt}] + history

    sent = await send_message(chat_id, CURSOR, parse_mode=None)
    if not sent.get("ok"):
        log.error(f"sendMessage placeholder failed: {sent}")
        return
    placeholder_id = sent["result"]["message_id"]

    buffer = ""
    last_edit = 0.0
    gpu_start = time.time()
    first_token_at = None
    error = None

    q = queue.Queue()

    def _producer():
        try:
            for chunk in generate_stream(messages, max_new, temperature, top_p):
                q.put(("chunk", chunk))
        except Exception as e:
            q.put(("error", repr(e)))
        finally:
            q.put(("done", None))

    threading.Thread(target=_producer, daemon=True).start()

    try:
        while True:
            try:
                kind, payload = q.get(timeout=QUEUE_TIMEOUT)
            except queue.Empty:
                error = "timeout"
                break

            if kind == "chunk":
                if first_token_at is None:
                    first_token_at = time.time()
                buffer += payload
                now = time.time()
                if now - last_edit >= EDIT_THROTTLE_SEC:
                    preview = buffer
                    if len(preview) > PREVIEW_LIMIT:
                        preview = preview[-PREVIEW_LIMIT:]
                    await _safe_edit_plain(chat_id, placeholder_id, preview + CURSOR)
                    last_edit = now
            elif kind == "error":
                error = payload
                break
            elif kind == "done":
                break
    except Exception as e:
        error = repr(e)
        log.exception("stream_reply failure")

    gpu_ms = int((time.time() - gpu_start) * 1000)
    latency_ms = int((time.time() - (first_token_at or gpu_start)) * 1000)

    if error:
        await _safe_edit_plain(chat_id, placeholder_id, f"⚠️ Error: {error}")
        dbm.log_gpu(user_id, gpu_ms, "error")
        return

    if not buffer.strip():
        await _safe_edit_plain(chat_id, placeholder_id, "⚠️ Empty response from model.")
        dbm.log_gpu(user_id, gpu_ms, "empty")
        return

    await delete_message(chat_id, placeholder_id)

    delivered = await _deliver_final(chat_id, buffer)
    if not delivered:
        log.error("Final delivery failed entirely")
        await send_message(
            chat_id,
            "⚠️ Could not deliver the response (formatting error). Full text below:",
            parse_mode=None,
        )
        for chunk in _chunk_plain(buffer, FALLBACK_PLAIN_LIMIT):
            await send_message(chat_id, chunk, parse_mode=None)

    # ---- Persist ----
    dbm.add_message(
        user_id, "assistant", buffer,
        latency_ms=latency_ms, gpu_time_ms=gpu_ms,
    )
    dbm.archive_append(user_id, "assistant", buffer, latency_ms=latency_ms)
    dbm.increment_user_turns(user_id)
    dbm.log_gpu(user_id, gpu_ms, "ok")

    user = dbm.get_user(user_id)
    if user and user.get("token"):
        dbm.bump_token_use(user["token"])


async def _deliver_final(chat_id, raw_text: str) -> bool:
    try:
        html_text = md_to_telegram_html(raw_text)
        chunks = split_for_telegram(html_text)
    except Exception:
        log.exception("md/split failed for final delivery")
        return False

    sent_count = 0
    for i, chunk in enumerate(chunks):
        r = await send_message(chat_id, chunk, parse_mode="HTML")
        if r.get("ok"):
            sent_count += 1
            continue
        if i == 0:
            log.warning(f"Final chunk 0 HTML failed: {r}")
            return False
        log.warning(f"Final chunk {i} HTML failed, retrying plain: {r}")
        r2 = await send_message(chat_id, _strip_html(chunk), parse_mode=None)
        if r2.get("ok"):
            sent_count += 1
            continue
        log.error(f"Plain fallback for chunk {i} also failed: {r2}")
        return False

    return sent_count == len(chunks)


def _strip_html(s: str) -> str:
    import re
    return re.sub(r"<[^>]+>", "", s)


def _chunk_plain(text: str, max_len: int):
    chunks = []
    remaining = text
    while len(remaining) > max_len:
        cut = remaining.rfind("\n", 0, max_len)
        if cut == -1 or cut < max_len * 0.5:
            cut = max_len
        chunks.append(remaining[:cut])
        remaining = remaining[cut:].lstrip("\n")
    if remaining:
        chunks.append(remaining)
    return chunks


async def _safe_edit_plain(chat_id, message_id, text):
    try:
        r = await edit_message(
            chat_id, message_id, text,
            parse_mode=None,
        )
        if not r.get("ok"):
            log.warning(f"edit (plain) failed: {r}")
    except Exception:
        log.exception("plain edit failed")