import asyncio
import json
import secrets
from pathlib import Path

from fastapi import (Depends, FastAPI, File, HTTPException, Request, Response,
                     UploadFile, WebSocket, WebSocketDisconnect)
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

from . import auth, db, registry, settings as settings_mod, store, sysinfo
from .runner import runner

app = FastAPI(title="MODELBEAST")
WEB_DIST = db.ROOT / "web" / "dist"
INBOX = db.DATA / "inbox"
MAX_UPLOAD = 1024 * 1024 * 1024  # 1 GB


@app.on_event("startup")
async def startup():
    runner.operators = registry.load_operators()
    con = db.connect()
    app.state.con = con
    runner.get_settings = lambda: settings_mod.get_all(con)
    auth.ensure_auth_secret(con)
    _bootstrap_owner(con)
    asyncio.create_task(runner.worker())
    asyncio.create_task(watch_inbox())


def _bootstrap_owner(con):
    if con.execute("SELECT 1 FROM users LIMIT 1").fetchone():
        return
    pw = secrets.token_urlsafe(12)
    auth.create_user(con, "monster", pw, role="owner", max_active_jobs=999)
    banner = "=" * 62
    print(f"\n{banner}\n MODELBEAST owner account created (first run):\n"
          f"   username: monster\n   password: {pw}\n"
          f"   change it:  MB_TOKEN=... mb user passwd monster   (or the UI)\n{banner}\n",
          flush=True)


def _con(request: Request):
    return request.app.state.con


def _with_usernames(con, rows: list[dict]) -> list[dict]:
    """Attach a 'username' field to job/asset dicts for display."""
    umap = {u["id"]: u["username"] for u in auth.list_users(con)}
    for r in rows:
        r["username"] = umap.get(r.get("user_id")) or ("—" if r.get("user_id") else "owner")
    return rows


# -- auth / session ------------------------------------------------------------
@app.post("/api/login")
async def login(payload: dict, request: Request, response: Response):
    ip = request.headers.get("x-forwarded-for", "").split(",")[0].strip() \
        or (request.client.host if request.client else "?")
    if not auth.rate_ok(ip):
        raise HTTPException(429, "too many attempts — wait a minute")
    con = _con(request)
    username = (payload.get("username") or "").strip().lower()
    password = payload.get("password") or ""
    u = auth.get_user_by_name(con, username)
    if not u or not auth.verify_pw(password, u["pw_hash"]):
        auth.note_fail(ip)
        raise HTTPException(401, "invalid username or password")
    token = auth.make_session(con, u["id"])
    secure = request.headers.get("x-forwarded-proto", "http") == "https"
    response.set_cookie(auth.SESSION_COOKIE, token, max_age=auth.SESSION_MAX_AGE,
                        httponly=True, samesite="lax", secure=secure, path="/")
    return auth.public_user(u)


@app.post("/api/logout")
async def logout(response: Response):
    response.delete_cookie(auth.SESSION_COOKIE, path="/")
    return {"ok": True}


@app.get("/api/me")
def me(user=Depends(auth.current_user)):
    return auth.public_user(user)


# -- users (owner only) --------------------------------------------------------
@app.get("/api/users")
def get_users(request: Request, owner=Depends(auth.require_owner)):
    return auth.list_users(_con(request))


@app.post("/api/users")
def add_user(payload: dict, request: Request, owner=Depends(auth.require_owner)):
    try:
        u = auth.create_user(_con(request), payload.get("username", ""),
                             payload.get("password", ""), role=payload.get("role", "guest"),
                             max_active_jobs=int(payload.get("max_active_jobs", 4)))
    except ValueError as e:
        raise HTTPException(400, str(e))
    return auth.public_user(u)


@app.patch("/api/users/{username}")
def patch_user(username: str, payload: dict, request: Request, owner=Depends(auth.require_owner)):
    con = _con(request)
    u = auth.get_user_by_name(con, username)
    if not u:
        raise HTTPException(404, "no such user")
    if payload.get("password"):
        try:
            auth.set_password(con, u["id"], payload["password"])
        except ValueError as e:
            raise HTTPException(400, str(e))
    if "max_active_jobs" in payload:
        con.execute("UPDATE users SET max_active_jobs = ? WHERE id = ?",
                    (int(payload["max_active_jobs"]), u["id"]))
        con.commit()
    return auth.public_user(auth.get_user(con, u["id"]))


@app.delete("/api/users/{username}")
def del_user(username: str, request: Request, owner=Depends(auth.require_owner)):
    con = _con(request)
    u = auth.get_user_by_name(con, username)
    if not u:
        raise HTTPException(404, "no such user")
    if u["role"] == "owner":
        raise HTTPException(400, "cannot delete an owner")
    auth.delete_user(con, u["id"])
    return {"ok": True}


@app.post("/api/tokens")
def make_token(payload: dict, request: Request, owner=Depends(auth.require_owner)):
    # owner mints a token for a user (default: self)
    con = _con(request)
    target = payload.get("username")
    u = auth.get_user_by_name(con, target) if target else owner
    if not u:
        raise HTTPException(404, "no such user")
    raw = auth.create_token(con, u["id"], payload.get("name", ""))
    return {"token": raw, "username": u["username"]}  # shown ONCE


# -- operators -----------------------------------------------------------------
@app.get("/api/operators")
def list_operators(user=Depends(auth.current_user)):
    is_owner = user["role"] == "owner"
    out = []
    for op in runner.operators.values():
        # guests are LOCAL-ONLY: cloud operators are filtered out entirely
        if not is_owner and op.get("requires_env"):
            continue
        out.append({k: v for k, v in op.items() if k != "dir"})
    return out


# -- settings (owner only) -----------------------------------------------------
@app.get("/api/settings")
def get_settings(request: Request, owner=Depends(auth.require_owner)):
    s = settings_mod.get_all(_con(request))
    view = settings_mod.public_view(s)
    view["_secret_keys"] = sorted(settings_mod.SECRET_KEYS)
    view["_env_keys"] = sorted(settings_mod.ENV_MAP.keys())
    view["_env_set"] = sorted(env for key, env in settings_mod.ENV_MAP.items() if s.get(key))
    return view


@app.put("/api/settings")
def put_settings(payload: dict, request: Request, owner=Depends(auth.require_owner)):
    clean = {k: v for k, v in payload.items()
             if not (k in settings_mod.SECRET_KEYS and v and set(str(v)) <= {"•"})}
    settings_mod.set_many(_con(request), clean)
    return get_settings(request, owner)


# -- assets --------------------------------------------------------------------
@app.get("/api/assets")
def list_assets(user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    return _with_usernames(con, store.list_assets(con))


@app.post("/api/assets")
async def upload_asset(request: Request, file: UploadFile = File(...),
                       user=Depends(auth.current_user)):
    cl = request.headers.get("content-length")
    if cl and int(cl) > MAX_UPLOAD:
        raise HTTPException(413, "file too large (1 GB max)")
    data = await file.read()
    if not data:
        raise HTTPException(400, "empty file")
    if len(data) > MAX_UPLOAD:
        raise HTTPException(413, "file too large (1 GB max)")
    return store.register_upload(_con(request), file.filename or "unnamed", data, user_id=user["id"])


@app.delete("/api/assets/{asset_id}")
def delete_asset(asset_id: str, user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    asset = store.get_asset(con, asset_id)
    if not asset:
        raise HTTPException(404, "no such asset")
    if user["role"] != "owner" and asset.get("user_id") != user["id"]:
        raise HTTPException(403, "not your asset")
    store.delete_asset(con, asset_id)
    return {"ok": True}


@app.get("/api/assets/{asset_id}/file")
def asset_file(asset_id: str, member: str | None = None,
               user=Depends(auth.current_user), request: Request = None):
    asset = store.get_asset(_con(request), asset_id)
    if not asset:
        raise HTTPException(404, "no such asset")
    path = Path(asset["path"])
    if member:
        target = (path / member).resolve()
        if not str(target).startswith(str(path.resolve())) or not target.is_file():
            raise HTTPException(404, "no such member")
        return FileResponse(target)
    if path.is_dir():
        files = sorted(p.name for p in path.iterdir() if p.is_file())
        return JSONResponse({"folder": True, "files": files})
    return FileResponse(path, filename=asset["name"])


# -- jobs ----------------------------------------------------------------------
@app.get("/api/jobs")
def list_jobs(user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    return _with_usernames(con, runner.list_jobs(con))


@app.get("/api/jobs/recent")
def recent_jobs(user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    return sysinfo.recent_jobs(con, runner)


@app.post("/api/jobs")
async def create_job(payload: dict, user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    operator = payload.get("operator")
    if operator not in runner.operators:
        raise HTTPException(400, f"unknown operator: {operator}")
    op = runner.operators[operator]
    # HARD RULE: guests run local models only — enforced here regardless of UI
    if op.get("requires_env") and user["role"] != "owner":
        raise HTTPException(403, "guests run local models only")
    if user["role"] != "owner" and runner.active_job_count(con, user["id"]) >= user["max_active_jobs"]:
        raise HTTPException(429, "too many active jobs — wait for some to finish")
    asset_ids = payload.get("asset_ids")
    if not asset_ids:
        asset_ids = [payload["asset_id"]] if payload.get("asset_id") else []
    for aid in asset_ids:
        if not store.get_asset(con, aid):
            raise HTTPException(400, f"unknown asset: {aid}")
    params = payload.get("params") or {}
    job = runner.create_job(con, operator, asset_ids, params, user_id=user["id"])
    await runner.broadcast({"type": "job", "job": job})
    return job


@app.get("/api/jobs/{job_id}")
def get_job(job_id: str, user=Depends(auth.current_user), request: Request = None):
    job = runner.get_job(_con(request), job_id)
    if not job:
        raise HTTPException(404, "no such job")
    return job


def _own_job_or_owner(con, job_id, user):
    job = runner.get_job(con, job_id)
    if not job:
        raise HTTPException(404, "no such job")
    if user["role"] != "owner" and job.get("user_id") != user["id"]:
        raise HTTPException(403, "not your job")
    return job


@app.post("/api/jobs/{job_id}/cancel")
async def cancel_job(job_id: str, user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    _own_job_or_owner(con, job_id, user)
    if not await runner.cancel(con, job_id):
        raise HTTPException(400, "job not cancellable")
    return {"ok": True}


@app.post("/api/jobs/{job_id}/retry")
async def retry_job(job_id: str, user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    job = _own_job_or_owner(con, job_id, user)
    op = runner.operators.get(job["operator"], {})
    if op.get("requires_env") and user["role"] != "owner":
        raise HTTPException(403, "guests run local models only")
    new = runner.create_job(con, job["operator"], job.get("asset_ids") or [],
                            job["params"], user_id=user["id"])
    await runner.broadcast({"type": "job", "job": new})
    return new


@app.delete("/api/jobs/{job_id}")
def delete_job(job_id: str, user=Depends(auth.current_user), request: Request = None):
    con = _con(request)
    _own_job_or_owner(con, job_id, user)
    if not runner.delete_job(con, job_id):
        raise HTTPException(400, "job running or missing")
    return {"ok": True}


# -- system / dashboard --------------------------------------------------------
@app.get("/api/system")
def system(user=Depends(auth.current_user), request: Request = None):
    return sysinfo.snapshot(_con(request), runner)


# -- inbox watch folder --------------------------------------------------------
async def watch_inbox():
    INBOX.mkdir(parents=True, exist_ok=True)
    seen: dict[str, int] = {}
    while True:
        await asyncio.sleep(3)
        try:
            for p in INBOX.iterdir():
                if p.name.startswith(".") or not p.exists():
                    continue
                size = p.stat().st_size if p.is_file() else sum(
                    f.stat().st_size for f in p.rglob("*") if f.is_file())
                if seen.get(str(p)) == size and size > 0:
                    asset = store.register_file(app.state.con, p, move=True,
                                                meta={"source": "inbox"})
                    seen.pop(str(p), None)
                    await runner.broadcast({"type": "assets_changed", "assets": [asset["id"]]})
                else:
                    seen[str(p)] = size
        except Exception as e:
            print(f"[inbox] {e}")


# -- websocket (auth: session cookie or ?token=) -------------------------------
@app.websocket("/ws")
async def ws(websocket: WebSocket):
    con = websocket.app.state.con
    user = None
    tok = websocket.cookies.get(auth.SESSION_COOKIE)
    if tok:
        uid = auth.read_session(con, tok)
        if uid:
            user = auth.get_user(con, uid)
    if not user:
        qtok = websocket.query_params.get("token")
        if qtok:
            user = auth.user_by_token(con, qtok)
    if not user:
        await websocket.close(code=1008)  # policy violation
        return
    await websocket.accept()
    runner.subscribers.add(websocket)
    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        runner.subscribers.discard(websocket)


# -- frontend (built) — public so the SPA can load and show the login screen ----
if WEB_DIST.exists():
    app.mount("/", StaticFiles(directory=WEB_DIST, html=True), name="web")