"""DeepSeek Harness (dsh) — per-user chat instance.

A single-user FastAPI chat web app that streams answers from OmniRoute
(http://192.168.20.13:20129/v1, OpenAI-compatible). One instance per family
member on .13 (sam:3081, jo:3082, harry:3083, finn:3084). No cross-user
account system — the console routes each user to their own instance.

Iframe-friendly so the console can embed it inline:
  - sets its own CSP `frame-ancestors https://console.lab.audasmedia.com.au`
  - does NOT send X-Frame-Options: DENY
"""
from __future__ import annotations

import html as _html
import json
import os
import base64
import re
import shutil
import subprocess
import threading
from pathlib import Path
from typing import Annotated
from uuid import uuid4
from datetime import datetime, timezone

import httpx
from fastapi import FastAPI, File, Form, Request, UploadFile
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates

BASE_DIR = Path(__file__).resolve().parent

# --- config from env ---
USER_NAME = os.getenv("DSH_USER", "user")
DEFAULT_SLOT = os.getenv("DSH_SLOT", "Welcome to DeepSeek Harness")
LLM_BASE = os.getenv("DSH_LLM_BASE", "http://192.168.20.13:20129/v1")
LLM_MODEL = os.getenv("DSH_LLM_MODEL", "opencode-go/deepseek-v4-flash")
LLM_KEY = os.getenv("DSH_LLM_KEY", "")
SYSTEM_PROMPT = os.getenv(
    "DSH_SYSTEM_PROMPT",
    f"You are {USER_NAME}'s helpful DeepSeek assistant on the family home lab. "
    "Be clear, safe and concise. Never expose server file contents or system "
    "secrets. If asked something unsafe, decline politely.",
)

# --- vision (image ingest) -> OpenRouter gpt-5 (works; OmniRoute auto/best-vision
# image route is currently unreliable). Sent as a separate provider call. ---
VISION_BASE = os.getenv("DSH_VISION_BASE", "https://openrouter.ai/api/v1")
VISION_MODEL = os.getenv("DSH_VISION_MODEL", "openai/gpt-5")
VISION_KEY = os.getenv("DSH_VISION_KEY", "")
IMAGE_MODEL = os.getenv("DSH_IMAGE_MODEL", "openai/gpt-image-1")
VIDEO_MODEL = os.getenv("DSH_VIDEO_MODEL", "z-ai/glm-5v-turbo")

# --- per-user conversations (multiple sessions, persist in the history volume) ---
HISTORY_LIMIT = 60
HISTORY_DIR = os.getenv("DSH_HISTORY_DIR", "/workspace")
SESSIONS_DIR = os.path.join(HISTORY_DIR, "sessions")
_history_lock = threading.Lock()


def _spath(sid):
    safe = re.sub(r"[^A-Za-z0-9_-]", "", sid or "") or "main"
    return os.path.join(SESSIONS_DIR, f"{safe}.json")


def _sdir():
    os.makedirs(SESSIONS_DIR, exist_ok=True)


def list_sessions():
    _sdir()
    out = []
    for f in os.listdir(SESSIONS_DIR):
        if not f.endswith(".json"):
            continue
        try:
            with open(os.path.join(SESSIONS_DIR, f), encoding="utf-8") as fh:
                d = json.load(fh)
            out.append({
                "id": d.get("id", f[:-5]),
                "name": d.get("name", "Conversation")[:40],
                "updated": d.get("updated", ""),
                "count": len(d.get("messages", [])),
            })
        except Exception:
            continue
    out.sort(key=lambda x: x["updated"] or "", reverse=True)
    return out


def load_session(sid):
    try:
        with open(_spath(sid), encoding="utf-8") as fh:
            d = json.load(fh)
        msgs = [m for m in d.get("messages", []) if m.get("role") in ("user", "assistant")]
        return {"id": d.get("id", sid or "main"), "name": d.get("name", "Conversation"),
                "messages": msgs[-HISTORY_LIMIT:],
                "model": d.get("model") or LLM_MODEL}
    except Exception:
        return {"id": sid or "main", "name": "Conversation", "messages": [],
                "model": LLM_MODEL}


def save_session(sid, msgs, model=None):
    _sdir()
    sid = (sid or "main")
    with _history_lock:
        d = load_session(sid)
        if model:
            d["model"] = model
        d["messages"] = msgs[-HISTORY_LIMIT:]
        d["updated"] = datetime.now(timezone.utc).isoformat()
        if not d.get("name") or d["name"] in ("Conversation", "New conversation"):
            for m in msgs:
                if m.get("role") == "user" and m.get("content"):
                    d["name"] = m["content"][:40]
                    break
        d["id"] = sid
        try:
            with open(_spath(sid), "w", encoding="utf-8") as fh:
                json.dump(d, fh, ensure_ascii=False)
        except Exception as exc:  # noqa: BLE001
            print(f"[dsh] session save failed: {exc}")


app = FastAPI(title=f"DeepSeek Harness — {USER_NAME}", docs_url=None)

# --- plugin-style tools (Phase 1: web + document summaries) ---
WORKSPACE_DIR = os.getenv("DSH_WORKSPACE", "/workspace")
_SUM_HEAD = (
    "You are a friendly family assistant. Summarise the given content into "
    "clear, concise bullet points in plain language. Keep it under ~200 words."
)


def _strip_html(s: str) -> str:
    s = re.sub(r"(?is)<(script|style|head|noscript|svg)[^>]*>.*?</\1>", " ", s)
    s = re.sub(r"(?s)<[^>]+>", " ", s)
    return _html.unescape(re.sub(r"\s+", " ", s)).strip()


async def _complete(messages) -> str:
    """Non-streaming OmniRoute call for summarisation tools."""
    payload = {"model": LLM_MODEL, "messages": messages, "stream": False, "temperature": 0.3}
    headers = {"Authorization": f"Bearer {LLM_KEY}"} if LLM_KEY else {}
    url = f"{LLM_BASE.rstrip('/')}/chat/completions"
    async with httpx.AsyncClient(timeout=180) as client:
        r = await client.post(url, json=payload, headers=headers)
        r.raise_for_status()
        data = r.json()
    return data["choices"][0]["message"]["content"]


async def _vision_complete(messages) -> str:
    """Non-streaming OpenRouter call to the multimodal video/vision model."""
    key = VISION_KEY
    base = VISION_BASE
    headers = {"Authorization": f"Bearer {key}"} if key else {}
    url = f"{base.rstrip('/')}/chat/completions"
    payload = {"model": VIDEO_MODEL, "messages": messages, "stream": False, "temperature": 0.3}
    async with httpx.AsyncClient(timeout=240) as client:
        r = await client.post(url, json=payload, headers=headers)
        r.raise_for_status()
    return r.json()["choices"][0]["message"]["content"]


def _extract_frames(video: str, max_frames: int = 6):
    """Sample up to N JPEG frames from a video with ffmpeg; return data URIs."""
    outdir = f"/tmp/frames_{uuid4().hex}"
    os.makedirs(outdir, exist_ok=True)
    try:
        dur = subprocess.run(
            ["ffprobe", "-v", "error", "-show_entries", "format=duration",
             "-of", "default=nw=1:nk=1", video],
            capture_output=True, text=True).stdout.strip()
        dur = float(dur) if dur else 10
        fps = max(0.1, min(1.0, max_frames / max(dur, 1)))
        subprocess.run(
            ["ffmpeg", "-v", "error", "-i", video, "-vf", f"fps={fps}",
             "-frames:v", str(max_frames), f"{outdir}/f%03d.jpg"],
            capture_output=True)
        frames = []
        for f in sorted(os.listdir(outdir)):
            p = os.path.join(outdir, f)
            if os.path.isfile(p) and os.path.getsize(p) > 0:
                with open(p, "rb") as fh:
                    frames.append("data:image/jpeg;base64," + base64.b64encode(fh.read()).decode())
        return frames
    finally:
        shutil.rmtree(outdir, ignore_errors=True)


@app.post("/api/summarize/web")
async def summarize_web(url: Annotated[str, Form()]) -> dict:
    """Fetch a URL and summarise its text content via the LLM."""
    try:
        async with httpx.AsyncClient(timeout=40, follow_redirects=True) as client:
            r = await client.get(url, headers={"User-Agent": "Mozilla/5.0 (dsh-summarizer)"})
            r.raise_for_status()
        text = _strip_html(r.text)
        if not text:
            return {"ok": False, "error": "Could not extract readable text."}
        content = text[:60000]
        summary = await _complete([
            {"role": "system", "content": _SUM_HEAD},
            {"role": "user", "content": f"URL: {url}\n\nCONTENT:\n{content}\n\nSummarise the key points."},
        ])
        return {"ok": True, "source": url, "summary": summary}
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "error": f"{type(exc).__name__}: {exc}"}


@app.post("/api/summarize/docs")
async def summarize_docs(path: Annotated[str, Form()]) -> dict:
    """Read a workspace file (txt/md/json/log) and summarise it via the LLM."""
    ws = str(Path(WORKSPACE_DIR).resolve())
    fp = (Path(ws) / path.lstrip("/")).resolve()
    if not str(fp).startswith(ws):
        return {"ok": False, "error": "Path must stay inside your workspace."}
    if not fp.is_file():
        return {"ok": False, "error": f"File not found: {path}"}
    try:
        data = fp.read_text(encoding="utf-8", errors="replace")[:60000]
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "error": f"Could not read: {exc}"}
    summary = await _complete([
        {"role": "system", "content": _SUM_HEAD},
        {"role": "user", "content": f"FILE: {path}\n\nCONTENT:\n{data}\n\nSummarise the key points."},
    ])
    return {"ok": True, "source": path, "summary": summary}
app.mount("/static", StaticFiles(directory=str(BASE_DIR / "static")), name="static")
templates = Jinja2Templates(directory=str(BASE_DIR / "templates"))
templates.env.globals["DSH_USER"] = USER_NAME
templates.env.globals["DSH_SLOT"] = DEFAULT_SLOT
templates.env.globals["DSH_MODEL"] = LLM_MODEL

# Frame-embedding policy: allow the console (and same origin) to frame us.
FRAME_POLICY = os.getenv("DSH_FRAME_POLICY", "https://console.lab.audasmedia.com.au")


@app.middleware("http")
async def frame_headers(request, call_next):
    response = await call_next(request)
    # Allow embedding from the console; refuse nothing else explicitly.
    # (CSP frame-ancestors is the modern control; no X-Frame-Options DENY.)
    response.headers.setdefault("Content-Security-Policy",
                                f"frame-ancestors 'self' {FRAME_POLICY}")
    response.headers.setdefault("X-Content-Type-Options", "nosniff")
    return response


@app.get("/", response_class=HTMLResponse)
async def index(request: Request) -> HTMLResponse:
    return templates.TemplateResponse(request, "chat.html", {})


@app.post("/api/chat")
async def chat(message: Annotated[str, Form()],
               image: Annotated[str, Form()] = "",
               session_id: Annotated[str, Form()] = "",
               model: Annotated[str, Form()] = "") -> StreamingResponse:
    """Stream an answer as SSE. When `image` is a data:image URI the vision
    model (OpenRouter) is used instead of the text LLM (image ingest)."""
    use_vision = bool(image and image.startswith("data:image/"))

    async def event_stream():
        sid = session_id or "main"
        session = load_session(sid)
        if model and model in MODEL_CHOICES:
            session["model"] = model
        hist = session["messages"]
        hist.append({"role": "user", "content": message})
        sel = session.get("model") or ""
        use_multimodal = use_vision or (sel in VISION_OVERRIDE_MODELS)
        if use_multimodal:
            base = VISION_BASE
            key = VISION_KEY
            mm_model = sel if (sel in VISION_OVERRIDE_MODELS) else VISION_MODEL
            content = [{"type": "text", "text": message or ("Describe this image." if use_vision else "")}]
            if use_vision:
                content.append({"type": "image_url", "image_url": {"url": image}})
            payload = {
                "model": mm_model,
                "messages": [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": content},
                ],
                "stream": True,
                "temperature": 0.4,
            }
        else:
            base = LLM_BASE
            key = LLM_KEY
            active_model = sel or LLM_MODEL
            if active_model not in MODEL_CHOICES:
                active_model = LLM_MODEL
            payload = {
                "model": active_model,
                "messages": [{"role": "system", "content": SYSTEM_PROMPT}] + hist,
                "stream": True,
                "temperature": 0.7,
            }
        headers = {"Authorization": f"Bearer {key}"} if key else {}
        url = f"{base.rstrip('/')}/chat/completions"
        full: list[str] = []
        try:
            async with httpx.AsyncClient(timeout=None) as client:
                async with client.stream("POST", url, json=payload, headers=headers) as r:
                    if r.status_code != 200:
                        body = (await r.aread())[:300].decode("utf-8", "replace")
                        yield f"data: {json.dumps({'e': f'{r.status_code} {body}'})}\n\n"
                        return
                    async for line in r.aiter_lines():
                        if not line or not line.startswith("data: "):
                            continue
                        data = line[6:].strip()
                        if data == "[DONE]":
                            break
                        try:
                            obj = json.loads(data)
                            delta = obj["choices"][0]["delta"].get("content")
                        except Exception:
                            continue
                        if delta:
                            full.append(delta)
                            yield f"data: {json.dumps({'c': delta})}\n\n"
        except Exception as exc:  # noqa: BLE001
            yield f"data: {json.dumps({'e': str(exc)})}\n\n"
            return
        if full:
            hist.append({"role": "assistant", "content": "".join(full)})
        else:
            yield f"data: {json.dumps({'e': 'Model returned no response — try auto/best-chat or auto/best-fast.'})}\n\n"
        save_session(sid, hist, session.get("model"))  # always persist user turn + chosen model

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )



@app.post("/api/tool/video")
async def analyze_video(file: Annotated[UploadFile, File()]) -> dict:
    """Analyze a video / screen recording: extract keyframes, send to the
    multimodal GLM model, return a description."""
    if file is None or not file.filename:
        return {"ok": False, "error": "No video chosen."}
    if not VISION_KEY:
        return {"ok": False, "error": "No vision API key configured."}
    vid = f"/tmp/video_{uuid4().hex}.mp4"
    try:
        data = await file.read()
        if len(data) > 200 * 1024 * 1024:
            return {"ok": False, "error": "Video too large (max 200 MB)."}
        with open(vid, "wb") as fh:
            fh.write(data)
        frames = _extract_frames(vid, max_frames=6)
        if not frames:
            return {"ok": False, "error": "Could not extract frames — is this a video?"}
        content = [
            {"type": "text", "text": "Analyze this video/screen recording: describe what is shown, "
                                     "the main actions, and any visible text or UI. Give a concise bullet summary."}
        ]
        for b64 in frames:
            content.append({"type": "image_url", "image_url": {"url": b64}})
        summary = await _vision_complete([{"role": "user", "content": content}])
        return {"ok": True, "frames": len(frames), "summary": summary}
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "error": f"{type(exc).__name__}: {exc}"}
    finally:
        if os.path.exists(vid):
            os.remove(vid)


@app.post("/api/tool/image")
async def generate_image(prompt: Annotated[str, Form()]) -> dict:
    """Generate an image from a prompt via OpenRouter (gpt-image-1); save to workspace."""
    if not VISION_KEY:
        return {"ok": False, "error": "No image API key configured."}
    prompt = prompt.strip()[:500]
    if not prompt:
        return {"ok": False, "error": "Prompt required."}
    try:
        headers = {"Authorization": f"Bearer {VISION_KEY}"}
        url = f"{VISION_BASE.rstrip('/')}/images/generations"
        payload = {"model": IMAGE_MODEL, "prompt": prompt, "n": 1, "response_format": "b64_json"}
        async with httpx.AsyncClient(timeout=180) as client:
            r = await client.post(url, json=payload, headers=headers)
            r.raise_for_status()
        b64 = r.json()["data"][0].get("b64_json")
        if not b64:
            return {"ok": False, "error": "No image returned."}
        img_dir = Path(WORKSPACE_DIR).resolve() / "images"
        img_dir.mkdir(parents=True, exist_ok=True)
        tag = uuid4().hex[:8]
        fname = img_dir / f"{tag}.png"
        fname.write_bytes(base64.b64decode(b64))
        return {"ok": True, "path": f"images/{tag}.png", "b64": "data:image/png;base64," + b64}
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "error": f"{type(exc).__name__}: {exc}"}



MODEL_CHOICES = ["opencode-go/deepseek-v4-flash",
                 "deepseek/deepseek-v4-flash",
                 "deepseek/deepseek-v4-pro",
                 "z-ai/glm-5v-turbo"]
MODEL_LABELS = {
    "opencode-go/deepseek-v4-flash": "Chat · opencode-go flash",
    "deepseek/deepseek-v4-flash": "Reasoning · deepseek flash",
    "deepseek/deepseek-v4-pro": "Coding · deepseek pro",
    "z-ai/glm-5v-turbo": "Multimodal · GLM vision",
}
VISION_OVERRIDE_MODELS = {"z-ai/glm-5v-turbo"}


@app.get("/api/models")
async def models_list() -> dict:
    return {"ok": True, "default": LLM_MODEL, "current": MODEL_CHOICES,
            "labels": MODEL_LABELS}


@app.get("/api/sessions")
async def sessions_list() -> dict:
    return {"ok": True, "sessions": list_sessions()}


@app.post("/api/sessions")
async def sessions_new() -> dict:
    _sdir()
    new_id = uuid4().hex[:12]
    now = datetime.now(timezone.utc).isoformat()
    try:
        with open(_spath(new_id), "w", encoding="utf-8") as fh:
            json.dump({"id": new_id, "name": "New conversation", "messages": [],
                       "created": now, "updated": now}, fh)
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "error": str(exc)}
    return {"ok": True, "id": new_id}


@app.get("/api/sessions/{sid}")
async def sessions_get(sid: str) -> dict:
    s = load_session(sid)
    return {"ok": True, "id": s["id"], "name": s["name"], "messages": s["messages"],
            "model": s.get("model") or LLM_MODEL}


@app.post("/api/sessions/{sid}/delete")
async def sessions_delete(sid: str) -> dict:
    try:
        p = _spath(sid)
        if os.path.exists(p):
            os.remove(p)
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "error": str(exc)}
    return {"ok": True}


@app.get("/healthz")
async def healthz() -> dict:
    return {"ok": True, "user": USER_NAME, "model": LLM_MODEL}