"""OpenCode Gateway & File Explorer for Hugging Face Space.

Features:
- Pure Full-Screen UI: Proxies directly to native OpenCode web UI.
- VS Code Style File Explorer API:
  * GET /__api/tree: Lists folders & files recursively with metadata.
  * GET /__api/file: Previews text/code content.
- Direct Download Endpoints:
  * GET /dl/zip: One-tap ZIP archive download of the entire workspace.
  * GET /dl/{path}: Direct file download with Content-Disposition attachment.
- Strict Security Gatekeeper:
  * Unauthenticated requests see a locked authentication challenge.
  * Accepts valid dynamic 8-digit PIN and/or Master Authorization Token.
  * Internal loopback Basic Auth automated injection (eliminates 401 Basic Realm prompt).
  * Frame Embedding: Strips X-Frame-Options & restrictive CSP.
- Full WebSocket & Streaming Support.
- Bucket persistence to kacapower/ghar-storage.
"""

import asyncio
import base64
import hashlib
import hmac
import json
import os
import time
import zipfile
from datetime import datetime, timezone, timedelta
from io import BytesIO
from pathlib import Path
from typing import Optional

import httpx
from fastapi import FastAPI, HTTPException, Header, Query, Request, Response, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, StreamingResponse
from starlette.background import BackgroundTask

try:
    import websockets
except ImportError:
    websockets = None

UPSTREAM = os.getenv("OPENCODE_URL", "http://127.0.0.1:4096").rstrip("/")
UPSTREAM_WS = UPSTREAM.replace("http://", "ws://").replace("https://", "wss://")

HOME = Path(os.getenv("HOME", "/home/node"))
ROOTS = {
    "workspace": Path(os.getenv("WORKSPACE_DIR", HOME / "workspace")),
    "sessions": Path(os.getenv("OPENCODE_DATA_DIR", HOME / ".local/share/opencode")) / "storage" / "session",
    "config": Path(os.getenv("OPENCODE_CONFIG_DIR", HOME / ".config/opencode")),
}

# Internal Loopback Basic Auth credentials
INTERNAL_USER = os.getenv("OPENCODE_SERVER_USERNAME", "opencode")
INTERNAL_PASS = os.getenv("OPENCODE_SERVER_PASSWORD", "loopback_internal_pass_777")
INTERNAL_BASIC_AUTH = "Basic " + base64.b64encode(f"{INTERNAL_USER}:{INTERNAL_PASS}".encode()).decode()

# Security Configuration
UNLOCK_SECRET = os.getenv("UNLOCK_SECRET") or os.getenv("VAULT_SECRET", "Ghar777").strip()
ACCESS_PIN = os.getenv("ACCESS_PIN", "").strip()
HF_READ_TOKEN = os.getenv("HF_READ_TOKEN", "").strip() or os.getenv("HF_TOKEN", "").strip()
DIRECT_SESSION_TIMEOUT = int(os.getenv("DIRECT_SESSION_TIMEOUT", "600"))

TEXT_LIMIT = 512 * 1024
ZIP_LIMIT = 500 * 1024 * 1024

state = {
    "last_direct_session": 0.0,
    "direct_mode_pause_external": os.getenv("DIRECT_MODE_PAUSE", "1") in {"1", "true", "True"},
}

HOP_HEADERS = {
    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
    "te", "trailer", "trailers", "transfer-encoding", "upgrade", "host",
    "content-length",
}
STRIP_HEADERS = {
    "x-frame-options",
    "content-security-policy",
    "content-security-policy-report-only",
    "www-authenticate",
    "date",
}

app = FastAPI(title="OpenCode Gateway", docs_url=None, redoc_url=None)

# Enable CORS for external frontends
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

client = httpx.AsyncClient(timeout=None, follow_redirects=False)


def compute_daily_pin(day_str: str) -> str:
    h = hmac.new(UNLOCK_SECRET.encode("utf-8"), day_str.encode("utf-8"), hashlib.sha256).hexdigest()
    return f"{int(h[:8], 16) % 100000000:08d}"


def get_valid_pins() -> set:
    now_utc = datetime.now(timezone.utc)
    today_str = now_utc.strftime("%Y-%m-%d")
    yesterday_str = (now_utc - timedelta(days=1)).strftime("%Y-%m-%d")

    pins = {compute_daily_pin(today_str), compute_daily_pin(yesterday_str)}
    if ACCESS_PIN:
        pins.add(ACCESS_PIN)
    return pins


def is_pin_valid(pin_val: Optional[str]) -> bool:
    if not (UNLOCK_SECRET or ACCESS_PIN):
        return True
    return bool(pin_val and pin_val.strip() in get_valid_pins())


def is_auth_valid(request: Request) -> bool:
    # 1. Check Query parameter
    pin = request.query_params.get("pin")
    token = request.query_params.get("token")
    if is_pin_valid(pin):
        return True
    if HF_READ_TOKEN and token == HF_READ_TOKEN:
        return True

    # 2. Check Cookie
    cookie_pin = request.cookies.get("opencode_auth")
    if is_pin_valid(cookie_pin):
        return True

    # 3. Check Headers
    header_pin = request.headers.get("X-Access-Pin")
    if is_pin_valid(header_pin):
        return True

    auth = request.headers.get("Authorization", "")
    if auth.lower().startswith("bearer "):
        bearer = auth[7:].strip()
        if HF_READ_TOKEN and bearer == HF_READ_TOKEN:
            return True

    return False


def get_lock_html(error_msg: str = "") -> str:
    err_html = f'<div style="color:#ef4444;font-size:13px;margin-bottom:12px;">{error_msg}</div>' if error_msg else ''
    return f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>OpenCode Access Locked</title>
<style>
body {{ background:#0d1117; color:#c9d1d9; font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,sans-serif; display:flex; align-items:center; justify-content:center; height:100vh; margin:0; }}
.card {{ background:#161b22; border:1px solid #30363d; border-radius:12px; padding:32px; width:90%; max-width:360px; text-align:center; box-shadow:0 8px 24px rgba(0,0,0,0.5); }}
h2 {{ color:#f0f6fc; font-size:18px; margin-top:0; }}
p {{ color:#8b949e; font-size:13px; margin-bottom:20px; }}
input {{ width:100%; box-sizing:border-box; background:#0d1117; border:1px solid #30363d; border-radius:8px; color:white; font-size:22px; letter-spacing:4px; text-align:center; padding:10px; margin-bottom:16px; outline:none; }}
input:focus {{ border-color:#58a6ff; }}
button {{ width:100%; background:#238636; color:white; border:none; padding:10px; border-radius:8px; font-weight:600; cursor:pointer; font-size:14px; }}
button:hover {{ background:#2ea043; }}
</style>
</head>
<body>
<div class="card">
<div style="font-size:36px;margin-bottom:12px;">🔒</div>
<h2>OpenCode Access Locked</h2>
<p>Authentication Required.</p>
{err_html}
<form method="GET">
<input type="password" name="pin" maxlength="8" pattern="[0-9]*" inputmode="numeric" placeholder="••••••••" autofocus required>
<button type="submit">Unlock</button>
</form>
</div>
</body>
</html>"""


# ---------------- PATH HELPERS ----------------

def _root(root: str = "workspace") -> Path:
    if root not in ROOTS:
        raise HTTPException(status_code=404, detail=f"unknown root {root!r}")
    return ROOTS[root].resolve()


def _resolve(root: str, rel: str) -> Path:
    base = _root(root)
    rel = rel.lstrip("/")
    target = (base / rel).resolve() if rel else base
    if target != base and base not in target.parents:
        raise HTTPException(status_code=403, detail="path escapes root")
    return target


def _rel(base: Path, target: Path) -> str:
    try:
        rel = target.relative_to(base).as_posix()
    except ValueError:
        return ""
    return "" if rel == "." else rel


def _entry(base: Path, p: Path) -> dict:
    try:
        st = p.stat()
    except OSError:
        st = None
    is_dir = p.is_dir()
    return {
        "name": p.name,
        "path": _rel(base, p),
        "type": "dir" if is_dir else "file",
        "size": 0 if is_dir or st is None else st.st_size,
        "mtime": 0 if st is None else int(st.st_mtime),
    }


# ---------------- HEALTH ENDPOINT ----------------

@app.get("/__health")
async def health(request: Request):
    ok = True
    try:
        async with httpx.AsyncClient(timeout=2.0) as c:
            await c.get(UPSTREAM + "/", headers={"authorization": INTERNAL_BASIC_AUTH})
    except Exception:
        ok = False

    now_utc = datetime.now(timezone.utc)
    return {
        "ok": ok,
        "upstream": UPSTREAM,
        "cycle_date": now_utc.strftime("%Y-%m-%d"),
    }


# ---------------- VS CODE FILE EXPLORER APIS ----------------

@app.get("/__api/tree")
async def api_tree(request: Request, path: str = ""):
    if not is_auth_valid(request):
        raise HTTPException(status_code=401, detail="Unauthorized")

    base = _root("workspace")
    target = _resolve("workspace", path)
    if not target.exists():
        raise HTTPException(status_code=404, detail="not found")

    entries = []
    try:
        children = list(target.iterdir())
    except OSError as exc:
        raise HTTPException(status_code=403, detail=str(exc))

    for child in children:
        # Hide internal git and cache folders from VS Code explorer view
        if child.name in {".git", "__pycache__", "node_modules"}:
            continue
        if not child.is_symlink() or child.exists():
            entries.append(_entry(base, child))

    entries.sort(key=lambda e: (e["type"] != "dir", e["name"].lower()))
    return {"path": _rel(base, target), "entries": entries}


@app.get("/__api/file")
async def api_file(request: Request, path: str = ""):
    if not is_auth_valid(request):
        raise HTTPException(status_code=401, detail="Unauthorized")

    target = _resolve("workspace", path)
    if not target.is_file():
        raise HTTPException(status_code=404, detail="not a file")

    size = target.stat().st_size
    with target.open("rb") as fh:
        head = fh.read(4096)
    binary = b"\x00" in head

    payload = {
        "name": target.name,
        "path": _rel(_root("workspace"), target),
        "size": size,
        "binary": binary,
        "content": None,
    }
    if not binary:
        with target.open("rb") as fh:
            raw = fh.read(TEXT_LIMIT + 1)
        payload["content"] = raw[:TEXT_LIMIT].decode("utf-8", errors="replace")
    return payload


# ---------------- DIRECT DOWNLOAD ROUTES ----------------

@app.get("/dl/zip")
async def download_workspace_zip(request: Request):
    if not is_auth_valid(request):
        raise HTTPException(status_code=401, detail="Unauthorized")

    base = _root("workspace")
    buf = BytesIO()
    total = 0
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        for child in sorted(base.rglob("*")):
            if child.is_file() and not any(part in {".git", "__pycache__", "node_modules"} for part in child.parts):
                total += child.stat().st_size
                if total > ZIP_LIMIT:
                    raise HTTPException(status_code=413, detail="Workspace archive exceeds size limit")
                zf.write(child, child.relative_to(base).as_posix())

    buf.seek(0)
    filename = f"workspace_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip"
    headers = {"Content-Disposition": f'attachment; filename="{filename}"'}
    return StreamingResponse(buf, media_type="application/zip", headers=headers)


@app.get("/dl/{path:path}")
async def direct_file_download(path: str, request: Request):
    if not is_auth_valid(request):
        raise HTTPException(status_code=401, detail="Unauthorized")

    target = _resolve("workspace", path)
    if not target.is_file():
        raise HTTPException(status_code=404, detail="File not found")

    return FileResponse(target, filename=target.name, media_type="application/octet-stream")


# ---------------- PROXY TO OPENCODE ----------------

def _proxy_headers(request: Request) -> dict:
    headers = {k: v for k, v in request.headers.items() if k.lower() not in HOP_HEADERS}
    headers["authorization"] = INTERNAL_BASIC_AUTH
    return headers


def _response_headers(upstream: httpx.Response) -> dict:
    return {
        k: v
        for k, v in upstream.headers.items()
        if k.lower() not in HOP_HEADERS and k.lower() not in STRIP_HEADERS
    }


@app.api_route(
    "/{path:path}",
    methods=["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"],
)
async def proxy(path: str, request: Request):
    if not is_auth_valid(request):
        if request.method == "GET" and "text/html" in request.headers.get("accept", ""):
            err = "Invalid PIN" if request.query_params.get("pin") else ""
            return HTMLResponse(get_lock_html(err), status_code=401 if err else 200)
        return JSONResponse({"error": "Unauthorized: 8-digit PIN or valid authorization required"}, status_code=401)

    params = [(k, v) for k, v in request.query_params.multi_items() if k not in {"pin", "token", "iframe"}]
    url = httpx.URL(f"{UPSTREAM}/{path}", params=params) if params else f"{UPSTREAM}/{path}"

    try:
        upstream_request = client.build_request(
            request.method,
            url,
            headers=_proxy_headers(request),
            content=request.stream(),
        )
        upstream = await client.send(upstream_request, stream=True)
    except httpx.HTTPError as exc:
        return JSONResponse(
            {"error": "opencode upstream unreachable", "detail": str(exc)},
            status_code=502,
        )

    resp = StreamingResponse(
        upstream.aiter_raw(),
        status_code=upstream.status_code,
        headers=_response_headers(upstream),
        background=BackgroundTask(upstream.aclose),
    )

    pin_val = request.query_params.get("pin") or request.cookies.get("opencode_auth")
    if pin_val and is_pin_valid(pin_val):
        resp.set_cookie("opencode_auth", pin_val, max_age=86400, httponly=True, samesite="none", secure=True)

    return resp


@app.websocket("/{path:path}")
async def ws_proxy(websocket: WebSocket, path: str):
    if websockets is None:
        await websocket.close(code=1011)
        return

    pin_val = websocket.query_params.get("pin") or websocket.cookies.get("opencode_auth")
    token_val = websocket.query_params.get("token")
    auth_ok = is_pin_valid(pin_val) or (HF_READ_TOKEN and token_val == HF_READ_TOKEN)

    if not auth_ok:
        await websocket.close(code=4401)
        return

    protocol = websocket.headers.get("sec-websocket-protocol")
    await websocket.accept(subprotocol=protocol)

    query = "&".join(f"{k}={v}" for k, v in websocket.query_params.items() if k not in {"pin", "token"})
    target = f"{UPSTREAM_WS}/{path}" + (f"?{query}" if query else "")
    subprotocols = [p.strip() for p in protocol.split(",")] if protocol else None
    extra_headers = {"Authorization": INTERNAL_BASIC_AUTH}

    try:
        async with websockets.connect(target, subprotocols=subprotocols, extra_headers=extra_headers, max_size=None) as upstream:

            async def pump_upstream():
                try:
                    async for message in upstream:
                        if isinstance(message, bytes):
                            await websocket.send_bytes(message)
                        else:
                            await websocket.send_text(message)
                except Exception:
                    pass

            task = asyncio.create_task(pump_upstream())
            try:
                while True:
                    message = await websocket.receive()
                    if message["type"] == "websocket.disconnect":
                        break
                    if message.get("bytes") is not None:
                        await upstream.send(message["bytes"])
                    elif message.get("text") is not None:
                        await upstream.send(message["text"])
            except WebSocketDisconnect:
                pass
            finally:
                task.cancel()
    except Exception:
        try:
            await websocket.close()
        except Exception:
            pass