#!/usr/bin/env python3
"""sjcode daemon.

Runs on the server, talks to the local ollama API (127.0.0.1) and serves
sjcode clients on its own port (from config.toml).

On startup it drops into a minimal admin CLI; while that CLI runs, an
HTTP server thread serves the client API:

    POST /login            {"username", "password"}  -> {"token"}
    POST /chat             {"token", "message", "model"?}
                           -> {"response"} or {"error"} with 4xx
    GET  /models?token=... -> {"models": [...]}

User accounts live in SQLite (sjcode.db) with:
    username, password_hash, last_logged_in, role, used_today (minutes).

Roles decide the daily usage limit (minutes) from config.toml [limits].
"""

from __future__ import annotations

import cmd
import hashlib
import hmac
import json
import os
import sqlite3
import sys
import threading
import time
import tomllib
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

ROOT = Path(__file__).resolve().parent
# SJCODE_CONFIG lets you run a second daemon (tests, staging) without
# touching the regular config; the database is created next to that file.
CONFIG_FILE = Path(os.environ.get("SJCODE_CONFIG") or (ROOT / "config.toml"))
DB_FILE = ROOT / "sjcode.db"  # overwritten by config value once loaded

PEPPER = b"sjcode-v1"
SESSION_TTL = 24 * 3600  # login token valid for one day

CONF: dict = {
    "server": {"host": "127.0.0.1", "port": 8420},
    "ollama": {"base_url": "http://127.0.0.1:11434", "model": "llama3.2",
               "num_ctx": 16384},
    "database": {"path": "sjcode.db"},
    "limits": {"user": 120, "admin": -1},
}


def load_config() -> None:
    """Read config.toml on top of the defaults."""
    global DB_FILE
    if not CONFIG_FILE.exists():
        sys.exit(
            f"error: {CONFIG_FILE.name} not found.\n"
            f"Copy {ROOT / 'config.toml.example'} to {CONFIG_FILE} and edit it."
        )
    try:
        with open(CONFIG_FILE, "rb") as fh:
            loaded = tomllib.load(fh)
    except tomllib.TOMLDecodeError as exc:
        sys.exit(f"error: invalid {CONFIG_FILE.name}: {exc}")

    for section, values in loaded.items():
        if isinstance(values, dict) and isinstance(CONF.get(section), dict):
            CONF[section].update(values)
    DB_FILE = CONFIG_FILE.parent / CONF["database"]["path"]


# --------------------------------------------------------------------------
# database
# --------------------------------------------------------------------------

SCHEMA = """
CREATE TABLE IF NOT EXISTS users (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    username        TEXT    NOT NULL UNIQUE,
    password_hash   TEXT    NOT NULL,
    role            TEXT    NOT NULL DEFAULT 'user',
    used_today      REAL    NOT NULL DEFAULT 0,
    last_logged_in  TEXT    NOT NULL DEFAULT '',
    created_at      TEXT    NOT NULL
);
CREATE TABLE IF NOT EXISTS sessions (
    token      TEXT PRIMARY KEY,
    username   TEXT NOT NULL,
    created_at INTEGER NOT NULL
);
"""


def db() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_FILE, timeout=10)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    return conn


def init_db() -> None:
    with db() as conn:
        conn.executescript(SCHEMA)


def hash_password(password: str) -> str:
    salt = f"{time.time_ns()}:{threading.get_ident()}".encode()
    digest = hashlib.scrypt(password.encode(), salt=PEPPER + salt, n=2**14, r=8, p=1)
    return f"{salt.decode()}${digest.hex()}"


def verify_password(password: str, stored: str) -> bool:
    try:
        salt_s, hexdigest = stored.split("$", 1)
    except ValueError:
        return False
    calc = hashlib.scrypt(
        password.encode(), salt=PEPPER + salt_s.encode(), n=2**14, r=8, p=1
    )
    return hmac.compare_digest(calc.hex(), hexdigest)


def new_token() -> str:
    import secrets

    return secrets.token_hex(32)


def limit_for(role: str) -> float:
    limits = CONF["limits"]
    return float(limits.get(role, limits.get("user", 0)))


def ask_password(prompt: str) -> str:
    """Hidden input on a terminal, plain stdin otherwise (scriptable)."""
    if sys.stdin.isatty():
        import getpass

        return getpass.getpass(prompt)
    print(prompt, end="", flush=True)
    return input("")


def usage_reset_if_needed(conn: sqlite3.Connection, row: sqlite3.Row) -> None:
    """used_today is a counter in minutes; reset it on a new calendar day (UTC)."""
    last_date = row["last_logged_in"][:10]
    if last_date and last_date != time.strftime("%Y-%m-%d", time.gmtime()):
        conn.execute(
            "UPDATE users SET used_today = 0 WHERE username = ?", (row["username"],)
        )
        conn.commit()


def get_user(username: str):
    with db() as conn:
        row = conn.execute(
            "SELECT * FROM users WHERE username = ?", (username,)
        ).fetchone()
        if row is not None:
            usage_reset_if_needed(conn, row)
            conn.commit()
        return row


def remaining_minutes(username: str) -> float:
    row = get_user(username)
    if row is None:
        return 0.0
    limit = limit_for(row["role"])
    if limit < 0:  # unlimited
        return float("inf")
    return max(0.0, limit - float(row["used_today"]))


def add_used_minutes(username: str, minutes: float) -> None:
    with db() as conn:
        conn.execute(
            "UPDATE users SET used_today = used_today + ? WHERE username = ?",
            (minutes, username),
        )


def set_last_login(username: str) -> None:
    stamp = time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime())
    with db() as conn:
        conn.execute(
            "UPDATE users SET last_logged_in = ? WHERE username = ?", (stamp, username)
        )


def create_session(username: str) -> str:
    token = new_token()
    with db() as conn:
        conn.execute(
            "INSERT INTO sessions (token, username, created_at) VALUES (?, ?, ?)",
            (token, username, int(time.time())),
        )
    return token


def session_user(token: str):
    """Return the username for a valid session token, else None."""
    with db() as conn:
        row = conn.execute(
            "SELECT username, created_at FROM sessions WHERE token = ?", (token,)
        ).fetchone()
    if row is None or int(time.time()) - row["created_at"] > SESSION_TTL:
        return None
    return row["username"]


# --------------------------------------------------------------------------
# ollama
# --------------------------------------------------------------------------

def ollama_json(path: str, payload: dict | None = None, timeout: int = 300) -> tuple[int, dict]:
    url = CONF["ollama"]["base_url"].rstrip("/") + path
    data = None if payload is None else json.dumps(payload).encode()
    req = urllib.request.Request(
        url, data=data, headers={"Content-Type": "application/json"}
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, json.loads(resp.read().decode())
    except urllib.error.HTTPError as exc:
        try:
            body = json.loads(exc.read().decode())
        except Exception:
            body = {"error": str(exc)}
        return exc.code, body
    except (urllib.error.URLError, TimeoutError, ConnectionError) as exc:
        return 502, {"error": f"ollama not reachable: {exc}"}


def ollama_stream(path: str, payload: dict) -> tuple[int, object]:
    """POST to ollama and return the raw response to iterate line by line."""
    url = CONF["ollama"]["base_url"].rstrip("/") + path
    data = json.dumps(payload).encode()
    req = urllib.request.Request(
        url, data=data, headers={"Content-Type": "application/json"}
    )
    try:
        resp = urllib.request.urlopen(req, timeout=600)
    except urllib.error.HTTPError as exc:
        try:
            body = json.loads(exc.read().decode())
        except Exception:
            body = {"error": str(exc)}
        return exc.code, str(body.get("error", "ollama error"))
    except (urllib.error.URLError, TimeoutError, ConnectionError) as exc:
        return 502, {"error": f"ollama not reachable: {exc}"}
    return resp.status, resp  # caller iterates the lines and closes it


def ollama_generate(prompt: str, model: str | None, stream: bool = False):
    """Generate a reply. Returns (status, text) or (200, token-iterator)."""
    options = {"num_ctx": int(CONF["ollama"].get("num_ctx", 16384))}
    payload = {
        "model": model or CONF["ollama"]["model"],
        "prompt": prompt,
        "stream": stream,
        "options": options,
    }
    if stream:
        status, body = ollama_stream("/api/generate", payload)
        if status != 200:
            return status, str(body)

        def tokens():
            try:
                for raw in body:
                    line = raw.decode() if isinstance(raw, bytes) else raw
                    if not line.strip():
                        continue
                    try:
                        obj = json.loads(line)
                    except json.JSONDecodeError:
                        continue
                    if obj.get("error"):
                        yield f"\n[ollama error: {obj['error']}]"
                        return
                    piece = obj.get("response", "")
                    if piece:
                        yield piece
            finally:
                body.close()

        return 200, tokens()

    status, body = ollama_json("/api/generate", payload)
    if status != 200:
        return status, str(body.get("error", "ollama error"))
    return 200, str(body.get("response", ""))


def list_models() -> tuple[int, list[str]]:
    status, body = ollama_json("/api/tags", timeout=10)
    if status != 200:
        return status, []
    names = [m.get("name", "") for m in body.get("models", [])]
    return 200, names


# --------------------------------------------------------------------------
# HTTP API
# --------------------------------------------------------------------------

class ApiHandler(BaseHTTPRequestHandler):
    server_version = "sjcode-daemon/0.2"
    protocol_version = "HTTP/1.1"  # required for chunked streamed replies

    # -- helpers ----------------------------------------------------------
    def _json(self, status: int, payload: dict) -> None:
        body = json.dumps(payload).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _read_body(self) -> dict:
        try:
            length = int(self.headers.get("Content-Length", "0"))
            return json.loads(self.rfile.read(length).decode())
        except (ValueError, json.JSONDecodeError):
            return {}

    def _respond_stream(self, tokens, model: str, started: float,
                        username: str) -> None:
        """Stream a reply as JSON lines: {"chunk": ...} ... {"response": ...}."""
        self.send_response(200)
        self.send_header("Content-Type", "application/x-ndjson")
        self.send_header("Transfer-Encoding", "chunked")
        self.end_headers()

        def send_line(obj: dict) -> None:
            data = (json.dumps(obj) + "\n").encode()
            self.wfile.write(f"{len(data):x}\r\n".encode() + data + b"\r\n")
            self.wfile.flush()

        parts: list[str] = []
        abandoned = False
        try:
            for piece in tokens:
                parts.append(piece)
                send_line({"chunk": piece})
        except (BrokenPipeError, ConnectionResetError):
            abandoned = True  # client went away; still bill the usage below
        except OSError as exc:
            parts.append(f"\n[stream error: {exc}]")
            send_line({"chunk": f"\n[stream error: {exc}]"})
        finally:
            text = "".join(parts)
            elapsed = time.monotonic() - started
            used_min = elapsed / 60.0
            add_used_minutes(username, used_min)
            if not abandoned:
                left_now = remaining_minutes(username)
                send_line({
                    "response": text,
                    "model": model,
                    "used_minutes": round(used_min, 4),
                    "remaining_minutes": None if left_now == float("inf")
                    else round(left_now, 2),
                })
                self.wfile.write(b"0\r\n\r\n")
                self.wfile.flush()
                print(f"[chat] {username}: {len(text)} chars in "
                      f"{elapsed:.1f}s (streamed)")

    # -- routes -----------------------------------------------------------
    def do_GET(self) -> None:  # noqa: N802 (http.server API)
        if self.path.startswith("/models"):
            status, models = list_models()
            self._json(200 if status == 200 else 502, {"models": models})
            return
        if self.path == "/status":
            self._json(200, {"status": "ok", "daemon": "sjcode"})
            return
        self._json(404, {"error": "not found"})

    def do_POST(self) -> None:  # noqa: N802 (http.server API)
        data = self._read_body()

        if self.path == "/login":
            self._handle_login(data)
            return
        if self.path == "/chat":
            self._handle_chat(data)
            return
        self._json(404, {"error": "not found"})

    # -- handlers ---------------------------------------------------------
    def _handle_login(self, data: dict) -> None:
        username = str(data.get("username", "")).strip()
        password = str(data.get("password", ""))
        if not username or not password:
            self._json(400, {"error": "username and password required"})
            return

        row = get_user(username)
        if row is None or not verify_password(password, row["password_hash"]):
            self._json(401, {"error": "invalid credentials"})
            return

        limit = limit_for(row["role"])
        left = remaining_minutes(username)
        if left <= 0:
            self._json(
                403,
                {"error": "daily limit reached", "limit_minutes": limit},
            )
            return

        set_last_login(username)
        token = create_session(username)
        self._json(
            200,
            {
                "token": token,
                "role": row["role"],
                "remaining_minutes": None if left == float("inf") else round(left, 1),
            },
        )

    def _handle_chat(self, data: dict) -> None:
        token = str(data.get("token", ""))
        username = session_user(token)
        if username is None:
            self._json(401, {"error": "not logged in"})
            return

        message = str(data.get("message", "")).strip()
        if not message:
            self._json(400, {"error": "empty message"})
            return

        left = remaining_minutes(username)
        if left <= 0:
            self._json(403, {"error": "daily limit reached"})
            return

        model = str(data.get("model") or "") or CONF["ollama"]["model"]
        want_stream = bool(data.get("stream"))

        started = time.monotonic()
        status, answer = ollama_generate(message, model, stream=want_stream)

        if want_stream and status == 200:
            self._respond_stream(answer, model, started, username)
            return

        elapsed = time.monotonic() - started

        if status != 200:
            self._json(status, {"error": answer})
            return

        used_min = elapsed / 60.0
        add_used_minutes(username, used_min)
        left_now = remaining_minutes(username)
        self._json(
            200,
            {
                "response": answer,
                "model": model,
                "used_minutes": round(used_min, 4),
                "remaining_minutes": None if left_now == float("inf") else round(left_now, 2),
            },
        )

    def log_message(self, fmt: str, *args) -> None:  # quiet by default
        return


# --------------------------------------------------------------------------
# admin CLI
# --------------------------------------------------------------------------

class DaemonCLI(cmd.Cmd):
    """Minimal admin shell shown after daemon startup."""

    intro = (
        "sjcode daemon CLI. Type 'help' for commands.\n"
        "Accounts are stored in SQLite; the HTTP API runs in the background."
    )
    prompt = "sjcode-admin> "

    def __init__(self, server_ref: dict) -> None:
        super().__init__()
        self.server_ref = server_ref

    # -- account management ----------------------------------------------
    def do_create(self, arg: str) -> None:
        """create <username> [role]  - create an account (asks for password)"""
        parts = arg.split()
        if not parts:
            print("usage: create <username> [role]")
            return
        username = parts[0]
        role = parts[1] if len(parts) > 1 else "user"
        if get_user(username) is not None:
            print(f"user '{username}' already exists")
            return
        pw1 = ask_password("password: ")
        pw2 = ask_password("password (again): ")
        if len(pw1) < 4:
            print("password too short (min 4 chars)")
            return
        if pw1 != pw2:
            print("passwords do not match")
            return
        with db() as conn:
            conn.execute(
                "INSERT INTO users (username, password_hash, role, created_at) "
                "VALUES (?, ?, ?, ?)",
                (
                    username,
                    hash_password(pw1),
                    role,
                    time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime()),
                ),
            )
        print(f"created '{username}' with role '{role}' "
              f"(limit: {limit_for(role)} min/day)")

    def do_list(self, arg: str) -> None:
        """list - show all accounts"""
        with db() as conn:
            rows = conn.execute(
                "SELECT username, role, used_today, last_logged_in "
                "FROM users ORDER BY username"
            ).fetchall()
        if not rows:
            print("no accounts yet. use 'create <username>'")
            return
        print(f"{'username':<20} {'role':<10} {'used today':>10}  last login")
        for r in rows:
            print(
                f"{r['username']:<20} {r['role']:<10} "
                f"{r['used_today']:>10.1f}  {r['last_logged_in'] or '-'}"
            )

    def do_delete(self, arg: str) -> None:
        """delete <username> - remove an account"""
        username = arg.strip().split()[0] if arg.strip() else ""
        if not username:
            print("usage: delete <username>")
            return
        with db() as conn:
            cur = conn.execute("DELETE FROM users WHERE username = ?", (username,))
        if cur.rowcount:
            print(f"deleted '{username}'")
        else:
            print(f"no user '{username}'")

    def do_role(self, arg: str) -> None:
        """role <username> <role> - change a user's role"""
        parts = arg.split()
        if len(parts) != 2:
            print("usage: role <username> <role>")
            return
        username, role = parts
        with db() as conn:
            cur = conn.execute(
                "UPDATE users SET role = ? WHERE username = ?", (role, username)
            )
        if cur.rowcount:
            print(f"'{username}' now has role '{role}' "
                  f"(limit: {limit_for(role)} min/day)")
        else:
            print(f"no user '{username}'")

    def do_passwd(self, arg: str) -> None:
        """passwd <username> - set a new password for a user"""
        username = arg.strip().split()[0] if arg.strip() else ""
        if not username:
            print("usage: passwd <username>")
            return
        row = get_user(username)
        if row is None:
            print(f"no user '{username}'")
            return
        pw = ask_password("new password: ")
        if len(pw) < 4:
            print("password too short (min 4 chars)")
            return
        with db() as conn:
            conn.execute(
                "UPDATE users SET password_hash = ? WHERE username = ?",
                (hash_password(pw), username),
            )
        print(f"password updated for '{username}'")

    def do_usage(self, arg: str) -> None:
        """usage [username] - show used minutes / limit for today"""
        with db() as conn:
            if arg.strip():
                rows = conn.execute(
                    "SELECT username, role, used_today FROM users "
                    "WHERE username = ?",
                    (arg.strip(),),
                ).fetchall()
            else:
                rows = conn.execute(
                    "SELECT username, role, used_today FROM users "
                    "ORDER BY username"
                ).fetchall()
        if not rows:
            print("no such user" if arg.strip() else "no accounts yet")
            return
        for r in rows:
            limit = limit_for(r["role"])
            lim = "unlimited" if limit < 0 else f"{limit:g} min"
            print(f"{r['username']:<20} {r['used_today']:>7.1f} min / {lim}")

    def do_reset(self, arg: str) -> None:
        """reset [username] - reset the used_today counter"""
        with db() as conn:
            if arg.strip():
                conn.execute(
                    "UPDATE users SET used_today = 0 WHERE username = ?",
                    (arg.strip(),),
                )
            else:
                conn.execute("UPDATE users SET used_today = 0")
        print("usage counter reset")

    # -- server control ---------------------------------------------------
    def do_status(self, arg: str) -> None:
        """status - show daemon settings"""
        host, port = CONF["server"]["host"], CONF["server"]["port"]
        print(f"listening on http://{host}:{port}")
        print(f"ollama at {CONF['ollama']['base_url']} (model: {CONF['ollama']['model']})")
        print(f"database at {DB_FILE}")

    def do_quit(self, arg: str) -> bool:
        """quit - stop the daemon"""
        print("shutting down...")
        self.server_ref["stop"] = True
        srv: ThreadingHTTPServer = self.server_ref["httpd"]
        srv.shutdown()
        return True

    def do_exit(self, arg: str) -> bool:
        """exit - stop the daemon (same as quit)"""
        return self.do_quit(arg)

    do_EOF = do_quit


# --------------------------------------------------------------------------
# main
# --------------------------------------------------------------------------

def main() -> None:
    load_config()
    init_db()

    host = CONF["server"]["host"]
    port = int(CONF["server"]["port"])

    try:
        httpd = ThreadingHTTPServer((host, port), ApiHandler)
    except OSError as exc:
        sys.exit(f"error: cannot listen on {host}:{port}: {exc}")

    ref = {"httpd": httpd, "stop": False}
    thread = threading.Thread(target=httpd.serve_forever, daemon=True)
    thread.start()

    print(f"sjcode daemon listening on http://{host}:{port}")
    print(f"ollama: {CONF['ollama']['base_url']}  model: {CONF['ollama']['model']}")
    print(f"database: {DB_FILE}")

    try:
        DaemonCLI(ref).cmdloop()
    except KeyboardInterrupt:
        print()
    finally:
        httpd.shutdown()
        httpd.server_close()
        print("bye")


if __name__ == "__main__":
    main()