"""
╔══════════════════════════════════════════════════════════════════════════╗
║  MYTHICAL UNIVERSAL SYSTEM — app.py [v7.1]                              ║
║                                                                          ║
║  v7.0 fixes carried forward +                                           ║
║  [FIX-A] Watchdog container-aware (cgroup v1+v2): no more constant      ║
║          flush loop on shared HF Spaces hosts                           ║
║  [FIX-B] Threshold sanity-check: env vars < 30% of total are ignored   ║
║  [FIX-C] Watchdog flush cooldown (30s) — reduces /slots/0 spam         ║
║  [FIX-D] THINKING_TIMEOUT (300s) — thinking mode no longer times out   ║
║  [FIX-E] Dynamic httpx timeout per-request (thinking vs normal)         ║
║  [FIX-F] Token factor uses wd.level, not hardcoded GB values           ║
║  [FIX-G] Health / metrics use container-aware RAM %                     ║
╚══════════════════════════════════════════════════════════════════════════╝
"""
from __future__ import annotations

import asyncio
import base64
import collections
import enum
import gc
import hashlib
import io
import logging
import math
import os
import re
import time
import urllib.parse
import uuid
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any, AsyncIterator

import httpx
import orjson
import psutil
import uvloop
from fastapi import FastAPI, File, Form, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response, StreamingResponse

from model_manager import ModelManager

uvloop.install()

# ─────────────────────────────────────────────────────────────────────────────
# CONFIG
# ─────────────────────────────────────────────────────────────────────────────
LLAMA_HOST      = os.getenv("LLAMA_HOST",    "127.0.0.1")
LLAMA_PORT      = os.getenv("LLAMA_PORT",    "8080")
WHISPER_HOST    = os.getenv("WHISPER_HOST",  "127.0.0.1")
WHISPER_PORT    = os.getenv("WHISPER_PORT",  "8081")
LLAMA_URL       = f"http://{LLAMA_HOST}:{LLAMA_PORT}"
WHISPER_URL     = f"http://{WHISPER_HOST}:{WHISPER_PORT}"

API_KEY          = os.getenv("API_KEY",          "change-this-key")
MAX_CTX_TOKENS   = int(os.getenv("MAX_CTX_TOKENS",  "14000"))
MAX_NEW_TOKENS   = int(os.getenv("MAX_NEW_TOKENS",   "2048"))
MAX_FILE_MB      = int(os.getenv("MAX_FILE_MB",      "50"))
IMAGE_MAX_PX     = int(os.getenv("IMAGE_MAX_PX",     "1120"))
VIDEO_MAX_FRAMES = int(os.getenv("VIDEO_MAX_FRAMES", "8"))
CACHE_TTL        = int(os.getenv("CACHE_TTL",        "60"))
RATE_LIMIT_RPM   = int(os.getenv("RATE_LIMIT_RPM",   "60"))
RATE_BURST       = int(os.getenv("RATE_LIMIT_BURST",  "10"))
RATE_VIP_IPS     = set(os.getenv("RATE_VIP_IPS", "127.0.0.1").split(","))
# RAM thresholds — auto-calculated as % of total RAM if not set explicitly
# RAM thresholds — computed dynamically in watchdog_task() based on actual machine RAM
RAM_WARN_GB   = float(os.getenv("RAM_WARN_GB",   "0"))  # 0 = auto 82% of total RAM
RAM_REJECT_GB = float(os.getenv("RAM_REJECT_GB", "0"))  # 0 = auto 90% of total RAM
RAM_FLUSH_GB  = float(os.getenv("RAM_FLUSH_GB",  "0"))  # 0 = auto 95% of total RAM
REQUEST_TIMEOUT  = float(os.getenv("REQUEST_TIMEOUT","90.0"))
THINKING_TIMEOUT = float(os.getenv("THINKING_TIMEOUT","600.0"))  # thinking mode — CPU needs time
GENERATE_TIMEOUT = float(os.getenv("GENERATE_TIMEOUT","600.0"))  # file generation (big code files)
QUEUE_TIMEOUT    = float(os.getenv("QUEUE_TIMEOUT",  "30.0"))
ENRICH_TIMEOUT   = float(os.getenv("ENRICH_TIMEOUT", "12.0"))
DOWNLOAD_TIMEOUT = float(os.getenv("DOWNLOAD_TIMEOUT","30.0"))

# ── Single source of truth for system prompt ─────────────────────────────────
# startup.sh reads this via: python3 -c "from app import DEFAULT_SYSTEM; print(DEFAULT_SYSTEM)"
# NEVER duplicate this string anywhere else.
DEFAULT_SYSTEM = (
    "You are a universal AI assistant. "
    "Handle text, images, audio, video, PDFs, and web URLs. "
    "When tools are provided emit tool_calls precisely — never execute them. "
    "Think step by step when needed. Respond in the user language. "
    "Be precise and concise."
)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%Y-%m-%dT%H:%M:%S",
)
logger = logging.getLogger("mythical")

# Ready flag: set True only after llama-server is confirmed healthy
_INFERENCE_READY = False


# ─────────────────────────────────────────────────────────────────────────────
# CONNECTION MANAGER
# ─────────────────────────────────────────────────────────────────────────────
class ConnManager:
    def __init__(self, base_url: str, name: str = "srv"):
        self._base = base_url; self._name = name
        self._c: httpx.AsyncClient | None = None
        self._lock = asyncio.Lock()
        self.reconnects = 0; self.healthy = False

    def _build(self) -> httpx.AsyncClient:
        return httpx.AsyncClient(
            base_url=self._base,
            timeout=httpx.Timeout(connect=10.0, read=REQUEST_TIMEOUT,
                                  write=10.0, pool=5.0),
            limits=httpx.Limits(max_connections=10,
                                max_keepalive_connections=5,
                                keepalive_expiry=60.0),
        )

    async def _get(self) -> httpx.AsyncClient:
        if not self._c or self._c.is_closed:
            async with self._lock:
                if not self._c or self._c.is_closed:
                    self._c = self._build()
        return self._c

    async def _reset(self):
        async with self._lock:
            if self._c and not self._c.is_closed:
                await self._c.aclose()
            self._c = None; self.healthy = False
        self.reconnects += 1
        logger.info(f"[conn:{self._name}] reset #{self.reconnects}")

    async def req(self, method: str, path: str, **kw) -> httpx.Response:
        for attempt in range(2):
            c = await self._get()
            try:
                r = await c.request(method, path, **kw)
                self.healthy = True; return r
            except (httpx.RemoteProtocolError, httpx.ConnectError) as e:
                if attempt == 0:
                    logger.warning(f"[conn:{self._name}] {type(e).__name__} — resetting")
                    await self._reset(); await asyncio.sleep(0.5)
                else:
                    raise
            except httpx.TimeoutException:
                raise

    def stream(self, path: str, data: bytes, hdrs: dict):
        mgr = self
        class _Ctx:
            async def __aenter__(s):
                c = await mgr._get()
                s._cm = c.stream("POST", path, content=data, headers=hdrs)
                return await s._cm.__aenter__()
            async def __aexit__(s, *a):
                return await s._cm.__aexit__(*a)
        return _Ctx()

    async def close(self):
        if self._c and not self._c.is_closed:
            await self._c.aclose()


# ─────────────────────────────────────────────────────────────────────────────
# [FIX-06] STREAM GUARD — replaces fragile _sem_released flag
# Context manager that guarantees semaphore released exactly once
# ─────────────────────────────────────────────────────────────────────────────
class StreamGuard:
    """Wraps a streaming generator, ensures semaphore released exactly once."""

    def __init__(self, gen: AsyncIterator[bytes], sem: "ObsSem"):
        self._gen = gen; self._sem = sem; self._released = False

    def release(self):
        if not self._released:
            self._released = True
            self._sem.release()

    async def __aiter__(self) -> AsyncIterator[bytes]:
        try:
            async for chunk in self._gen:
                yield chunk
        finally:
            self.release()


# ─────────────────────────────────────────────────────────────────────────────
# OBSERVABLE SEMAPHORE
# ─────────────────────────────────────────────────────────────────────────────
class ObsSem:
    def __init__(self, n: int):
        self._s = asyncio.Semaphore(n)
        self.active = 0; self.waiting = 0

    async def acquire(self):
        self.waiting += 1
        try:
            await self._s.acquire()
        finally:
            self.waiting -= 1
        self.active += 1

    def release(self):
        if self.active > 0:
            self.active -= 1
            self._s.release()


# ─────────────────────────────────────────────────────────────────────────────
# RATE LIMITER
# ─────────────────────────────────────────────────────────────────────────────
class RateLimiter:
    def __init__(self, rpm: int, burst: int):
        self._rpm = rpm; self._burst = burst
        self._w: dict[str, collections.deque] = {}
        self.blocked = 0

    def check(self, ip: str) -> tuple[bool, float]:
        if ip in RATE_VIP_IPS: return True, 0.0
        now = time.monotonic(); ws = now - 60.0
        if ip not in self._w: self._w[ip] = collections.deque()
        dq = self._w[ip]
        while dq and dq[0] < ws: dq.popleft()
        if len(dq) >= self._rpm + self._burst:
            retry = round(60.0 - (now - dq[0]) + 0.5, 1)
            self.blocked += 1; return False, max(retry, 1.0)
        dq.append(now); return True, 0.0

    def cleanup(self):
        now = time.monotonic()
        stale = [ip for ip, dq in self._w.items()
                 if not dq or dq[-1] < now - 120]
        for ip in stale: del self._w[ip]
        return len(stale)


# ─────────────────────────────────────────────────────────────────────────────
# EXACT CACHE
# ─────────────────────────────────────────────────────────────────────────────
class ExactCache:
    def __init__(self, ttl: int, max_size: int = 200):
        self._s: dict[str, tuple[bytes, float]] = {}
        self._ttl = ttl; self._max = max_size
        self.hits = 0; self.misses = 0

    def _key(self, msgs: list, tools: list, n: int) -> str:
        return hashlib.sha256(
            orjson.dumps({"m": msgs, "t": tools, "n": n},
                         option=orjson.OPT_SORT_KEYS)
        ).hexdigest()[:16]

    def get(self, k: str) -> bytes | None:
        if k not in self._s: self.misses += 1; return None
        data, exp = self._s[k]
        if time.monotonic() > exp:
            del self._s[k]; self.misses += 1; return None
        self.hits += 1; return data

    def set(self, k: str, data: bytes):
        if len(self._s) >= self._max:
            oldest = min(self._s, key=lambda x: self._s[x][1])
            del self._s[oldest]
        self._s[k] = (data, time.monotonic() + self._ttl)

    def cleanup(self) -> int:
        now = time.monotonic()
        exp = [k for k, (_, t) in self._s.items() if now > t]
        for k in exp: del self._s[k]
        return len(exp)

    @property
    def hit_rate(self) -> float:
        total = self.hits + self.misses
        return (self.hits / total * 100) if total else 0.0


# ─────────────────────────────────────────────────────────────────────────────
# SEMANTIC CACHE  [FIX-05: uses CACHE_TTL env var]
# ─────────────────────────────────────────────────────────────────────────────
class SemCache:
    THRESHOLD = float(os.getenv("SEMANTIC_THRESHOLD", "0.92"))

    def __init__(self, max_size: int = 200):
        self._e: list = []
        self._max = max_size
        self._ttl = CACHE_TTL * 2   # uses env var
        self.hits = 0; self.misses = 0

    @staticmethod
    def _vec(text: str) -> dict:
        t = text.lower().strip(); ng: dict[str, int] = {}
        for i in range(max(0, len(t) - 3)):
            g = t[i:i+4]; ng[g] = ng.get(g, 0) + 1
        total = sum(ng.values()) or 1
        return {k: v/total for k, v in ng.items()}

    @staticmethod
    def _cos(a: dict, b: dict) -> float:
        dot = sum(a.get(k, 0) * v for k, v in b.items())
        ma = math.sqrt(sum(v*v for v in a.values()))
        mb = math.sqrt(sum(v*v for v in b.values()))
        return dot / (ma * mb) if ma and mb else 0.0

    def _text(self, msgs: list) -> str:
        parts = []
        for m in msgs:
            c = m.get("content", "")
            if isinstance(c, str): parts.append(c)
            elif isinstance(c, list):
                for b in c:
                    if b.get("type") == "text":
                        parts.append(b.get("text", ""))
        return " ".join(parts)[:2000]

    def lookup(self, msgs: list) -> bytes | None:
        now = time.monotonic(); text = self._text(msgs)
        if len(text) < 8: self.misses += 1; return None
        qv = self._vec(text); best_s = 0.0; best_r = None
        for vec, resp, ts in reversed(self._e):
            if now - ts > self._ttl: continue
            s = self._cos(qv, vec)
            if s > best_s: best_s, best_r = s, resp
        if best_s >= self.THRESHOLD and best_r:
            self.hits += 1
            logger.info(f"[semantic] HIT sim={best_s:.3f}")
            return best_r
        self.misses += 1; return None

    def store(self, msgs: list, resp: bytes):
        text = self._text(msgs)
        if not text: return
        self._e.append((self._vec(text), resp, time.monotonic()))
        if len(self._e) > self._max: self._e.pop(0)

    def cleanup(self) -> int:
        now = time.monotonic(); before = len(self._e)
        self._e = [(v, r, t) for v, r, t in self._e if now - t <= self._ttl]
        return before - len(self._e)

    @property
    def hit_rate(self) -> float:
        total = self.hits + self.misses
        return (self.hits / total * 100) if total else 0.0


# ─────────────────────────────────────────────────────────────────────────────
# DEDUPLICATOR
# ─────────────────────────────────────────────────────────────────────────────
class Dedup:
    def __init__(self):
        self._f: dict[str, asyncio.Future] = {}
        self.count = 0

    async def run_once(self, key: str, fn) -> tuple[Any, bool]:
        if key in self._f:
            self.count += 1
            try:
                r = await asyncio.wait_for(asyncio.shield(self._f[key]), 120.0)
                return r, True
            except (asyncio.TimeoutError, asyncio.CancelledError):
                pass
        fut = asyncio.get_running_loop().create_future()
        self._f[key] = fut
        try:
            r = await fn(); fut.set_result(r); return r, False
        except Exception as e:
            if not fut.done(): fut.set_exception(e)
            raise
        finally:
            self._f.pop(key, None)


# ─────────────────────────────────────────────────────────────────────────────
# JOB QUEUE  [FIX-11: asyncio.Lock on submit]
# ─────────────────────────────────────────────────────────────────────────────
class JobStatus(str, enum.Enum):
    PENDING = "pending"; RUNNING = "running"
    DONE    = "done";    FAILED  = "failed"

class Job:
    __slots__ = ("id","status","created_at","started_at","finished_at",
                 "result","error","payload")
    def __init__(self, jid: str, payload: dict):
        self.id = jid; self.status = JobStatus.PENDING
        self.created_at = time.monotonic()
        self.started_at = self.finished_at = None
        self.result = self.error = None
        self.payload = payload

    def to_dict(self) -> dict:
        return {
            "job_id":    self.id,
            "status":    self.status.value,
            "elapsed_s": round((self.finished_at or time.monotonic()) - self.created_at, 2),
            "error":     self.error,
        }

class JobQ:
    TTL = int(os.getenv("JOB_TTL_SECONDS", "600"))
    MAX = int(os.getenv("MAX_JOBS", "50"))

    def __init__(self):
        self._jobs: dict[str, Job] = {}
        self._q    = asyncio.Queue()
        self._lock = asyncio.Lock()   # [FIX-11]
        self._task: asyncio.Task | None = None
        self.submitted = self.done = self.failed = 0

    def start(self):
        self._task = asyncio.create_task(self._worker(), name="job_worker")

    async def _worker(self):
        logger.info("[jobs] Worker started.")
        while True:
            try:
                jid = await self._q.get()
                job = self._jobs.get(jid)
                if not job: continue
                job.status = JobStatus.RUNNING
                job.started_at = time.monotonic()
                try:
                    resp = await asyncio.wait_for(
                        llama.req("POST", "/v1/chat/completions",
                                  content=orjson.dumps(job.payload),
                                  headers={"Content-Type": "application/json"}),
                        timeout=300.0)
                    if resp.status_code != 200:
                        raise RuntimeError(f"HTTP {resp.status_code}: {resp.text[:120]}")
                    job.result = resp.content
                    job.status = JobStatus.DONE
                    self.done += 1
                except Exception as e:
                    job.error  = str(e)
                    job.status = JobStatus.FAILED
                    self.failed += 1
                finally:
                    job.finished_at = time.monotonic()
                    elapsed = job.finished_at - (job.started_at or job.finished_at)
                    logger.info(f"[jobs] {jid} → {job.status} ({elapsed:.1f}s)")
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"[jobs] worker error: {e}", exc_info=True)

    async def submit(self, payload: dict) -> str:
        async with self._lock:          # [FIX-11] atomic check+insert
            self._evict()
            if len(self._jobs) >= self.MAX:
                raise RuntimeError(f"Job queue full ({self.MAX} max). Retry later.")
            jid = f"job-{uuid.uuid4().hex[:12]}"
            self._jobs[jid] = Job(jid, payload)
        await self._q.put(jid)
        self.submitted += 1
        return jid

    def get(self, jid: str) -> Job | None:
        return self._jobs.get(jid)

    def _evict(self):
        now = time.monotonic()
        done = {JobStatus.DONE, JobStatus.FAILED}
        old = [k for k, j in self._jobs.items()
               if j.status in done and (now - (j.finished_at or 0)) > self.TTL]
        for k in old: del self._jobs[k]

    def stop(self):
        if self._task: self._task.cancel()


# ─────────────────────────────────────────────────────────────────────────────
# OOM WATCHDOG
# ─────────────────────────────────────────────────────────────────────────────

# ── Container-aware memory helpers ───────────────────────────────────────────
def _get_container_mem_limit_gb() -> float:
    """Read container memory limit from cgroup (Docker / HF Spaces).
    Returns 0.0 if running on bare metal or limit is 'unlimited'."""
    checks = [
        ("/sys/fs/cgroup/memory.max",                 "max"),   # cgroup v2
        ("/sys/fs/cgroup/memory/memory.limit_in_bytes", None),  # cgroup v1
    ]
    for path, unlimited_sentinel in checks:
        try:
            raw = Path(path).read_text().strip()
            if raw == unlimited_sentinel:
                continue
            val = int(raw)
            # Sanity: must be between 256 MB and 512 GB to be a real limit
            if 256 * 1024 * 1024 <= val <= 512 * (1024 ** 3):
                return val / (1024 ** 3)
        except Exception:
            pass
    return 0.0


def _get_mem_used_gb() -> float:
    """Current memory usage — reads from cgroup when in a container,
    falls back to psutil system-wide."""
    for path in (
        "/sys/fs/cgroup/memory.current",                # cgroup v2
        "/sys/fs/cgroup/memory/memory.usage_in_bytes",  # cgroup v1
    ):
        try:
            return int(Path(path).read_text()) / (1024 ** 3)
        except Exception:
            pass
    return psutil.virtual_memory().used / (1024 ** 3)


class _WD:
    level = "ok"; ram_gb = 0.0; flushes = 0; rejects = 0
wd = _WD()

_last_flush_attempt: float = 0.0
_flush_backoff: float = 10.0   # start at 10s, doubles on failure, max 120s

async def _flush_kv() -> bool:
    global _last_flush_attempt, _flush_backoff
    now = time.monotonic()
    # Backoff: don't hammer llama-server on repeated failures
    if now - _last_flush_attempt < _flush_backoff:
        return False
    _last_flush_attempt = now

    # Try 1: /slots/0 {"action":"erase"} (llama.cpp >= v0.3.x)
    try:
        r = await llama.req("POST", "/slots/0",
                            content=orjson.dumps({"action": "erase"}),
                            headers={"Content-Type": "application/json"},
                            timeout=httpx.Timeout(5.0))
        if r.status_code in (200, 204):
            logger.info("[wd] KV flushed via /slots/0")
            _flush_backoff = 10.0  # reset on success
            return True
        # 400 = endpoint not supported in this build
        if r.status_code == 400:
            logger.debug("[wd] /slots/0 not supported (400) — KV flush unavailable in this llama.cpp build")
            _flush_backoff = min(_flush_backoff * 2, 120.0)
            return False
    except Exception as e:
        logger.debug(f"[wd] /slots/0 failed: {e}")

    # Try 2: /cache/clear (older llama.cpp)
    try:
        r = await llama.req("POST", "/cache/clear", timeout=httpx.Timeout(5.0))
        if r.status_code in (200, 204):
            logger.info("[wd] KV flushed via /cache/clear")
            _flush_backoff = 10.0
            return True
    except Exception:
        pass

    _flush_backoff = min(_flush_backoff * 2, 120.0)
    return False

async def watchdog_task():
    # ── Step 1: Determine the "effective total" for threshold math ────────────
    _sys_total   = psutil.virtual_memory().total / (1024 ** 3)
    _container   = _get_container_mem_limit_gb()
    _base        = _container if _container else _sys_total

    # ── Step 2: Auto-thresholds as % of effective base ────────────────────────
    _auto_warn   = round(_base * 0.82, 1)
    _auto_reject = round(_base * 0.90, 1)
    _auto_flush  = round(_base * 0.95, 1)

    # ── Step 3: Accept env-var overrides ONLY when they look sensible ─────────
    # Any env var < 30% of base is almost certainly a stale tiny value from an
    # old config (e.g. 11 / 13 / 14 GB on a 124 GB host) → ignore it.
    _floor  = _base * 0.30
    _warn   = RAM_WARN_GB   if RAM_WARN_GB   >= _floor else _auto_warn
    _reject = RAM_REJECT_GB if RAM_REJECT_GB >= _floor else _auto_reject
    _flush  = RAM_FLUSH_GB  if RAM_FLUSH_GB  >= _floor else _auto_flush

    logger.info(
        f"[wd] sys={_sys_total:.1f}G "
        f"{'container=' + f'{_container:.1f}G ' if _container else ''}"
        f"warn={_warn}G reject={_reject}G flush={_flush}G"
    )

    _flush_cooldown = 30.0   # minimum seconds between flush attempts
    _last_flush_t   = 0.0

    while True:
        try:
            await asyncio.sleep(2.0)
            used     = _get_mem_used_gb()
            wd.ram_gb = used

            if used >= _flush:
                prev = wd.level
                if prev != "flush":
                    logger.warning(f"[wd] 🚨 FLUSH {used:.1f}/{_base:.0f}GB")
                wd.level = "flush"; wd.flushes += 1
                # Only attempt GC + KV flush if cooldown has elapsed
                now = time.monotonic()
                if now - _last_flush_t >= _flush_cooldown:
                    _last_flush_t = now
                    gc.collect(2)
                    exact_cache.cleanup(); sem_cache.cleanup()
                    rl.cleanup(); jobs._evict()
                    await _flush_kv()
                    after = _get_mem_used_gb()
                    freed = used - after
                    if abs(freed) > 0.02:  # only log if something actually moved
                        logger.info(f"[wd] post-flush {after:.1f}GB (freed {freed:.1f}GB)")
                # Transition back if flush worked
                if _get_mem_used_gb() < _reject:
                    wd.level = "ok"

            elif used >= _reject:
                if wd.level not in ("reject", "flush"):
                    logger.warning(f"[wd] ⚠ REJECT {used:.1f}/{_base:.0f}GB")
                wd.level = "reject"; wd.rejects += 1; gc.collect(1)

            elif used >= _warn:
                if wd.level == "ok":
                    logger.info(f"[wd] ⚡ WARN {used:.1f}/{_base:.0f}GB")
                wd.level = "warn"

            else:
                if wd.level != "ok":
                    logger.info(f"[wd] ✅ OK {used:.1f}GB")
                wd.level = "ok"

        except asyncio.CancelledError:
            break
        except Exception as e:
            logger.error(f"[wd] {e}", exc_info=True)


# ─────────────────────────────────────────────────────────────────────────────
# GLOBAL INSTANCES
# ─────────────────────────────────────────────────────────────────────────────
llama     = ConnManager(LLAMA_URL,   "llama")
whisper_c = ConnManager(WHISPER_URL, "whisper")
sem       = ObsSem(2)
rl        = RateLimiter(RATE_LIMIT_RPM, RATE_BURST)
exact_cache = ExactCache(CACHE_TTL)
sem_cache   = SemCache()
dedup       = Dedup()
jobs        = JobQ()
mgr         = ModelManager()


# ─────────────────────────────────────────────────────────────────────────────
# MEDIA PROCESSING
# ─────────────────────────────────────────────────────────────────────────────
def _resize_img(data: bytes) -> str | None:
    try:
        from PIL import Image
        img = Image.open(io.BytesIO(data)).convert("RGB")
        if max(img.size) > IMAGE_MAX_PX:
            img.thumbnail((IMAGE_MAX_PX, IMAGE_MAX_PX), Image.LANCZOS)
        buf = io.BytesIO()
        img.save(buf, "JPEG", quality=85, optimize=True)
        return base64.b64encode(buf.getvalue()).decode()
    except Exception as e:
        logger.warning(f"[img] {e}"); return None

async def resize_img(data: bytes) -> str | None:
    loop = asyncio.get_running_loop()
    return await loop.run_in_executor(None, _resize_img, data)

# [FIX-02] async chunked read to avoid blocking event loop
async def read_upload_chunked(file: UploadFile, max_mb: int = MAX_FILE_MB) -> bytes | None:
    """Read upload in chunks on thread executor — never blocks asyncio loop."""
    max_bytes = max_mb * 1024 * 1024
    loop = asyncio.get_running_loop()
    chunks = []
    total = 0
    while True:
        chunk = await loop.run_in_executor(None, file.file.read, 65536)
        if not chunk:
            break
        total += len(chunk)
        if total > max_bytes:
            logger.warning(f"[upload] File exceeds {max_mb}MB limit")
            return None
        chunks.append(chunk)
    return b"".join(chunks)

async def convert_audio(data: bytes, fmt: str = "mp3") -> str | None:
    if len(data) > MAX_FILE_MB * 1024 * 1024: return None
    Path("/tmp/media").mkdir(parents=True, exist_ok=True)
    inp = f"/tmp/media/{uuid.uuid4().hex}.{fmt}"
    out = f"/tmp/media/{uuid.uuid4().hex}.wav"
    try:
        Path(inp).write_bytes(data)
        p = await asyncio.create_subprocess_exec(
            "ffmpeg", "-y", "-i", inp,
            "-ar", "16000", "-ac", "1", "-c:a", "pcm_s16le", "-t", "60",
            out,
            stdout=asyncio.subprocess.DEVNULL,
            stderr=asyncio.subprocess.DEVNULL)
        await asyncio.wait_for(p.communicate(), timeout=45.0)
        if p.returncode != 0: return None
        loop = asyncio.get_running_loop()
        wav = await loop.run_in_executor(None, Path(out).read_bytes)
        return base64.b64encode(wav).decode()
    except Exception as e:
        logger.warning(f"[audio] {e}"); return None
    finally:
        for f in (inp, out):
            try: os.unlink(f)
            except: pass

async def extract_frames(data: bytes) -> list[str]:
    if len(data) > MAX_FILE_MB * 1024 * 1024: return []
    import shutil
    td = Path(f"/tmp/media/{uuid.uuid4().hex}")
    td.mkdir(parents=True)
    try:
        (td / "v.mp4").write_bytes(data)
        fps = VIDEO_MAX_FRAMES / 60.0
        p = await asyncio.create_subprocess_exec(
            "ffmpeg", "-y", "-i", str(td / "v.mp4"),
            "-vf", f"fps={fps:.4f},scale=560:-1", "-q:v", "4",
            str(td / "f%04d.jpg"),
            stdout=asyncio.subprocess.DEVNULL,
            stderr=asyncio.subprocess.DEVNULL)
        await asyncio.wait_for(p.communicate(), timeout=60.0)
        frames: list[str] = []
        loop = asyncio.get_running_loop()
        for f in sorted(td.glob("f*.jpg"))[:VIDEO_MAX_FRAMES]:
            raw = await loop.run_in_executor(None, f.read_bytes)
            b64 = await resize_img(raw)
            if b64: frames.append(f"data:image/jpeg;base64,{b64}")
        return frames
    except Exception as e:
        logger.warning(f"[video] {e}"); return []
    finally:
        shutil.rmtree(td, ignore_errors=True)

# [FIX-08] graceful handling for encrypted/corrupt PDFs
def _pdf_text(data: bytes, max_chars: int = 20000) -> str | None:
    try:
        from pypdf import PdfReader
        try:
            reader = PdfReader(io.BytesIO(data))
        except Exception:
            return "[PDF Error: file is encrypted, corrupted, or not a valid PDF]"
        if reader.is_encrypted:
            # try empty password
            try:
                reader.decrypt("")
            except Exception:
                return "[PDF Error: file is password-protected. Please provide an unlocked PDF]"
        text = ""
        for page in reader.pages:
            try:
                text += (page.extract_text() or "")
            except Exception:
                continue
            if len(text) > max_chars:
                break
        return text[:max_chars] if text.strip() else "[PDF: no extractable text found (may be image-based)]"
    except Exception as e:
        return f"[PDF processing error: {e}]"

def _yt_id(url: str) -> str | None:
    p = urllib.parse.urlparse(url)
    if p.hostname == "youtu.be": return p.path[1:]
    if p.hostname in ("www.youtube.com", "youtube.com"):
        if p.path == "/watch":
            return urllib.parse.parse_qs(p.query).get("v", [None])[0]
        if p.path.startswith("/shorts/"):
            return p.path.split("/")[2]
    return None

async def _yt_transcript(vid: str) -> str | None:
    try:
        from youtube_transcript_api import YouTubeTranscriptApi
        loop = asyncio.get_running_loop()
        t = await asyncio.wait_for(
            loop.run_in_executor(None, lambda: YouTubeTranscriptApi.get_transcript(
                vid, languages=["ar", "ar-SA", "en", "en-US"])),
            timeout=10.0)
        return " ".join(x["text"] for x in t)[:15000]
    except asyncio.TimeoutError:
        return "[YouTube transcript timed out]"
    except Exception:
        return None

async def _fetch_url(url: str) -> str | None:
    try:
        import aiohttp
        from bs4 import BeautifulSoup
        async with aiohttp.ClientSession() as s:
            async with s.get(url,
                             headers={"User-Agent": "Mozilla/5.0 (compatible; MythicalBot/7.0)"},
                             timeout=aiohttp.ClientTimeout(total=8),
                             allow_redirects=True) as r:
                if r.status == 200:
                    soup = BeautifulSoup(await r.text(errors="replace"), "html.parser")
                    for t in soup(["script", "style", "nav", "footer", "aside"]): t.decompose()
                    return soup.get_text(" ", strip=True)[:15000]
    except Exception:
        return None

async def web_search(query: str, max_results: int = 5) -> list[dict]:
    """DuckDuckGo HTML search — free, no API key."""
    try:
        import aiohttp
        from bs4 import BeautifulSoup
        url = f"https://html.duckduckgo.com/html/?q={urllib.parse.quote(query)}"
        async with aiohttp.ClientSession() as s:
            async with s.get(url,
                             headers={"User-Agent": "Mozilla/5.0"},
                             timeout=aiohttp.ClientTimeout(total=10)) as resp:
                if resp.status != 200: return []
                soup = BeautifulSoup(await resp.text(), "html.parser")
                results = []
                for r in soup.select(".result__body")[:max_results]:
                    title   = r.select_one(".result__title")
                    snippet = r.select_one(".result__snippet")
                    link    = r.select_one(".result__url")
                    results.append({
                        "title":   title.get_text(strip=True) if title else "",
                        "snippet": snippet.get_text(strip=True) if snippet else "",
                        "url":     link.get_text(strip=True) if link else "",
                    })
                return results
    except Exception as e:
        logger.warning(f"[search] {e}"); return []


# [FIX-04] enrich_error creates user-visible message for empty search
async def enrich(messages: list[dict]) -> list[dict]:
    """Process all media/URL content with hard timeout."""
    try:
        return await asyncio.wait_for(_enrich_inner(messages), timeout=ENRICH_TIMEOUT)
    except asyncio.TimeoutError:
        logger.warning(f"[enrich] Timeout after {ENRICH_TIMEOUT}s")
        return messages

async def _enrich_inner(messages: list[dict]) -> list[dict]:
    result = []
    for msg in messages:
        content = msg.get("content")

        if isinstance(content, str):
            urls = re.findall(r"https?://\S+", content)
            if urls:
                url = urls[0]
                yt  = _yt_id(url)
                if yt:
                    t = await _yt_transcript(yt)
                    if t: content += f"\n\n[YouTube Transcript]:\n{t}"
                elif any(url.lower().endswith(e) for e in
                         [".jpg",".jpeg",".png",".webp",".gif"]):
                    try:
                        import aiohttp
                        async with aiohttp.ClientSession() as s:
                            async with s.get(url, timeout=aiohttp.ClientTimeout(total=10)) as r:
                                if r.status == 200:
                                    b64 = await resize_img(await r.read())
                                    if b64:
                                        result.append({**msg, "content": [
                                            {"type": "text", "text": content},
                                            {"type": "image_url",
                                             "image_url": {"url": f"data:image/jpeg;base64,{b64}"}}
                                        ]})
                                        continue
                    except Exception:
                        pass
                else:
                    web = await _fetch_url(url)
                    if web:
                        content += f"\n\n[Web Content from {url}]:\n{web}"
            result.append({**msg, "content": content})
            continue

        if isinstance(content, list):
            new = []
            for block in content:
                bt  = block.get("type", "")
                if bt == "image_url":
                    uv = block.get("image_url", {}).get("url", "")
                    if uv.startswith("data:"):
                        try:
                            _, b = uv.split(",", 1)
                            raw = base64.b64decode(b)
                            if len(raw) > MAX_FILE_MB * 1024 * 1024:
                                new.append({"type":"text","text":f"[Image rejected: exceeds {MAX_FILE_MB}MB limit]"})
                                continue
                            b64 = await resize_img(raw)
                            new.append({"type":"image_url","image_url":
                                        {"url":f"data:image/jpeg;base64,{b64 or b}"}})
                        except Exception as e:
                            new.append({"type":"text","text":f"[Image error: {e}]"})
                    else:
                        new.append({"type":"text","text":"[External image URLs not supported — please convert to base64 data URI]"})

                elif bt == "input_audio":
                    info     = block.get("input_audio", {})
                    fmt      = info.get("format", "wav")
                    data_b64 = info.get("data", "")
                    if data_b64 and fmt != "wav":
                        wav = await convert_audio(base64.b64decode(data_b64), fmt)
                        new.append({"type":"input_audio",
                                    "input_audio":{"data": wav or data_b64, "format":"wav"}})
                    else:
                        new.append(block)

                elif bt == "video_url":
                    uv = block.get("video_url", {}).get("url", "")
                    if uv.startswith("data:"):
                        try:
                            _, b = uv.split(",", 1)
                            frames = await extract_frames(base64.b64decode(b))
                            if frames:
                                new.append({"type":"text","text":f"[Video: {len(frames)} frames extracted]"})
                                for i, f in enumerate(frames):
                                    new.append({"type":"text","text":f"Frame {i+1}/{len(frames)}:"})
                                    new.append({"type":"image_url","image_url":{"url":f}})
                            else:
                                new.append({"type":"text","text":"[Video: could not extract frames — check format/size]"})
                        except Exception as e:
                            new.append({"type":"text","text":f"[Video error: {e}]"})
                    else:
                        new.append(block)

                elif "pdf" in block.get("image_url", {}).get("url", ""):
                    try:
                        _, b   = block["image_url"]["url"].split(",", 1)
                        loop   = asyncio.get_running_loop()
                        text   = await loop.run_in_executor(None, _pdf_text, base64.b64decode(b))
                        new.append({"type":"text","text":f"[PDF Content]:\n{text}"})
                    except Exception as e:
                        new.append({"type":"text","text":f"[PDF error: {e}]"})

                else:
                    new.append(block)
            result.append({**msg, "content": new})
            continue

        result.append(msg)
    return result


def count_tokens(messages: list) -> int:
    total = 0
    for m in messages:
        c = m.get("content", "")
        if isinstance(c, str): total += len(c)
        elif isinstance(c, list):
            for b in c:
                if b.get("type") == "text":     total += len(b.get("text",""))
                elif b.get("type") == "image_url":   total += 1500
                elif b.get("type") == "input_audio": total += 3000
        total += 16
    return int(total / 3.5)

def surgeon(messages: list, max_tok: int) -> tuple[list, int, int]:
    orig = count_tokens(messages)
    if orig <= max_tok: return messages, orig, orig
    sys_m = [m for m in messages if m.get("role") == "system"]
    conv  = [m for m in messages if m.get("role") != "system"]
    while len(conv) > 2:
        conv.pop(len(conv) // 2)
        if count_tokens(sys_m + conv) <= max_tok: break
    if count_tokens(sys_m + conv) > max_tok and conv:
        conv = conv[-1:]
    res = sys_m + conv
    return res, orig, count_tokens(res)


# ─────────────────────────────────────────────────────────────────────────────
# LIFESPAN
# ─────────────────────────────────────────────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
    global _INFERENCE_READY
    logger.info("══ MYTHICAL UNIVERSAL SYSTEM v7.0 starting ══")
    wd_t = asyncio.create_task(watchdog_task(), name="oom_watchdog")
    jobs.start()
    mgr.start_updater()

    # [FIX-01] Poll until llama-server is actually healthy before accepting inference
    async def _wait_for_inference():
        global _INFERENCE_READY
        for _ in range(120):   # max 6 minutes
            try:
                r = await llama.req("GET", "/health", timeout=httpx.Timeout(3.0))
                if r.status_code == 200:
                    _INFERENCE_READY = True
                    logger.info("✅ Inference engine ready — accepting all requests")
                    return
            except Exception:
                pass
            await asyncio.sleep(3)
        logger.warning("⚠ Inference engine not ready after 6min — requests will fail gracefully")
        _INFERENCE_READY = True   # allow through, will get 502 with proper error

    asyncio.create_task(_wait_for_inference(), name="inference_ready_probe")
    logger.info("Services: Watchdog ✓ | Jobs ✓ | ModelUpdater ✓ | ReadyProbe ✓")
    yield

    logger.info("Shutting down...")
    wd_t.cancel(); jobs.stop(); mgr.stop()
    try: await asyncio.wait_for(asyncio.shield(wd_t), 3.0)
    except (asyncio.CancelledError, asyncio.TimeoutError): pass
    await llama.close(); await whisper_c.close()
    logger.info("Clean shutdown complete.")


app = FastAPI(title="Mythical Universal System", version="7.0.0",
              lifespan=lifespan, default_response_class=Response)
app.add_middleware(CORSMiddleware, allow_origins=["*"],
                   allow_methods=["*"], allow_headers=["*"],
                   allow_credentials=True)


def J(data: Any, code: int = 200) -> Response:
    return Response(orjson.dumps(data), code, media_type="application/json")

def E(msg: str, code: int = 500, t: str = "error", rid: str | None = None) -> Response:
    return J({"error": {"message": msg, "type": t, "code": code},
               "id": rid or f"err-{uuid.uuid4().hex[:8]}"}, code)

def auth(r: Request) -> bool:
    return r.headers.get("Authorization","").replace("Bearer ","").strip() == API_KEY

def client_ip(r: Request) -> str:
    return r.headers.get("X-Forwarded-For",
                         r.client.host or "0.0.0.0").split(",")[0].strip()


# ─────────────────────────────────────────────────────────────────────────────
# HEALTH & MONITORING
# ─────────────────────────────────────────────────────────────────────────────
@app.get("/health", response_model=None)
@app.get("/", response_model=None)
async def health() -> Response:
    used        = _get_mem_used_gb()
    _climit     = _get_container_mem_limit_gb()
    _sys_total  = psutil.virtual_memory().total / (1024 ** 3)
    _total      = _climit if _climit else _sys_total
    ram_pct     = round(used / _total * 100, 1) if _total > 0 else 0.0
    ll_ok = False
    try:
        r = await llama.req("GET", "/health", timeout=httpx.Timeout(3.0))
        ll_ok = r.status_code == 200
    except Exception:
        pass
    ok = ll_ok and wd.level in ("ok","warn","flush")  # flush = high RAM but still working
    return J({
        "status":           "healthy" if ok else "degraded",
        "inference_ready":  _INFERENCE_READY,
        "llama_server":     "ok" if ll_ok else "error",
        "ram_gb":           round(used, 2),
        "ram_pct":          ram_pct,
        "watchdog":         wd.level,
        "queue_active":     sem.active,
        "queue_waiting":    sem.waiting,
        "cache_hit_rate":   f"{exact_cache.hit_rate:.1f}%",
        "sem_hit_rate":     f"{sem_cache.hit_rate:.1f}%",
        "conn_reconnects":  llama.reconnects,
        "model":            mgr.status()["active"],
        "capabilities": {
            "text":True, "images":True, "audio":True, "video":True,
            "pdf":True, "youtube":True, "urls":True, "web_search":True,
            "tool_calls":True, "streaming":True, "thinking_mode":True,
            "async_jobs":True, "file_upload":True, "file_generation":True,
        },
    }, 200 if ok else 503)

# [FIX-01] /ready endpoint — used by n8n/orchestrators to wait for inference
@app.get("/ready", response_model=None)
async def ready() -> Response:
    """Returns 200 only when inference engine is confirmed healthy."""
    if not _INFERENCE_READY:
        return J({"ready": False, "message": "Inference engine still loading..."}, 503)
    ll_ok = False
    try:
        r = await llama.req("GET", "/health", timeout=httpx.Timeout(2.0))
        ll_ok = r.status_code == 200
    except Exception:
        pass
    if ll_ok:
        return J({"ready": True})
    return J({"ready": False, "message": "Inference engine not responding"}, 503)

@app.get("/metrics", response_model=None)
async def metrics() -> Response:
    used   = _get_mem_used_gb()
    _cl    = _get_container_mem_limit_gb()
    _tot   = _cl if _cl else psutil.virtual_memory().total / (1024 ** 3)
    pct    = round(used / _tot * 100, 1) if _tot > 0 else 0.0
    lines = [
        f"ram_gb {used:.3f}", f"ram_pct {pct}",
        f"watchdog_flushes {wd.flushes}", f"watchdog_rejects {wd.rejects}",
        f"queue_active {sem.active}", f"queue_waiting {sem.waiting}",
        f"exact_cache_hits {exact_cache.hits}", f"exact_cache_misses {exact_cache.misses}",
        f"sem_cache_hits {sem_cache.hits}", f"rate_blocked {rl.blocked}",
        f"jobs_submitted {jobs.submitted}", f"jobs_done {jobs.done}",
        f"jobs_failed {jobs.failed}", f"conn_reconnects {llama.reconnects}",
        f"dedup_count {dedup.count}", f"inference_ready {int(_INFERENCE_READY)}",
    ]
    return Response("\n".join(lines), media_type="text/plain")

@app.get("/v1/models", response_model=None)
async def model_list() -> Response:
    return J({"object":"list","data":[
        {"id":"mythical","object":"model","owned_by":"mythical-system",
         "capabilities":["text","vision","audio","function_calling"]}]})


# ─────────────────────────────────────────────────────────────────────────────
# CHAT COMPLETIONS
# ─────────────────────────────────────────────────────────────────────────────
@app.post("/v1/chat/completions", response_model=None)
async def chat(request: Request) -> Response | StreamingResponse:
    if not auth(request): return E("Unauthorized",401,"auth_error")

    rid = f"chatcmpl-{uuid.uuid4().hex[:12]}"; t0 = time.monotonic()
    ip  = client_ip(request)

    # Rate limit
    allowed, retry = rl.check(ip)
    if not allowed:
        return Response(
            orjson.dumps({"error":{"message":f"Rate limit exceeded. Retry in {retry}s."}}),
            429, media_type="application/json",
            headers={"Retry-After": str(int(retry))})

    # OOM gate
    if wd.level == "reject":
        return E(f"Memory pressure ({wd.ram_gb:.1f}GB). Retry in 30s.", 503, "server_overloaded")

    # [FIX-01] Inference ready gate
    if not _INFERENCE_READY:
        return E("Inference engine is still loading. Check /ready endpoint.", 503, "not_ready")

    try:
        body = orjson.loads(await request.body())
    except Exception as e:
        return E(f"Invalid JSON: {e}", 400, "invalid_request")

    messages: list = body.get("messages", [])
    if not messages:
        return E("'messages' field is required", 400, "invalid_request")

    # Vision pre-check: reject image content early if model has no mmproj
    # (avoids llama-server returning a confusing 500)
    _has_image = any(
        isinstance(m.get("content"), list) and
        any(p.get("type") in ("image_url","image") for p in m["content"])
        for m in messages
    )
    if _has_image and not mgr.cfg.mmproj_path:
        return E(
            "Vision not available: this model was loaded without an mmproj file. "
            "Use a vision-capable model or send text-only messages.",
            400, "unsupported_media_type"
        )

    if not any(m.get("role") == "system" for m in messages):
        messages = [{"role":"system","content":DEFAULT_SYSTEM}] + messages

    messages = await enrich(messages)
    messages, orig, final = surgeon(messages, MAX_CTX_TOKENS)
    if orig > final:
        logger.info(f"[{rid}] surgeon {orig}→{final} tokens")

    req_max  = int(body.get("max_tokens") or MAX_NEW_TOKENS)
    # Token budget: reduce under memory pressure (based on watchdog level, not hardcoded GBs)
    factor   = 0.5 if wd.level in ("flush", "reject") else 0.75 if wd.level == "warn" else 1.0
    max_tok  = max(int(min(req_max, MAX_NEW_TOKENS) * factor), 128)
    is_stream= bool(body.get("stream", False))
    temp     = float(body.get("temperature", 0.6))
    use_cache= (temp == 0.0 or bool(body.get("use_cache"))) and not is_stream

    logger.info(f"[{rid}] {ip} | {final}tok | "
                f"tools={len(body.get('tools',[]))} | stream={is_stream} | "
                f"thinking={body.get('thinking',False)} | q={sem.active+sem.waiting}")

    # Cache check
    cache_key: str | None = None
    if use_cache:
        cache_key = exact_cache._key(messages, body.get("tools",[]), max_tok)
        if (hit := exact_cache.get(cache_key)):
            logger.info(f"[{rid}] EXACT HIT ({exact_cache.hit_rate:.0f}%)")
            return Response(hit, 200, media_type="application/json",
                            headers={"X-Request-ID":rid,"X-Cache":"EXACT-HIT"})
        if (sh := sem_cache.lookup(messages)):
            return Response(sh, 200, media_type="application/json",
                            headers={"X-Request-ID":rid,"X-Cache":"SEMANTIC-HIT"})

    # Build payload
    # Thinking mode OFF by default — Qwen3 thinks for EVERY request otherwise
    # User must explicitly pass "thinking": true to enable reasoning
    thinking_requested = body.get("thinking", False)

    payload: dict = {
        "messages":     messages,
        "cache_prompt": True,
        "id_slot":      0,
        "max_tokens":   max_tok,
        "temperature":  temp,
        "top_p":        float(body.get("top_p", 0.95)),
        "top_k":        int(body.get("top_k", 20)),
        "min_p":        float(body.get("min_p", 0.0)),
        "stream":       is_stream,
        # Always set thinking explicitly to avoid Qwen3 auto-enabling it
        "chat_template_kwargs": {"enable_thinking": bool(thinking_requested)},
    }
    if thinking_requested:
        _think_budget = int(body.get("thinking_budget", 512))  # 512 default — 4096 is too slow on CPU
        payload["reasoning_budget"] = _think_budget
        # CRITICAL: max_tokens must cover thinking budget + actual answer
        # Without this, the <think> block eats all tokens and content is empty
        max_tok = max(max_tok, _think_budget + 256)
        payload["max_tokens"] = max_tok
    if body.get("tools"):
        payload["tools"]               = body["tools"]
        payload["tool_choice"]         = body.get("tool_choice", "auto")
        payload["parallel_tool_calls"] = body.get("parallel_tool_calls", True)
    for p in ("stop","presence_penalty","frequency_penalty","seed"):
        if p in body: payload[p] = body[p]

    # Semaphore
    try:
        await asyncio.wait_for(sem.acquire(), timeout=QUEUE_TIMEOUT)
    except asyncio.TimeoutError:
        return E(f"Queued {QUEUE_TIMEOUT:.0f}s — server busy. Retry shortly.",
                 503, "server_overloaded", rid)

    raw = orjson.dumps(payload)
    ct  = {"Content-Type": "application/json"}

    try:
        # [FIX-06] Streaming uses StreamGuard — semaphore released exactly once
        if is_stream:
            async def _raw_stream() -> AsyncIterator[bytes]:
                async with llama.stream("/v1/chat/completions", raw, ct) as resp:
                    if resp.status_code != 200:
                        b = await resp.aread()
                        yield b"data: " + orjson.dumps({
                            "error":{"message":"Upstream error","code":resp.status_code}
                        }) + b"\n\n"
                        return
                    async for chunk in resp.aiter_bytes(256):
                        if chunk: yield chunk

            guard = StreamGuard(_raw_stream(), sem)

            async def _guarded_stream() -> AsyncIterator[bytes]:
                try:
                    async for chunk in guard:
                        yield chunk
                except httpx.TimeoutException:
                    yield b"data: [DONE]\n\n"
                except Exception as e:
                    logger.error(f"[{rid}] stream error: {e}")
                    yield b"data: [DONE]\n\n"
                finally:
                    guard.release()   # idempotent via StreamGuard
                    logger.info(f"[{rid}] stream done {time.monotonic()-t0:.2f}s")

            return StreamingResponse(_guarded_stream(),
                                     media_type="text/event-stream",
                                     headers={"X-Request-ID":rid,
                                              "Cache-Control":"no-cache"})

        # Blocking
        try:
            # Thinking mode needs a much longer budget — use THINKING_TIMEOUT
            _req_timeout = THINKING_TIMEOUT if thinking_requested else REQUEST_TIMEOUT
            _http_timeout = httpx.Timeout(
                connect=10.0, read=_req_timeout + 30.0, write=10.0, pool=5.0
            )

            async def _do():
                return await asyncio.wait_for(
                    llama.req("POST", "/v1/chat/completions",
                              content=raw, headers=ct, timeout=_http_timeout),
                    timeout=_req_timeout)
            resp, deduped = await dedup.run_once(cache_key or rid, _do)
            if deduped: logger.info(f"[{rid}] DEDUP HIT")
        except asyncio.TimeoutError:
            await llama._reset()  # close broken connection before next request uses it
            return E(f"Inference timeout {_req_timeout:.0f}s", 504, "timeout", rid)
        except httpx.RequestError as e:
            return E(f"Upstream error: {e}",502,"server_error",rid)

        elapsed = time.monotonic() - t0
        logger.info(f"[{rid}] done {elapsed:.2f}s HTTP {resp.status_code}")

        if use_cache and cache_key and resp.status_code == 200:
            exact_cache.set(cache_key, resp.content)
            sem_cache.store(messages, resp.content)

        return Response(resp.content, resp.status_code,
                        media_type="application/json",
                        headers={"X-Request-ID":rid,
                                 "X-Time":f"{elapsed:.3f}",
                                 "X-Cache":"MISS"})

    except Exception as e:
        logger.error(f"[{rid}] unhandled: {e}", exc_info=True)
        return E(f"Internal: {type(e).__name__}", 500, "server_error", rid)
    finally:
        if not is_stream:
            sem.release()   # safe: ObsSem.release guards active > 0


# ─────────────────────────────────────────────────────────────────────────────
# FILE UPLOAD  [FIX-02: chunked read]
# ─────────────────────────────────────────────────────────────────────────────
@app.post("/v1/files", response_model=None)
async def file_upload(
    request:  Request,
    file:     UploadFile = File(...),
    prompt:   str = Form(default="Analyze this file and describe its contents."),
    thinking: str = Form(default="false"),
) -> Response:
    """
    Direct file upload — no base64 encoding needed.
    Supports: images, PDFs, audio, video, text/code files.

    curl -H "Authorization: Bearer KEY" \\
         -F "file=@document.pdf" \\
         -F "prompt=Summarize this document" \\
         https://YOUR-SPACE.hf.space/v1/files
    """
    if not auth(request): return E("Unauthorized",401)
    if not _INFERENCE_READY: return E("Inference engine still loading.",503,"not_ready")

    try:
        # [FIX-02] chunked read — never blocks asyncio loop
        data = await read_upload_chunked(file, MAX_FILE_MB)
        if data is None:
            return E(f"File too large. Maximum size: {MAX_FILE_MB}MB", 413, "file_too_large")

        ct_in = file.content_type or ""
        fname = file.filename or "upload"
        b64   = base64.b64encode(data).decode()

        if ct_in.startswith("image/") or any(fname.lower().endswith(e)
             for e in [".jpg",".jpeg",".png",".webp",".gif",".bmp"]):
            content = [
                {"type":"text","text":prompt},
                {"type":"image_url","image_url":{"url":f"data:{ct_in};base64,{b64}"}},
            ]
        elif ct_in == "application/pdf" or fname.lower().endswith(".pdf"):
            content = [
                {"type":"image_url","image_url":{"url":f"data:application/pdf;base64,{b64}"}},
                {"type":"text","text":prompt},
            ]
        elif ct_in.startswith("audio/") or any(fname.lower().endswith(e)
             for e in [".mp3",".wav",".ogg",".m4a",".webm",".flac"]):
            ext = fname.rsplit(".",1)[-1].lower() if "." in fname else "mp3"
            content = [
                {"type":"input_audio","input_audio":{"data":b64,"format":ext}},
                {"type":"text","text":prompt},
            ]
        elif ct_in.startswith("video/") or any(fname.lower().endswith(e)
             for e in [".mp4",".mov",".avi",".mkv",".webm"]):
            content = [
                {"type":"video_url","video_url":{"url":f"data:{ct_in};base64,{b64}"}},
                {"type":"text","text":prompt},
            ]
        else:
            # Text / code file
            try:
                text_content = data.decode("utf-8","replace")[:15000]
            except Exception:
                text_content = "[Binary file — cannot display as text]"
            content = [{"type":"text",
                        "text":f"File: {fname} ({len(data)} bytes)\n\n{text_content}\n\n{prompt}"}]

        messages_raw = [
            {"role":"system","content":DEFAULT_SYSTEM},
            {"role":"user","content":content},
        ]
        messages_enriched = await enrich(messages_raw)
        messages_final, _, _ = surgeon(messages_enriched, MAX_CTX_TOKENS)

        pl = {
            "messages":     messages_final,
            "max_tokens":   MAX_NEW_TOKENS,
            "temperature":  0.6,
            "cache_prompt": False,
            "stream":       False,
            # Always set explicitly — without this Qwen3 may silently enter thinking mode
            "chat_template_kwargs": {"enable_thinking": thinking.lower() == "true"},
        }
        if thinking.lower() == "true":
            _think_budget = 512  # same conservative default as chat endpoint
            pl["reasoning_budget"] = _think_budget
            pl["max_tokens"] = max(MAX_NEW_TOKENS, _think_budget + 256)

        _file_timeout = THINKING_TIMEOUT if thinking.lower() == "true" else REQUEST_TIMEOUT
        resp = await asyncio.wait_for(
            llama.req("POST","/v1/chat/completions",
                      content=orjson.dumps(pl),
                      headers={"Content-Type":"application/json"},
                      timeout=httpx.Timeout(connect=10.0, read=_file_timeout + 30.0,
                                            write=10.0, pool=5.0)),
            timeout=_file_timeout)

        return Response(resp.content, resp.status_code,
                        media_type="application/json",
                        headers={"X-File":fname,"X-Bytes":str(len(data))})

    except asyncio.TimeoutError:
        await llama._reset()
        return E("File processing timed out", 504, "timeout")
    except Exception as e:
        logger.error(f"[files] {e}", exc_info=True)
        return E(f"File upload error: {e}", 500)


# ─────────────────────────────────────────────────────────────────────────────
# FILE GENERATION  [FIX-07: better fence removal]
# ─────────────────────────────────────────────────────────────────────────────
VALID_FORMATS = {"py","md","html","json","txt","sh","csv","js","ts",
                 "yaml","yml","sql","r","cpp","c","java","go","rs","jsx","vue"}

@app.post("/v1/generate", response_model=None)
async def generate_file(request: Request) -> Response:
    """
    Ask the AI to generate a file and download it directly.

    {
      "prompt": "Write a Python script to download all images from a webpage",
      "format": "py",
      "filename": "image_downloader.py",
      "thinking": false
    }
    """
    if not auth(request): return E("Unauthorized",401)
    if not _INFERENCE_READY: return E("Inference engine still loading.",503,"not_ready")

    try: body = orjson.loads(await request.body())
    except Exception as e: return E(f"Invalid JSON: {e}",400)

    prompt   = body.get("prompt","").strip()
    fmt      = body.get("format","txt").lower().strip(".")
    filename = body.get("filename") or f"output_{uuid.uuid4().hex[:6]}.{fmt}"
    thinking = bool(body.get("thinking",False))
    # Cap max_tokens: 1024 default (enough for most scripts; user can override up to 2048)
    gen_max_tokens = min(int(body.get("max_tokens", 1024)), MAX_NEW_TOKENS)

    if not prompt: return E("'prompt' required",400)
    if fmt not in VALID_FORMATS:
        return E(f"Invalid format '{fmt}'. Supported: {sorted(VALID_FORMATS)}",400)

    gen_system = (
        f"You are a file generator. Output ONLY the raw content of a .{fmt} file. "
        f"No explanations, no markdown code blocks, no preamble. "
        f"Output the file content directly, ready to save."
    )
    pl = {
        "messages": [
            {"role":"system","content":gen_system},
            {"role":"user","content":prompt},
        ],
        "max_tokens":  gen_max_tokens,
        "temperature": 0.3,
        "cache_prompt":False,
        "stream":      False,
        # CRITICAL: always set thinking explicitly — without this Qwen3 may enter thinking
        # mode silently, consuming all tokens in <think> blocks and producing empty files
        "chat_template_kwargs": {"enable_thinking": bool(thinking)},
    }
    if thinking:
        pl["reasoning_budget"] = 4096

    try:
        resp = await asyncio.wait_for(
            llama.req("POST","/v1/chat/completions",
                      content=orjson.dumps(pl),
                      headers={"Content-Type":"application/json"},
                      timeout=httpx.Timeout(connect=10.0, read=GENERATE_TIMEOUT + 30.0,
                                            write=10.0, pool=5.0)),
            timeout=GENERATE_TIMEOUT)

        if resp.status_code != 200:
            return E(f"Generation failed: HTTP {resp.status_code}",500)

        data     = orjson.loads(resp.content)
        raw_text = data["choices"][0]["message"]["content"].strip()

        # [FIX-07] Remove only first/last fence lines, preserve internal ```
        lines = raw_text.splitlines()
        if lines and lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        file_content = "\n".join(lines).strip()

        tokens_used = data.get("usage",{}).get("completion_tokens","?")

        # Guard: if content is empty the model likely entered thinking mode silently
        if not file_content:
            logger.warning(f"[generate] model returned empty content (tokens={tokens_used}) — "
                           "possible silent thinking mode. Check enable_thinking=False is applied.")
            return E("Model returned empty file content — the generation produced no output. "
                     "Try rephrasing the prompt.", 500)

        return Response(
            content=file_content.encode("utf-8"),
            media_type="application/octet-stream",
            headers={
                "Content-Disposition": f'attachment; filename="{filename}"',
                "X-Format":           fmt,
                "X-Tokens-Used":      str(tokens_used),
                "X-Lines":            str(len(file_content.splitlines())),
            })
    except asyncio.TimeoutError:
        await llama._reset()  # prevent broken connection on next request
        return E("Generation timed out",504)
    except Exception as e:
        logger.error(f"[generate] {e}",exc_info=True)
        return E(f"Generation error: {e}",500)


# ─────────────────────────────────────────────────────────────────────────────
# WEB SEARCH  [FIX-04: empty results shown clearly]
# ─────────────────────────────────────────────────────────────────────────────
@app.post("/v1/search", response_model=None)
@app.get("/v1/search", response_model=None)
async def search_endpoint(request: Request) -> Response:
    """
    Web search via DuckDuckGo — free, no API key.

    GET  /v1/search?q=your+query&n=5&ask=summarize
    POST {"query":"...","max_results":5,"ask":"analyze these results"}
    """
    if not auth(request): return E("Unauthorized",401)

    if request.method == "GET":
        params  = dict(request.query_params)
        query   = params.get("q","")
        max_res = int(params.get("n","5"))
        ask     = params.get("ask","")
    else:
        try: body = orjson.loads(await request.body())
        except: body = {}
        query   = body.get("query","")
        max_res = int(body.get("max_results",5))
        ask     = body.get("ask","")

    if not query: return E("'query' or 'q' parameter required",400)

    results = await web_search(query, min(max_res, 10))

    # [FIX-04] Explicit message when no results
    if not results:
        no_results_msg = (
            f"No results found for '{query}'. "
            "This may be due to network restrictions or a very specific query. "
            "Try rephrasing or broadening your search."
        )
        return J({"query":query,"results":[],"count":0,
                   "message":no_results_msg})

    if not ask:
        return J({"query":query,"results":results,"count":len(results)})

    results_text = "\n\n".join(
        f"[{i+1}] {r['title']}\n{r['snippet']}\nURL: {r['url']}"
        for i, r in enumerate(results))
    pl = {
        "messages":[
            {"role":"system","content":"You are a research assistant. Analyze web search results accurately."},
            {"role":"user","content":f"Query: {query}\n\nResults:\n{results_text}\n\nTask: {ask}"},
        ],
        "max_tokens":1024,"temperature":0.3,"cache_prompt":False,
    }
    analysis = ""
    try:
        r = await asyncio.wait_for(
            llama.req("POST","/v1/chat/completions",
                      content=orjson.dumps(pl),
                      headers={"Content-Type":"application/json"}),
            timeout=60.0)
        if r.status_code == 200:
            analysis = orjson.loads(r.content)["choices"][0]["message"]["content"]
    except Exception:
        analysis = "[Analysis unavailable]"

    return J({"query":query,"results":results,"count":len(results),"analysis":analysis})


# ─────────────────────────────────────────────────────────────────────────────
# AUDIO TRANSCRIPTION
# ─────────────────────────────────────────────────────────────────────────────
@app.post("/v1/audio/transcriptions", response_model=None)
async def audio_transcription(
    request:  Request,
    file:     UploadFile = File(...),
    language: str = Form(default=""),
    prompt:   str = Form(default=""),
) -> Response:
    if not auth(request): return E("Unauthorized",401)
    try:
        import aiohttp
        data = await read_upload_chunked(file, MAX_FILE_MB)  # [FIX-02]
        if data is None: return E(f"Audio file too large (max {MAX_FILE_MB}MB)",413)

        form = aiohttp.FormData()
        form.add_field("file",data,filename=file.filename or "audio.wav",
                       content_type=file.content_type or "audio/wav")
        form.add_field("model","whisper-base")
        if language: form.add_field("language",language)
        if prompt:   form.add_field("prompt",prompt)

        async with aiohttp.ClientSession() as s:
            async with s.post(f"{WHISPER_URL}/v1/audio/transcriptions",
                              data=form,timeout=aiohttp.ClientTimeout(total=120)) as r:
                ct_resp = r.headers.get("content-type","")
                if "json" in ct_resp:
                    return J(await r.json(), r.status)
                text = await r.text()
                return J({"text": text}, r.status)
    except Exception as e:
        logger.error(f"[transcription] {e}",exc_info=True)
        return E(f"Transcription error: {e}",500,"audio_error")


# ─────────────────────────────────────────────────────────────────────────────
# JOB QUEUE ENDPOINTS
# ─────────────────────────────────────────────────────────────────────────────
@app.post("/v1/jobs/submit", response_model=None)
async def job_submit(request: Request) -> Response:
    if not auth(request): return E("Unauthorized",401)
    if not _INFERENCE_READY: return E("Inference engine still loading.",503,"not_ready")
    try: body = orjson.loads(await request.body())
    except Exception as e: return E(f"Invalid JSON: {e}",400)

    body["cache_prompt"] = True
    body["max_tokens"]   = min(int(body.get("max_tokens") or MAX_NEW_TOKENS), MAX_NEW_TOKENS)
    body.pop("model",None); body.pop("stream",None)

    msgs = body.get("messages",[])
    if msgs and not any(m.get("role")=="system" for m in msgs):
        body["messages"] = [{"role":"system","content":DEFAULT_SYSTEM}] + msgs

    try:
        jid = await jobs.submit(body)
        return J({"job_id":jid,"status":"pending",
                   "poll_url":f"/v1/jobs/{jid}",
                   "result_url":f"/v1/jobs/{jid}/result"})
    except RuntimeError as e:
        return E(str(e),503,"server_overloaded")


@app.get("/v1/jobs/{job_id}", response_model=None)
async def job_status(job_id: str, request: Request) -> Response:
    if not auth(request): return E("Unauthorized",401)
    j = jobs.get(job_id)
    if not j: return E(f"Job '{job_id}' not found",404,"not_found")
    return J(j.to_dict())


@app.get("/v1/jobs/{job_id}/result", response_model=None)
async def job_result(job_id: str, request: Request) -> Response:
    if not auth(request): return E("Unauthorized",401)
    j = jobs.get(job_id)
    if not j: return E(f"Job '{job_id}' not found",404,"not_found")
    if j.status in (JobStatus.PENDING, JobStatus.RUNNING):
        return J({"job_id":job_id,"status":j.status.value,
                   "elapsed_s":round(time.monotonic()-j.created_at,1),
                   "message":"Job is still processing."}, 202)
    if j.status == JobStatus.FAILED:
        return E(f"Job failed: {j.error}",500,"job_failed")
    if j.result is None:
        return E("Job completed but result is empty",500,"job_empty")
    return Response(j.result,200,media_type="application/json",
                    headers={"X-Job-ID":job_id,
                             "X-Elapsed":str(j.to_dict()["elapsed_s"])})


# ─────────────────────────────────────────────────────────────────────────────
# MODEL MANAGER ENDPOINTS
# ─────────────────────────────────────────────────────────────────────────────
@app.get("/v1/model-manager/status", response_model=None)
async def mm_status(request: Request) -> Response:
    if not auth(request): return E("Unauthorized",401)
    return J(mgr.status())

@app.get("/v1/model-manager/catalog", response_model=None)
async def mm_catalog(request: Request) -> Response:
    if not auth(request): return E("Unauthorized",401)
    return J({"catalog":mgr.catalog(),"active":mgr.cfg.active_id})

@app.post("/v1/model-manager/check-updates", response_model=None)
async def mm_check(request: Request) -> Response:
    if not auth(request): return E("Unauthorized",401)
    try:
        await mgr._smart_upgrade()
        return J({"status":"checked","pending":mgr.cfg.pending_id or None})
    except Exception as e: return E(str(e),500)

@app.post("/v1/model-manager/switch", response_model=None)
async def mm_switch(request: Request) -> Response:
    if not auth(request): return E("Unauthorized",401)
    try:
        body = orjson.loads(await request.body())
        mid  = body.get("model_id","")
        from model_manager import CATALOG_BY_ID as _C
        if mid not in _C:
            return E(f"Unknown model_id: '{mid}'. Valid: {list(_C)}",400)
        m    = _C[mid]
        path,mmproj = await mgr.download(m)
        if not path: return E(f"Failed to download '{mid}'",500)
        mgr.cfg.pending_id   = mid
        mgr.cfg.pending_path = path
        mgr.cfg.save()
        return J({"status":"pending_restart","model":m.name,
                   "message":"Downloaded. Will activate on next container restart."})
    except Exception as e: return E(str(e),500)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0",
                port=int(os.getenv("API_PORT","7860")),
                loop="uvloop", http="h11", workers=1)