"""LLM stream normalization, idle timeouts, and token-bucket limiting."""
from __future__ import annotations

import logging
import os
import re
import threading
import time
import unicodedata
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime

_log = logging.getLogger(__name__)

_LLM_RETRY_MAX_DELAY = 30.0
_LLM_RETRY_ATTEMPTS = 3
_LLM_RETRY_BASE_DELAY = 1.0  # seconds

_STREAM_IDLE_TIMEOUT = 180  # seconds — base max silence between chunks before treating stream as stalled
_STREAM_IDLE_TIMEOUT_PER_10K_TOKENS = 40  # extra seconds per 10K estimated tokens (prompt eval takes longer on big contexts)
_STREAM_IDLE_TIMEOUT_MAX = 600  # hard cap — never wait more than 10 minutes for a single chunk
# Time-to-first-token cap (subsequent chunks use scaled idle). Keeps UI from hanging when MLX is contended.
_STREAM_FIRST_CHUNK_TIMEOUT_MLX = 120
_STREAM_FIRST_CHUNK_TIMEOUT_DEFAULT = 90

# ─── LLM token-bucket rate limiter ────────────────────────────────────────────
# Prevents burst hammering of the remote LLM endpoint during multi-round research
# or rapid tool-call loops.  Capacity = 5 allows short bursts; refill = 1/2s gives
# a sustained rate of 30 requests/minute, well within normal server limits.
_LLM_BUCKET_CAPACITY = 5.0     # max tokens in the bucket
_LLM_BUCKET_REFILL_RATE = 0.5  # tokens per second  (1 token per 2 s)
_llm_bucket_tokens: float = _LLM_BUCKET_CAPACITY
_llm_bucket_last_refill: float = time.monotonic()
_llm_bucket_lock = threading.Lock()


def _llm_bucket_acquire() -> None:
    """Block until one token is available in the LLM rate-limit bucket."""
    global _llm_bucket_tokens, _llm_bucket_last_refill
    while True:
        with _llm_bucket_lock:
            now = time.monotonic()
            elapsed = now - _llm_bucket_last_refill
            _llm_bucket_tokens = min(
                _LLM_BUCKET_CAPACITY,
                _llm_bucket_tokens + elapsed * _LLM_BUCKET_REFILL_RATE,
            )
            _llm_bucket_last_refill = now
            if _llm_bucket_tokens >= 1.0:
                _llm_bucket_tokens -= 1.0
                return
        # Bucket empty; sleep for roughly the time to produce one token, then retry.
        time.sleep(1.0 / _LLM_BUCKET_REFILL_RATE * 0.5)


class StreamIdleTimeout(Exception):
    """Raised when no chunk arrives within _STREAM_IDLE_TIMEOUT seconds."""


def _scaled_idle_timeout(estimated_tokens: int) -> int:
    """Return a stream idle timeout scaled to context size.

    Local models need significantly more time for prompt evaluation on large
    contexts. Base timeout handles normal cases; extra time is added per 10K tokens.
    """
    extra = int((estimated_tokens / 10_000) * _STREAM_IDLE_TIMEOUT_PER_10K_TOKENS)
    return min(_STREAM_IDLE_TIMEOUT + extra, _STREAM_IDLE_TIMEOUT_MAX)


def _first_chunk_idle_timeout(scaled_idle: int, *, is_mlx: bool) -> int:
    cap = _STREAM_FIRST_CHUNK_TIMEOUT_MLX if is_mlx else _STREAM_FIRST_CHUNK_TIMEOUT_DEFAULT
    return min(int(scaled_idle), cap)


def _iter_stream_with_timeout(
    stream,
    idle_timeout=_STREAM_IDLE_TIMEOUT,
    *,
    first_chunk_timeout: int | None = None,
):
    """Wrap an OpenAI stream iterator with a per-chunk idle deadline.

    Yields chunks from *stream*.  If the underlying iterator blocks for longer
    than the active deadline without producing a value, raises
    ``StreamIdleTimeout`` so the caller can salvage partial results.

    *first_chunk_timeout* applies only to the first chunk (time-to-first-token);
    later chunks use *idle_timeout*.
    """
    it = iter(stream)
    executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="stream_timeout")
    first_chunk = True
    try:
        while True:
            deadline = (
                first_chunk_timeout
                if first_chunk and first_chunk_timeout is not None
                else idle_timeout
            )
            future = executor.submit(next, it)
            try:
                yield future.result(timeout=deadline)
                first_chunk = False
            except StopIteration:
                return
            except TimeoutError:
                future.cancel()
                raise StreamIdleTimeout(
                    f"LLM stream stalled — no chunk received for {deadline}s"
                )
    finally:
        executor.shutdown(wait=False)


def _is_retryable_llm_error(e: BaseException) -> bool:
    """True if we should retry the LLM call (5xx, 429, connection, timeout)."""
    tname = type(e).__name__
    if tname in ("APIConnectionError", "APITimeoutError", "ConnectionError", "TimeoutError"):
        return True
    if "Timeout" in tname or "Connection" in tname:
        return True
    code = getattr(e, "status_code", None)
    if code is not None:
        return code == 429 or (500 <= code < 600)
    err_str = str(e).lower()
    if "timeout" in err_str or "connection" in err_str or "502" in err_str or "503" in err_str or "429" in err_str:
        return True
    return False


def _is_retryable_stream_error(e: BaseException) -> bool:
    """True if a stream crash (e.g. connection reset) may be due to context overflow and we should trim and retry."""
    tname = type(e).__name__
    if tname in ("ReadError", "ConnectError", "APIError") or "ReadError" in tname or "ConnectError" in tname:
        return True
    err_str = str(e).lower()
    if (
        "connection reset" in err_str
        or "errno 54" in err_str
        or "connection" in err_str
        or "an error occurred during streaming" in err_str
        or "incomplete chunked read" in err_str
    ):
        return True
    return False


def _is_kraken_error(result: str) -> bool:
    """True if the OCR result is a Kraken failure message (so we can fallback to complex OCR)."""
    if not result or not isinstance(result, str):
        return False
    s = result.strip()
    return (
        s.startswith("Kraken OCR failed")
        or s.startswith("Kraken OCR (venv) failed")
        or s.startswith("Kraken OCR is not available")
        or s.startswith("Kraken OCR timed out")
    )


_INLINE_CODE_PLACEHOLDER = "\uE000"
_MARKDOWN_LINK_PLACEHOLDER = "\uE001"
_FILEPATH_PLACEHOLDER = "\uE002"
_BARE_URL_PLACEHOLDER = "\uE003"

try:
    from persistence import MUGI_LOGS_DIR
except ImportError:
    MUGI_LOGS_DIR = os.path.join(os.path.expanduser("~"), ".mugi", "logs")


# Stream merges: literal (wrong, right) only. Add new pairs when you see them; no regex word-splitting.
# General approach: structural fixes (punctuation, markdown, spaces) + NFKC + this small list.
# Technical terms and Kraken/OCR error text often get split by streaming tokenization (e.g. "Kr aken", "Num Py").
_STREAM_MERGE_FIXES = (
    ("adistributed", "a distributed"),
    ("aversion control", "a version control"),  # phrase only; keeps "aversion" (dislike) intact
    ("witnessto", "witness to"),
    ("tobecome", "to become"),
    ("becausememory", "because memory"),
    ("resolutionwith", "resolution with"),
    ("systemthat", "system that"),
    ("justreinforcingit", "just reinforcing it"),
    ("Tryaged", "Try aged"),
    ("SSDat", "SSD at"),
    ("andrefurbished", "and refurbished"),
    ("war rior", "warrior"),
    ("mo ody", "moody"),
    # Kraken/OCR and Python env error text (streaming often splits camelCase / PascalCase)
    ("Kr aken", "Kraken"),
    ("Num Py", "NumPy"),
    ("Sci Py", "SciPy"),
    ("c math", "cmath"),
    ("there 's", "there's"),
    ("there ' re", "there're"),
    ("it 's", "it's"),
    ("don 't", "don't"),
    ("can 't", "can't"),
    ("isn 't", "isn't"),
    ("won 't", "won't"),
)

# Compiled regexes for normalize_stream_text (avoid recompiling on every call).
_RE_INLINE_CODE = re.compile(r"`[^`]*`")
_RE_MARKDOWN_LINK = re.compile(r"\[[^\]]*\]\([^)]*\)")
_RE_BARE_URL = re.compile(r"https?://[^\s\)\]>]+")
_RE_BARE_FILEPATH = re.compile(r"(?<![:\w/])~?/[\w.\-/]+\w")
_RE_FILENAME_EXT = re.compile(
    r"\b[\w][\w.\-]*\."
    r"(?:md|txt|json|py|swift|js|ts|jsx|tsx|yaml|yml|toml|csv|pdf|zip|rtf|html|css"
    r"|log|xml|env|cfg|ini|sh|zsh|rs|go|rb|java|c|h|cpp|hpp|gz|tar|bz2|png|jpg"
    r"|jpeg|gif|svg|webp|mp4|wav|doc|docx|xls|xlsx|sql|db|sqlite|lock)\b",
    re.IGNORECASE,
)
_RE_SENTENCE_SPACE = re.compile(r"([.!?:\uFF1F\uFF01])([A-Za-z])")
_RE_SPACE_BEFORE_PUNCT = re.compile(r"(\S)\s+([.,!?;:\u2014\uFF0E\uFF0C\uFF1F\uFF01])")
_RE_CONTRACTION_APOSTROPHE = re.compile(r"(\w+)\s+['\u2019\u2018]\s*(\w+)")
_RE_LIST_PERIOD_LETTER = re.compile(r"(\d)\.([A-Za-z\uE000\uE001])")
_RE_WORD_LIST_NUM = re.compile(r"([A-Za-z])(\d+\.)(?=\s|$|[A-Za-z\uE000\uE001])")
_RE_HYPHEN_WORDS = re.compile(r"(\w)[^\S\n]+-[^\S\n]*(\w)")
_RE_ITALIC_SPACES = re.compile(r"\*\s+(\w+)\s+\*")
_RE_COLLAPSE_SPACES = re.compile(r" {2,}")
_RE_LIST_MARKER_SPACE = re.compile(r"(\n|^)[^\S\n]*([-*+])\s+", re.MULTILINE)
_RE_LIST_MARKER_NO_SPACE = re.compile(r"(\n|^)[^\S\n]*([-+]|\*(?!\*))([^\s\uE000\uE001])", re.MULTILINE)
_RE_HEADING_SPACE = re.compile(r"(\n|^)(#{1,6})([A-Za-z0-9*_])", re.MULTILINE)
_RE_PAREN_LETTER = re.compile(r"\)([A-Za-z-])")
_RE_WORD_OPEN_PAREN = re.compile(r"([A-Za-z])\(")
_RE_DIGIT_OPEN_PAREN = re.compile(r"(\d)\(")
_RE_THE_NUMBER = re.compile(r"\bthe(\d+)")
_RE_A_NUMBER = re.compile(r"\ba(\d{2,})")
_RE_WORD_TWO_DIGITS = re.compile(r"([A-Za-z])(\d{2,})")
_RE_WORD_PERCENT = re.compile(r"([A-Za-z])(\d+%)")
_RE_BLOCKQUOTE_SPACE = re.compile(r"(\n|^)\s*>\s+", re.MULTILINE)
_RE_BOLD_AFTER = re.compile(r"(?<=\s)\*\* +(?=\w)")
_RE_BOLD_BEFORE = re.compile(r"(?<=\w) +\*\*(?=\s|$)")
_RE_UNDER_AFTER = re.compile(r"(?<=\s)__ +(?=\w)")
_RE_UNDER_BEFORE = re.compile(r"(?<=\w) +__(?=\s|$)")
_RE_LIST_NUM_BOLD = re.compile(r"(\d)\.(\*+)")
_RE_A_AND_BOLD = re.compile(r"\b(a|and)(\*\*)")
_RE_HEADING_BOLD = re.compile(r"(\n|^)(#{1,6})(\*\*|__)", re.MULTILINE)
_RE_LIST_MARKER_BOLD = re.compile(r"(\n|^)([-*+])(\*\*|__)", re.MULTILINE)


def normalize_stream_text(text: str) -> str:
    """Lightweight normalization to fix spacing that breaks markdown in WKWebView.
    Streamed chunks often lack spaces at token boundaries (e.g. 'bonfire.Quick', 'idea:Tryaged').
    Idempotent and safe to call on every streamed chunk. Only cosmetic whitespace changes."""
    # Coerce to str so normalizer never raises on type (e.g. when run from py2app bundle).
    if text is None:
        return ""
    if not isinstance(text, str):
        text = str(text)
    if not text or not text.strip():
        return text
    # 0. Unicode normalization so all steps see a consistent form (spaces, accents, compatibility chars)
    s = unicodedata.normalize("NFKC", text)
    # 0b. Normalize Unicode apostrophes/quotes to ASCII so contraction and punctuation rules match (e.g. ' (U+2019) -> ')
    s = s.replace("\u2019", "'").replace("\u2018", "'")
    # 1. Protect inline code: mask `...` segments so we don't alter spaces inside them
    placeholders = []
    def _mask(m):
        placeholders.append(m.group(0))
        return _INLINE_CODE_PLACEHOLDER

    s = _RE_INLINE_CODE.sub(_mask, s)
    # 1b. Protect markdown links [text](url) so we don't insert spaces in URLs or link text (e.g. "domain.com" -> "domain. com")
    link_placeholders = []
    def _mask_link(m):
        link_placeholders.append(m.group(0))
        return _MARKDOWN_LINK_PLACEHOLDER
    s = _RE_MARKDOWN_LINK.sub(_mask_link, s)
    # 1b2. Protect bare https?:// URLs so sentence-spacing rules don't break
    # domains (e.g. "leftfieldnyc.com" -> "leftfieldnyc. com").
    url_placeholders = []
    def _mask_url(m):
        url_placeholders.append(m.group(0))
        return _BARE_URL_PLACEHOLDER
    s = _RE_BARE_URL.sub(_mask_url, s)
    # 1c. Protect bare file paths (/path/to/file.ext, ~/dir/file.ext) so
    # period-spacing rules don't corrupt extensions (e.g. "file.md" -> "file. md").
    filepath_placeholders = []
    def _mask_filepath(m):
        filepath_placeholders.append(m.group(0))
        return _FILEPATH_PLACEHOLDER
    s = _RE_BARE_FILEPATH.sub(_mask_filepath, s)
    # 1d. Protect bare filenames with known extensions (e.g. MEMORY.md, drives.json)
    s = _RE_FILENAME_EXT.sub(_mask_filepath, s)
    # 2. Ensure space after sentence-ending punctuation and colon when followed by a letter (streaming often drops it); include fullwidth ？ ！
    s = _RE_SENTENCE_SPACE.sub(r"\1 \2", s)
    # 3. Remove space before common punctuation (e.g. "word ." -> "word.", "flow ?" -> "flow?"); include fullwidth ？ ！ ， ．
    s = _RE_SPACE_BEFORE_PUNCT.sub(r"\1\2", s)
    # 4. Fix space before apostrophe in contractions (e.g. "I 'm" -> "I'm", "you 're" -> "you're", "what ' s" -> "what's"); match ASCII or Unicode apostrophe
    s = _RE_CONTRACTION_APOSTROPHE.sub(r"\1'\2", s)
    # 5. Ensure space after period when used as list punctuation (e.g. "1.Use" -> "1. Use")
    s = _RE_LIST_PERIOD_LETTER.sub(r"\1. \2", s)
    # 5b. Space before list number when glued to previous word (e.g. "Picks1." -> "Picks 1.")
    s = _RE_WORD_LIST_NUM.sub(r"\1 \2", s)
    # 6. Literal stream-merge fixes only (add new pairs when seen; no regex word-splitting)
    for wrong, right in _STREAM_MERGE_FIXES:
        s = s.replace(wrong, right)
    # 6b. Remove space before hyphen when hyphen is followed by word (e.g. "poet - warrior" -> "poet-warrior")
    s = _RE_HYPHEN_WORDS.sub(r"\1-\2", s)
    # 6c. Italic/emphasis: remove spaces between single * and word (e.g. "* see *" -> "*see*")
    s = _RE_ITALIC_SPACES.sub(r"*\1*", s)
    # 7. Collapse runs of 2+ spaces to single space
    s = _RE_COLLAPSE_SPACES.sub(" ", s)
    # 8. Fix list markers: (\n|^) optional spaces then - or * or + then one+ spaces -> \1\2 single space
    s = _RE_LIST_MARKER_SPACE.sub(r"\1\2 ", s)
    # 9b. List marker at line start with no space after (e.g. "-6K" -> "- 6K")
    s = _RE_LIST_MARKER_NO_SPACE.sub(r"\1\2 \3", s)
    # 9c. Space after markdown heading # when followed by letter/digit (e.g. "###Top" -> "### Top")
    s = _RE_HEADING_SPACE.sub(r"\1\2 \3", s)
    # 9d. Space after closing paren before letter or dash (e.g. "Match)-Why" -> "Match) - Why")
    s = _RE_PAREN_LETTER.sub(r") \1", s)
    # 9e. Space between word and following open paren (e.g. "gamut(ideal" -> "gamut (ideal")
    s = _RE_WORD_OPEN_PAREN.sub(r"\1 (", s)
    # 9e2. Space between digit and open paren (e.g. "$849(Amazon)" -> "$849 (Amazon)")
    s = _RE_DIGIT_OPEN_PAREN.sub(r"\1 (", s)
    # 9e3. Space between "the" and following number (e.g. "the13-inch" -> "the 13-inch")
    s = _RE_THE_NUMBER.sub(r"the \1", s)
    # 9e4. Space between "a" and following number (e.g. "a13-inch" -> "a 13-inch")
    s = _RE_A_NUMBER.sub(r"a \1", s)
    # 9f. Space between word and number (e.g. "resolutionwith98%" -> "resolution with 98%")
    s = _RE_WORD_TWO_DIGITS.sub(r"\1 \2", s)
    s = _RE_WORD_PERCENT.sub(r"\1 \2", s)
    # 10. Fix blockquotes: (\n|^) optional spaces then > then one+ spaces -> \1> single space
    s = _RE_BLOCKQUOTE_SPACE.sub(r"\1> ", s)
    # 11. Fix bold markers: remove space immediately after ** and before **
    s = _RE_BOLD_AFTER.sub("**", s)
    s = _RE_BOLD_BEFORE.sub("**", s)
    s = _RE_UNDER_AFTER.sub("__", s)
    s = _RE_UNDER_BEFORE.sub("__", s)
    # 12. Ensure space after list number when followed by bold/italic (e.g. "1.**Renew**" -> "1. **Renew**")
    s = _RE_LIST_NUM_BOLD.sub(r"\1. \2", s)
    # 12b. Ensure space between heading markers and bold (e.g. "###**Title**" -> "### **Title**")
    s = _RE_HEADING_BOLD.sub(r"\1\2 \3", s)
    # 12c. Ensure space between list marker and bold (e.g. "-**Item**" -> "- **Item**")
    s = _RE_LIST_MARKER_BOLD.sub(r"\1\2 \3", s)
    # 13. Space before ** when glued to "a" or "and" so bold renders
    s = _RE_A_AND_BOLD.sub(r"\1 \2", s)
    # 14. Final pass: remove space before sentence-ending punctuation (e.g. "accordingly ." -> "accordingly.", "flow ?" -> "flow?")
    s = _RE_SPACE_BEFORE_PUNCT.sub(r"\1\2", s)
    # Restore bare file paths
    for p in filepath_placeholders:
        s = s.replace(_FILEPATH_PLACEHOLDER, p, 1)
    # Restore bare URLs
    for p in url_placeholders:
        s = s.replace(_BARE_URL_PLACEHOLDER, p, 1)
    # Restore markdown links first (they may contain backticks)
    for p in link_placeholders:
        s = s.replace(_MARKDOWN_LINK_PLACEHOLDER, p, 1)
    # Restore inline code segments
    for p in placeholders:
        s = s.replace(_INLINE_CODE_PLACEHOLDER, p, 1)
    return s


def _log_ts():
    """Return current timestamp for log lines (YYYY-MM-DD HH:MM:SS.mmm)."""
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
