#!/usr/bin/env python3
"""csroom: work on CS Room assignments from your own machine.

Single file, standard library only. Requires Python 3.10+.

Commands:
    login                          store an API token (create one on the
                                   CS Room website under API tokens)
    classrooms                     list your classes
    assignments <class>            list assignments in a class
    pull <class> <assignment>      download your working copy
    push <class> <assignment>      upload the current directory as your
                                   working copy
    test <class> <assignment>      run the assignment's public tests
    submit <class> <assignment>    submit your work for grading

<class> is a class name or id. Push uploads your working copy only;
run `csroom submit` when you want the teacher to grade it.

Configuration:
    CSROOM_TOKEN   overrides the stored token
    CSROOM_API     overrides the API base URL (default https://api.csroom.org)
    Token and API base are stored in ~/.config/csroom/config (mode 0600).
"""

from __future__ import annotations

import argparse
import getpass
import io
import json
import os
import secrets
import sys
import urllib.error
import urllib.parse
import urllib.request
import zipfile
from pathlib import Path


def _seg(s: str) -> str:
    """Quote one URL path segment (assignment names may hold spaces, # or ?)."""
    return urllib.parse.quote(s, safe="")


DEFAULT_API_BASE = "https://api.csroom.org"
CONFIG_PATH = Path.home() / ".config" / "csroom" / "config"

# Directories never included in a push. Mirrors the server's submission
# snapshot filter, so what you push is what would have been collected anyway.
PUSH_IGNORE_DIRS = {".git", "__pycache__", "node_modules", ".venv", ".deps"}

PUSH_MAX_BYTES = 50 * 1024 * 1024


class CliError(Exception):
    """Fatal, user-facing. Printed without a traceback."""


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------


def read_config() -> dict:
    try:
        return json.loads(CONFIG_PATH.read_text())
    except (OSError, ValueError):
        return {}


def write_config(cfg: dict) -> None:
    CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
    try:
        CONFIG_PATH.parent.chmod(0o700)
    except OSError:
        pass
    # Write then chmod is a window where the file is group-readable on some
    # umasks; create it private from the start instead.
    fd = os.open(CONFIG_PATH, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, "w") as fh:
        json.dump(cfg, fh, indent=2)
        fh.write("\n")
    try:
        CONFIG_PATH.chmod(0o600)
    except OSError:
        pass


def get_token() -> str:
    token = os.environ.get("CSROOM_TOKEN") or read_config().get("token")
    if not token:
        raise CliError("No API token. Run `csroom login` first (or set CSROOM_TOKEN).")
    return token


def get_api_base(args: argparse.Namespace) -> str:
    return (
        getattr(args, "api", None)
        or os.environ.get("CSROOM_API")
        or read_config().get("api_base")
        or DEFAULT_API_BASE
    ).rstrip("/")


# ---------------------------------------------------------------------------
# HTTP
# ---------------------------------------------------------------------------


def request(
    method: str,
    base: str,
    path: str,
    token: str | None = None,
    body: bytes | None = None,
    content_type: str | None = None,
) -> tuple[int, bytes]:
    req = urllib.request.Request(base + path, data=body, method=method)
    if token:
        req.add_header("Authorization", f"Bearer {token}")
    if content_type:
        req.add_header("Content-Type", content_type)
    try:
        with urllib.request.urlopen(req) as resp:
            return resp.status, resp.read()
    except urllib.error.HTTPError as e:
        detail = ""
        try:
            payload = json.loads(e.read().decode("utf-8", "replace"))
            detail = payload.get("detail", "")
            if not isinstance(detail, str):
                detail = json.dumps(detail)
        except Exception:
            pass
        if e.code == 401:
            raise CliError(
                "Not authenticated. Your token may be revoked or expired. "
                "Run `csroom login` with a fresh token."
            )
        raise CliError(f"{e.code} {e.reason}" + (f": {detail}" if detail else ""))
    except urllib.error.URLError as e:
        raise CliError(f"Could not reach {base}: {e.reason}")


def request_json(method: str, base: str, path: str, token: str | None = None,
                 body: dict | None = None) -> dict:
    raw = json.dumps(body).encode() if body is not None else None
    ctype = "application/json" if body is not None else None
    _status, data = request(method, base, path, token, raw, ctype)
    try:
        return json.loads(data.decode("utf-8"))
    except ValueError:
        raise CliError("Server returned an unexpected response.")


# ---------------------------------------------------------------------------
# Class / assignment resolution
# ---------------------------------------------------------------------------


def fetch_classrooms(base: str, token: str) -> list[dict]:
    data = request_json("GET", base, "/api/classrooms/", token)
    rooms: list[dict] = []
    for key, role in (("joined", "student"), ("owner", "instructor")):
        for c in data.get(key, []):
            rooms.append({"id": c["id"], "name": c.get("name", ""), "role": role})
    return rooms


def resolve_classroom(base: str, token: str, ref: str) -> dict:
    rooms = fetch_classrooms(base, token)
    for room in rooms:
        if room["id"] == ref:
            return room
    matches = [r for r in rooms if r["name"].lower() == ref.lower()]
    if len(matches) == 1:
        return matches[0]
    if len(matches) > 1:
        ids = ", ".join(r["id"] for r in matches)
        raise CliError(f'Multiple classes named "{ref}". Use an id: {ids}')
    raise CliError(f'No class "{ref}". Run `csroom classrooms` to see yours.')


# ---------------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------------


def cmd_login(args: argparse.Namespace) -> int:
    base = get_api_base(args)
    token = args.token or os.environ.get("CSROOM_TOKEN")
    if not token:
        token = getpass.getpass("API token (from the CS Room API tokens page): ").strip()
    if not token:
        raise CliError("No token given.")

    me = request_json("GET", base, "/api/auth/me", token)
    # Multi-region: if this account lives on a different region, pin that
    # region's API so pull/push read the same disk the IDE writes.
    relocate = me.get("relocate_to")
    if relocate:
        base = relocate.rstrip("/")
        me = request_json("GET", base, "/api/auth/me", token)

    cfg = read_config()
    cfg["token"] = token
    cfg["api_base"] = base
    write_config(cfg)
    print(f"Signed in as {me.get('email')} ({base})")
    print(f"Token stored in {CONFIG_PATH}")
    return 0


def cmd_classrooms(args: argparse.Namespace) -> int:
    base = get_api_base(args)
    rooms = fetch_classrooms(base, get_token())
    if not rooms:
        print("You are not in any classes yet.")
        return 0
    width = max(len(r["name"]) for r in rooms)
    for r in rooms:
        print(f"{r['name']:<{width}}  {r['role']:<10}  {r['id']}")
    return 0


def cmd_assignments(args: argparse.Namespace) -> int:
    base = get_api_base(args)
    token = get_token()
    room = resolve_classroom(base, token, args.classroom)
    data = request_json("GET", base, "/api/classrooms/assignments", token)
    for c in data.get("classrooms", []):
        if c["id"] == room["id"]:
            templates = c.get("templates", [])
            if not templates:
                break
            for t in templates:
                print(t["name"])
            return 0
    print(f'No assignments in "{room["name"]}" yet.')
    return 0


def _safe_zip_members(zf: zipfile.ZipFile, dest: Path) -> list[zipfile.ZipInfo]:
    """Refuse zip entries that would land outside *dest*."""
    members = []
    for info in zf.infolist():
        name = info.filename
        if name.startswith("/") or "\\" in name or ".." in name.split("/"):
            raise CliError(f"Refusing unsafe path in zip: {name}")
        target = (dest / name).resolve()
        if not target.is_relative_to(dest.resolve()):
            raise CliError(f"Refusing unsafe path in zip: {name}")
        members.append(info)
    return members


def cmd_pull(args: argparse.Namespace) -> int:
    base = get_api_base(args)
    token = get_token()
    room = resolve_classroom(base, token, args.classroom)
    dest = Path(args.dir or args.assignment)
    if dest.exists() and any(dest.iterdir()) and not args.force:
        raise CliError(
            f"{dest} is not empty. Pull overwrites files; pass --force to continue."
        )

    _status, data = request(
        "GET", base,
        f"/api/classrooms/{room['id']}/assignments/{_seg(args.assignment)}/pull",
        token,
    )
    dest.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(io.BytesIO(data)) as zf:
        members = _safe_zip_members(zf, dest)
        zf.extractall(dest, members=members)
        count = sum(1 for m in members if not m.is_dir())
    print(f"Pulled {count} files into {dest}/")
    return 0


def _zip_directory(src: Path) -> bytes:
    buf = io.BytesIO()
    total = 0
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        for root, dirs, files in os.walk(src):
            dirs[:] = [d for d in dirs if d not in PUSH_IGNORE_DIRS
                       and not (Path(root) / d).is_symlink()]
            for fn in files:
                full = Path(root) / fn
                if full.is_symlink() or not full.is_file():
                    continue
                total += full.stat().st_size
                if total > PUSH_MAX_BYTES:
                    raise CliError(
                        f"Directory exceeds {PUSH_MAX_BYTES // (1024 * 1024)} MB; "
                        "remove build artifacts and try again."
                    )
                zf.write(full, str(full.relative_to(src)))
    return buf.getvalue()


def cmd_push(args: argparse.Namespace) -> int:
    base = get_api_base(args)
    token = get_token()
    room = resolve_classroom(base, token, args.classroom)
    src = Path(args.dir or ".").resolve()
    payload = _zip_directory(src)
    if not payload or not zipfile.ZipFile(io.BytesIO(payload)).namelist():
        raise CliError(f"Nothing to push in {src}")

    boundary = "----csroom" + secrets.token_hex(16)
    body = (
        f"--{boundary}\r\n"
        'Content-Disposition: form-data; name="file"; filename="push.zip"\r\n'
        "Content-Type: application/zip\r\n\r\n"
    ).encode() + payload + f"\r\n--{boundary}--\r\n".encode()

    result = json.loads(request(
        "POST", base,
        f"/api/classrooms/{room['id']}/assignments/{_seg(args.assignment)}/push",
        token, body, f"multipart/form-data; boundary={boundary}",
    )[1])
    print(f"Pushed {result.get('files', '?')} files to {room['name']} / {args.assignment}")
    print("This updates your working copy. Run `csroom submit` to turn it in.")
    return 0


def cmd_submit(args: argparse.Namespace) -> int:
    base = get_api_base(args)
    token = get_token()
    room = resolve_classroom(base, token, args.classroom)
    result = request_json(
        "POST", base,
        f"/api/classrooms/{room['id']}/assignments/{_seg(args.assignment)}/submit",
        token, body={},
    )
    files = result.get("files", "?")
    print(f"Submitted {files} files for {args.assignment}.")
    skipped = result.get("skipped") or []
    if skipped:
        print(f"Skipped (too large or symlink): {', '.join(skipped)}")
    return 0


def cmd_test(args: argparse.Namespace) -> int:
    base = get_api_base(args)
    token = get_token()
    room = resolve_classroom(base, token, args.classroom)
    result = request_json(
        "POST", base,
        f"/api/classrooms/{room['id']}/templates/{_seg(args.assignment)}/run-my-tests",
        token, body={},
    )
    output = result.get("output") or ""
    if output:
        print(output.rstrip())
        print()
    print(f"Score: {result.get('passed', 0)}/{result.get('total', 0)} tests passed")
    return 0


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="csroom",
        description="Work on CS Room assignments from your own machine.",
    )
    parser.add_argument(
        "--api",
        help=f"API base URL (default {DEFAULT_API_BASE}, or CSROOM_API)",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    p = sub.add_parser("login", help="store an API token")
    p.add_argument("--token", help="the token; prompted for if omitted")
    p.set_defaults(func=cmd_login)

    p = sub.add_parser("classrooms", help="list your classes")
    p.set_defaults(func=cmd_classrooms)

    p = sub.add_parser("assignments", help="list assignments in a class")
    p.add_argument("classroom", metavar="class")
    p.set_defaults(func=cmd_assignments)

    p = sub.add_parser("pull", help="download your working copy of an assignment")
    p.add_argument("classroom", metavar="class")
    p.add_argument("assignment")
    p.add_argument("--dir", help="target directory (default: ./<assignment>)")
    p.add_argument("--force", action="store_true",
                   help="pull into a non-empty directory, overwriting files")
    p.set_defaults(func=cmd_pull)

    p = sub.add_parser("push", help="upload this directory as your working copy")
    p.add_argument("classroom", metavar="class")
    p.add_argument("assignment")
    p.add_argument("--dir", help="directory to push (default: current directory)")
    p.set_defaults(func=cmd_push)

    p = sub.add_parser("submit", help="submit your working copy for grading")
    p.add_argument("classroom", metavar="class")
    p.add_argument("assignment")
    p.set_defaults(func=cmd_submit)

    p = sub.add_parser("test", help="run the assignment's public tests")
    p.add_argument("classroom", metavar="class")
    p.add_argument("assignment")
    p.set_defaults(func=cmd_test)

    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        return args.func(args)
    except CliError as e:
        print(f"error: {e}", file=sys.stderr)
        return 1
    except KeyboardInterrupt:
        return 130


if __name__ == "__main__":
    sys.exit(main())