backend/main.py

"""
TreyOS Backend — FastAPI
Real Linux filesystem, WebSocket terminal, system stats, Playwright browser
"""

import asyncio
import json
import os
import shutil
import stat
import subprocess
import time
from datetime import datetime
from pathlib import Path

import psutil
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from watchfiles import awatch

app = FastAPI(title="TreyOS", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Path handling ─────────────────────────────────────────────────────────────
# TREY_ROOT jails access. Default "/" means full filesystem access.
# The jail check is intentionally lenient — on HF the whole container is fair game.
ROOT_JAIL = Path(os.environ.get("TREY_ROOT", "/")).resolve()

def safe_path(requested: str) -> Path:
    """Resolve path. If absolute, use as-is (within jail). If relative, join to ROOT_JAIL."""
    requested = requested.strip()
    if requested.startswith("/"):
        p = Path(requested).resolve()
    else:
        p = (ROOT_JAIL / requested).resolve()
    # Jail check
    try:
        p.relative_to(ROOT_JAIL)
    except ValueError:
        raise HTTPException(403, f"Access denied: {p} is outside jail {ROOT_JAIL}")
    return p

def entry_info(p: Path) -> dict:
    try:
        s = p.stat()
        return {
            "name": p.name,
            "path": str(p),
            "type": "dir" if p.is_dir() else "file",
            "size": s.st_size,
            "modified": datetime.fromtimestamp(s.st_mtime).isoformat(),
            "permissions": oct(stat.S_IMODE(s.st_mode)),
            "extension": p.suffix.lower() if p.is_file() else None,
        }
    except (PermissionError, OSError):
        return {"name": p.name, "path": str(p), "type": "unknown",
                "size": 0, "modified": None, "permissions": None, "extension": None}


# ── Filesystem API ────────────────────────────────────────────────────────────

@app.get("/api/files")
async def list_files(path: str = "/", show_hidden: bool = False):
    p = safe_path(path)
    if not p.exists():
        raise HTTPException(404, f"Path not found: {path}")
    if not p.is_dir():
        raise HTTPException(400, "Not a directory")
    try:
        entries = sorted(p.iterdir(), key=lambda x: (not x.is_dir(), x.name.lower()))
        if not show_hidden:
            entries = [e for e in entries if not e.name.startswith(".")]
        return {
            "path": str(p),
            "parent": str(p.parent) if str(p) != "/" else None,
            "entries": [entry_info(e) for e in entries],
            "count": len(entries),
        }
    except PermissionError:
        raise HTTPException(403, "Permission denied")


@app.get("/api/files/read")
async def read_file(path: str):
    p = safe_path(path)
    if not p.is_file():
        raise HTTPException(404, "File not found")
    try:
        return {"path": str(p), "content": p.read_text(errors="replace"), "size": p.stat().st_size}
    except Exception as e:
        raise HTTPException(500, str(e))


@app.post("/api/files/write")
async def write_file(path: str, content: str = ""):
    p = safe_path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(content)
    return {"ok": True, "path": str(p), "size": p.stat().st_size}


@app.post("/api/files/mkdir")
async def make_dir(path: str):
    p = safe_path(path)
    p.mkdir(parents=True, exist_ok=True)
    return {"ok": True, "path": str(p)}


@app.delete("/api/files/delete")
async def delete_entry(path: str):
    p = safe_path(path)
    if not p.exists():
        raise HTTPException(404, "Not found")
    if p.is_dir():
        shutil.rmtree(p)
    else:
        p.unlink()
    return {"ok": True, "deleted": str(p)}


@app.post("/api/files/rename")
async def rename_entry(path: str, new_name: str):
    p = safe_path(path)
    dest = p.parent / new_name
    p.rename(dest)
    return {"ok": True, "path": str(dest)}


@app.post("/api/files/copy")
async def copy_entry(src: str, dst: str):
    s, d = safe_path(src), safe_path(dst)
    shutil.copytree(s, d) if s.is_dir() else shutil.copy2(s, d)
    return {"ok": True, "destination": str(d)}


@app.post("/api/files/move")
async def move_entry(src: str, dst: str):
    s, d = safe_path(src), safe_path(dst)
    shutil.move(str(s), str(d))
    return {"ok": True, "destination": str(d)}


@app.post("/api/files/upload")
async def upload_file(path: str, file: UploadFile = File(...)):
    dest = safe_path(path) / file.filename
    dest.parent.mkdir(parents=True, exist_ok=True)
    with open(dest, "wb") as f:
        f.write(await file.read())
    return {"ok": True, "path": str(dest), "size": dest.stat().st_size}


@app.get("/api/files/download")
async def download_file(path: str):
    p = safe_path(path)
    if not p.is_file():
        raise HTTPException(404, "File not found")
    return FileResponse(str(p), filename=p.name)


@app.get("/api/files/search")
async def search_files(path: str = "/", query: str = "", limit: int = 50):
    p = safe_path(path)
    results = []
    try:
        for entry in p.rglob(f"*{query}*"):
            results.append(entry_info(entry))
            if len(results) >= limit:
                break
    except (PermissionError, OSError):
        pass
    return {"query": query, "results": results, "count": len(results)}


# ── System Stats ──────────────────────────────────────────────────────────────

@app.get("/api/system")
async def system_stats():
    disk = psutil.disk_usage("/")
    net  = psutil.net_io_counters()
    uptime_s = int(time.time() - psutil.boot_time())
    h, r = divmod(uptime_s, 3600)
    m = r // 60
    return {
        "cpu":    {"percent": psutil.cpu_percent(interval=0.1), "cores": psutil.cpu_count(),
                   "freq_mhz": psutil.cpu_freq().current if psutil.cpu_freq() else None},
        "memory": {"total_gb": round(psutil.virtual_memory().total/1e9,1),
                   "used_gb":  round(psutil.virtual_memory().used/1e9,1),
                   "percent":  psutil.virtual_memory().percent},
        "disk":   {"total_gb": round(disk.total/1e9,1), "used_gb": round(disk.used/1e9,1),
                   "free_gb":  round(disk.free/1e9,1),  "percent": disk.percent},
        "network":{"bytes_sent_mb": round(net.bytes_sent/1e6,1), "bytes_recv_mb": round(net.bytes_recv/1e6,1)},
        "uptime": f"{h}h {m}m",
        "boot_time": datetime.fromtimestamp(psutil.boot_time()).isoformat(),
        "hostname": os.uname().nodename,
        "os": f"{os.uname().sysname} {os.uname().release}",
    }


@app.get("/api/processes")
async def list_processes(limit: int = 20):
    procs = []
    for p in psutil.process_iter(["pid","name","cpu_percent","memory_percent","status"]):
        try: procs.append(p.info)
        except (psutil.NoSuchProcess, psutil.AccessDenied): pass
    procs.sort(key=lambda x: x.get("cpu_percent", 0), reverse=True)
    return {"processes": procs[:limit]}


# ── WebSocket Terminal ────────────────────────────────────────────────────────

class TerminalSession:
    def __init__(self, cwd: str = "/"):
        self.cwd = cwd if Path(cwd).is_dir() else "/"
        self.env = {**os.environ, "TERM": "xterm-256color"}

    async def run(self, command: str, ws: WebSocket):
        cmd = command.strip()
        # Handle cd internally
        if cmd.startswith("cd"):
            target = cmd[2:].strip() or str(Path.home())
            new_cwd = str((Path(self.cwd) / target).resolve())
            if os.path.isdir(new_cwd):
                self.cwd = new_cwd
            else:
                await ws.send_json({"type":"output","data":f"cd: {target}: No such directory\r\n"})
            await ws.send_json({"type":"prompt","cwd":self.cwd})
            return
        try:
            proc = await asyncio.create_subprocess_shell(
                cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.STDOUT,
                cwd=self.cwd, env=self.env,
            )
            async for chunk in proc.stdout:
                await ws.send_json({"type":"output","data":chunk.decode(errors="replace")})
            await proc.wait()
            await ws.send_json({"type":"exit","code":proc.returncode,"cwd":self.cwd})
        except Exception as e:
            await ws.send_json({"type":"error","data":str(e)})
        finally:
            await ws.send_json({"type":"prompt","cwd":self.cwd})


@app.websocket("/ws/terminal")
async def terminal_ws(ws: WebSocket, cwd: str = "/"):
    await ws.accept()
    # Use ROOT_JAIL as default cwd if it exists, else /
    default_cwd = str(ROOT_JAIL) if ROOT_JAIL.is_dir() else "/"
    session = TerminalSession(cwd=cwd if Path(cwd).is_dir() else default_cwd)
    await ws.send_json({"type":"welcome","message":"TreyOS Terminal Ready","cwd":session.cwd})
    await ws.send_json({"type":"prompt","cwd":session.cwd})
    try:
        while True:
            msg = await ws.receive_json()
            if msg.get("type") == "command":
                await session.run(msg["data"], ws)
            elif msg.get("type") == "ping":
                await ws.send_json({"type":"pong"})
    except WebSocketDisconnect:
        pass


@app.websocket("/ws/watch")
async def watch_ws(ws: WebSocket, path: str = "/"):
    await ws.accept()
    p = safe_path(path)
    await ws.send_json({"type":"watching","path":str(p)})
    try:
        async for changes in awatch(str(p)):
            events = [{"change": str(c[0].name), "path": c[1]} for c in changes]
            await ws.send_json({"type":"changes","events":events})
    except WebSocketDisconnect:
        pass


@app.websocket("/ws/system")
async def system_ws(ws: WebSocket, interval: float = 2.0):
    await ws.accept()
    try:
        while True:
            stats = await system_stats()
            await ws.send_json({"type":"stats","data":stats})
            await asyncio.sleep(max(interval, 1.0))
    except WebSocketDisconnect:
        pass


# ── Playwright Browser Proxy ──────────────────────────────────────────────────
# Streams a real Chromium screenshot as JPEG for each navigation.
# Install: pip install playwright && playwright install chromium

_browser = None
_playwright = None

async def get_browser():
    global _browser, _playwright
    if _browser is None:
        try:
            from playwright.async_api import async_playwright
            _playwright = await async_playwright().start()
            _browser = await _playwright.chromium.launch(
                headless=True,
                args=["--no-sandbox","--disable-setuid-sandbox","--disable-dev-shm-usage",
                      "--disable-gpu","--single-process"]
            )
        except Exception as e:
            raise HTTPException(503, f"Playwright not available: {e}")
    return _browser


@app.get("/api/browser/screenshot")
async def browser_screenshot(url: str, width: int = 390, height: int = 700):
    """Take a screenshot of a URL and return it as JPEG."""
    if not url.startswith(("http://","https://")):
        url = "https://" + url
    try:
        browser = await get_browser()
        page = await browser.new_page(viewport={"width": width, "height": height})
        await page.goto(url, wait_until="domcontentloaded", timeout=20000)
        screenshot = await page.screenshot(type="jpeg", quality=80, full_page=False)
        await page.close()
        from fastapi.responses import Response
        return Response(content=screenshot, media_type="image/jpeg",
                        headers={"X-Final-URL": page.url if not page.is_closed() else url})
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(502, f"Browser error: {e}")


@app.websocket("/ws/browser")
async def browser_ws(ws: WebSocket):
    """
    WebSocket browser session.
    Client sends: {"type":"navigate","url":"https://..."}
    Server sends: {"type":"screenshot","data":"<base64 jpeg>","url":"...","title":"..."}
    """
    await ws.accept()
    await ws.send_json({"type":"status","message":"Browser session ready"})
    page = None
    try:
        browser = await get_browser()
        page = await browser.new_page(viewport={"width":390,"height":700})

        async def send_screenshot(p):
            try:
                shot = await p.screenshot(type="jpeg", quality=75)
                import base64
                await ws.send_json({
                    "type": "screenshot",
                    "data": base64.b64encode(shot).decode(),
                    "url":  p.url,
                    "title": await p.title(),
                })
            except Exception as e:
                await ws.send_json({"type":"error","message":str(e)})

        while True:
            msg = await ws.receive_json()
            if msg.get("type") == "navigate":
                url = msg["url"]
                if not url.startswith(("http://","https://")):
                    url = "https://" + url
                await ws.send_json({"type":"loading","url":url})
                try:
                    await page.goto(url, wait_until="domcontentloaded", timeout=20000)
                    await send_screenshot(page)
                except Exception as e:
                    await ws.send_json({"type":"error","message":str(e)})
            elif msg.get("type") == "click":
                await page.mouse.click(msg.get("x",0), msg.get("y",0))
                await asyncio.sleep(0.5)
                await send_screenshot(page)
            elif msg.get("type") == "scroll":
                await page.mouse.wheel(0, msg.get("delta",200))
                await asyncio.sleep(0.2)
                await send_screenshot(page)
            elif msg.get("type") == "back":
                await page.go_back()
                await send_screenshot(page)
            elif msg.get("type") == "forward":
                await page.go_forward()
                await send_screenshot(page)
            elif msg.get("type") == "refresh":
                await page.reload()
                await send_screenshot(page)
    except WebSocketDisconnect:
        pass
    finally:
        if page and not page.is_closed():
            await page.close()


# ── Docker ────────────────────────────────────────────────────────────────────

def run_docker(*args) -> dict:
    try:
        r = subprocess.run(["docker",*args], capture_output=True, text=True, timeout=10)
        return {"ok": r.returncode==0, "stdout": r.stdout, "stderr": r.stderr}
    except FileNotFoundError:
        return {"ok":False,"error":"Docker not installed"}
    except subprocess.TimeoutExpired:
        return {"ok":False,"error":"Timed out"}


@app.get("/api/docker/containers")
async def docker_containers():
    r = run_docker("ps","-a","--format","json")
    containers = []
    for line in r.get("stdout","").strip().splitlines():
        try: containers.append(json.loads(line))
        except: pass
    return {"containers": containers, "error": r.get("error") or r.get("stderr") if not r["ok"] else None}


@app.post("/api/docker/start/{cid}")
async def docker_start(cid: str): return run_docker("start", cid)

@app.post("/api/docker/stop/{cid}")
async def docker_stop(cid: str): return run_docker("stop", cid)

@app.get("/api/docker/images")
async def docker_images():
    r = run_docker("images","--format","json")
    images = []
    for line in r.get("stdout","").strip().splitlines():
        try: images.append(json.loads(line))
        except: pass
    return {"images": images}


# ── Health ────────────────────────────────────────────────────────────────────

@app.get("/api/health")
async def health():
    return {"status":"ok","name":"TreyOS","version":"1.0.0","root":str(ROOT_JAIL)}


# ── Serve frontend ────────────────────────────────────────────────────────────
# HuggingFace Spaces loads the root URL "/" directly in an iframe, so the
# TreyOS UI must be served there — not just at /app — or the Space will
# render blank.

frontend_path = Path(__file__).parent / "frontend"
if not frontend_path.exists():
    frontend_path = Path(__file__).parent.parent / "frontend"

if frontend_path.exists():
    index_file = frontend_path / "index.html"

    @app.get("/")
    async def serve_root():
        return FileResponse(str(index_file), media_type="text/html")

    @app.get("/app")
    async def serve_app():
        return FileResponse(str(index_file), media_type="text/html")

    app.mount("/static", StaticFiles(directory=str(frontend_path)), name="static")
else:
    @app.get("/")
    async def serve_root_fallback():
        return JSONResponse({"status":"ok","name":"TreyOS","warning":"frontend/ directory not found"})