CSOJ-cli.py

#!/usr/bin/env python3

import asyncio
import json
import os
import sys
from pathlib import Path
import base64
import fnmatch

import click
import requests
import websockets
from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeRemainingColumn
from rich.table import Table
from rich.syntax import Syntax
from rich.text import Text
from rich.markdown import Markdown

# --- Configuration ---
CONFIG_DIR = Path.home() / ".config" / "CSOJ-cli"
CONFIG_FILE = CONFIG_DIR / "config.json"
console = Console()

# --- Command Aliases ---
COMMAND_ALIASES = {
    "lb": "leaderboard",
    "lsc": "ls-contests",
    "lsp": "ls-problems",
    "lss": "ls-submissions",
    "reg": "register",
    "show": "show-problem",
    "st": "status",
    "sub": "submit",
}

class AliasedGroup(click.Group):
    """A click Group subclass that supports command aliases."""
    def get_command(self, ctx, cmd_name):
        # Try the standard command name first
        rv = click.Group.get_command(self, ctx, cmd_name)
        if rv is not None:
            return rv
        # If that fails, check if it's an alias
        if cmd_name in COMMAND_ALIASES:
            actual_cmd = COMMAND_ALIASES[cmd_name]
            return click.Group.get_command(self, ctx, actual_cmd)
        return None

# --- Helper Functions ---
def ensure_config_dir():
    """Ensures the configuration directory exists."""
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)


def save_config(domain, jwt):
    """Saves domain and JWT to the config file."""
    ensure_config_dir()
    with open(CONFIG_FILE, "w") as f:
        json.dump({"domain": domain.rstrip('/'), "jwt": jwt}, f)
    CONFIG_FILE.chmod(0o600)


def load_config():
    """Loads domain and JWT from the config file."""
    if not CONFIG_FILE.is_file():
        console.print(
            "[bold red]Error:[/bold red] You are not logged in. "
            "Please run '[bold cyan]CSOJ-cli login <domain> <jwt>[/bold cyan]' first."
        )
        sys.exit(1)
    try:
        with open(CONFIG_FILE, "r") as f:
            return json.load(f)
    except (json.JSONDecodeError, IOError) as e:
        console.print(f"[bold red]Error reading config file {CONFIG_FILE}:[/bold red] {e}")
        sys.exit(1)


def get_api_headers(cookie=None):
    """Gets the necessary headers for authenticated API calls."""
    config = load_config()
    headers = {"Authorization": f"Bearer {config['jwt']}"}
    if cookie:
        headers['Cookie'] = cookie
    return headers


def handle_api_error(response):
    """Prints a formatted error from an API response and exits."""
    try:
        data = response.json()
        message = data.get("message", "An unknown error occurred.")
    except json.JSONDecodeError:
        message = response.text
    console.print(f"[bold red]API Error (HTTP {response.status_code}):[/bold red] {message}")
    sys.exit(1)


# --- CLI Command Group ---
@click.group(cls=AliasedGroup)
@click.option('--cookie', help='Optional cookie string to include in all requests.')
@click.pass_context
def cli(ctx, cookie):
    """
    CSOJ-cli: A command-line tool for the CSOJ Online Judge platform.
    
    Supports both full command names and short aliases (e.g., 'leaderboard' or 'lb').
    The --cookie option can be used with any command.
    """
    ctx.ensure_object(dict)
    ctx.obj['cookie'] = cookie


# --- CLI Commands ---
@cli.command()
@click.argument("domain")
@click.argument("jwt")
def login(domain, jwt):
    """
    Log in to a CSOJ instance and save credentials.

    DOMAIN: The base URL of the CSOJ instance (e.g., https://oj.example.com)
    JWT: Your JSON Web Token from the CSOJ web interface.
    """
    save_config(domain, jwt)
    console.print(f"[bold green]✓ Successfully logged in to {domain}[/bold green]")
    console.print(f"Configuration saved to [cyan]{CONFIG_FILE}[/cyan]")


@cli.command("ls-contests", help="List all available contests. (alias: lsc)")
@click.pass_context
def list_contests(ctx):
    config = load_config()
    domain = config["domain"]
    cookie = ctx.obj.get('cookie')
    headers = get_api_headers(cookie)
    url = f"{domain}/api/v1/contests"

    try:
        with console.status("[bold green]Fetching contests...[/]"):
            response = requests.get(url, headers=headers, timeout=10)
        
        if response.status_code != 200:
            handle_api_error(response)
        
        contests = response.json().get("data", {})
        if not contests:
            console.print("[yellow]No contests found.[/yellow]")
            return

        table = Table(title="[bold]Available Contests[/bold]", title_style="default", header_style="bold magenta")
        table.add_column("ID", style="cyan", no_wrap=True)
        table.add_column("Name", style="bold")
        table.add_column("Start Time", style="green")
        table.add_column("End Time", style="yellow")
        
        sorted_contests = sorted(contests.values(), key=lambda x: x.get('starttime', ''))

        for contest in sorted_contests:
            table.add_row(
                contest.get("id", "N/A"),
                contest.get("name", "N/A"),
                contest.get("starttime", "N/A").replace("T", " ").split(".")[0],
                contest.get("endtime", "N/A").replace("T", " ").split(".")[0],
            )
        console.print(table)

    except requests.exceptions.RequestException as e:
        console.print(f"[bold red]Network Error:[/bold red] Could not connect to {domain}. {e}")
        sys.exit(1)


@cli.command("ls-problems", help="List all problems for a given contest. (alias: lsp)")
@click.argument("contest_id")
@click.pass_context
def list_problems(ctx, contest_id):
    config = load_config()
    domain = config["domain"]
    cookie = ctx.obj.get('cookie')
    headers = get_api_headers(cookie)
    contest_url = f"{domain}/api/v1/contests/{contest_id}"

    try:
        with console.status("[bold green]Fetching contest details...[/]"):
            contest_response = requests.get(contest_url, headers=headers, timeout=10)
        
        if contest_response.status_code != 200:
            handle_api_error(contest_response)
        
        contest_data = contest_response.json().get("data", {})
        problem_ids = contest_data.get("problem_ids", [])

        if not problem_ids:
            console.print(f"No problems found for contest '[bold cyan]{contest_id}[/bold cyan]'. The contest might not have started yet.")
            return

        table = Table(title=f"[bold]Problems in Contest: {contest_data.get('name', contest_id)}[/bold]", header_style="bold magenta")
        table.add_column("ID", style="cyan", no_wrap=True)
        table.add_column("Name", style="bold")

        progress = Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            TimeRemainingColumn(),
            console=console
        )
        with progress:
            task = progress.add_task("[green]Fetching problems...", total=len(problem_ids))
            for prob_id in problem_ids:
                problem_url = f"{domain}/api/v1/problems/{prob_id}"
                problem_response = requests.get(problem_url, headers=headers, timeout=10)
                if problem_response.status_code == 200:
                    problem_data = problem_response.json().get("data", {})
                    table.add_row(problem_data.get("id", "N/A"), problem_data.get("name", "N/A"))
                else:
                     table.add_row(prob_id, "[red]Error fetching name[/red]")
                progress.update(task, advance=1)
        console.print(table)

    except requests.exceptions.RequestException as e:
        console.print(f"[bold red]Network Error:[/bold red] {e}")
        sys.exit(1)


@cli.command(help="Register for a specific contest. (alias: reg)")
@click.argument("contest_id")
@click.pass_context
def register(ctx, contest_id):
    config = load_config()
    domain = config["domain"]
    cookie = ctx.obj.get('cookie')
    headers = get_api_headers(cookie)
    url = f"{domain}/api/v1/contests/{contest_id}/register"

    try:
        with console.status(f"[bold green]Registering for contest [cyan]{contest_id}[/cyan]...[/]"):
            response = requests.post(url, headers=headers, timeout=15)

        if response.status_code == 200:
            console.print(f"[bold green]✓ Successfully registered for contest '{contest_id}'![/bold green]")
        elif response.status_code == 409:
            console.print(f"[bold yellow]✓ You are already registered for contest '{contest_id}'.[/bold yellow]")
        else:
            handle_api_error(response)

    except requests.exceptions.RequestException as e:
        console.print(f"[bold red]Network Error:[/bold red] Could not connect to {domain}. {e}")
        sys.exit(1)


@cli.command(help="Submit a file or folder to a problem. (alias: sub)")
@click.argument("problem_id")
@click.argument("path", type=click.Path(exists=True, resolve_path=True, file_okay=True, dir_okay=True))
@click.pass_context
def submit(ctx, problem_id, path):
    config = load_config()
    domain = config["domain"]
    cookie = ctx.obj.get('cookie')
    headers = get_api_headers(cookie)
    submit_url = f"{domain}/api/v1/problems/{problem_id}/submit"
    problem_url = f"{domain}/api/v1/problems/{problem_id}"
    submit_path = Path(path)
    skipped_files = []
    files_to_upload_info = []
    prepared_files_for_request = []

    try:
        # Step 1: Fetch problem details to get upload rules
        with console.status("[bold green]Fetching problem details for upload rules...[/]"):
            problem_response = requests.get(problem_url, headers=headers, timeout=10)
        
        if problem_response.status_code != 200:
            console.print(f"[bold red]Error fetching problem details for '{problem_id}'. Cannot verify file rules.[/bold red]")
            handle_api_error(problem_response)

        problem_data = problem_response.json().get("data", {})
        upload_rules = problem_data.get("upload", {}).get("upload_files", [])

        console.print(f"Packaging submission from [cyan]{submit_path}[/cyan]...")
        
        # Step 2: Determine files to check and filter them based on rules
        all_files_to_check = []
        if submit_path.is_dir():
            all_files_to_check = [p for p in submit_path.rglob('*') if p.is_file()]
        else:
            all_files_to_check = [submit_path]

        for file_path in all_files_to_check:
            if submit_path.is_dir():
                relative_path = file_path.relative_to(submit_path.parent)
            else:
                relative_path = Path(file_path.name)
            
            # Default to allowed if no rules are specified on the server
            is_allowed = not upload_rules
            if upload_rules:
                # Check if the file matches any of the allowed patterns
                is_allowed = any(fnmatch.fnmatch(str(relative_path), pattern) for pattern in upload_rules)

            if is_allowed:
                files_to_upload_info.append({'absolute_path': file_path, 'relative_path': relative_path})
            else:
                skipped_files.append(str(relative_path))

        # Step 3: Handle skipped files and abort if no valid files remain
        if skipped_files:
            console.print("[bold yellow]Warning:[/bold yellow] The following files were skipped as they do not match the allowed patterns:")
            for skipped in skipped_files:
                console.print(f"  - [dim]{skipped}[/dim]")

        if not files_to_upload_info:
            console.print("[bold red]Error:[/bold red] No valid files found for submission after filtering. Aborting.")
            sys.exit(1)

        # Step 4: Prepare files for upload by opening them
        for file_info in files_to_upload_info:
            file_obj = open(file_info['absolute_path'], 'rb')
            b64_name = base64.b64encode(file_info['relative_path'].as_posix().encode()).decode()
            prepared_files_for_request.append(('files', (b64_name, file_obj)))
        
        console.print(f"Uploading {len(prepared_files_for_request)} file(s) to problem [cyan]{problem_id}[/cyan]...")
        response = requests.post(submit_url, headers=headers, files=prepared_files_for_request, timeout=60)

        if response.status_code == 403:
            try:
                error_data = response.json()
                if "banned" in error_data.get("message", "").lower():
                    handle_api_error(response)
            except json.JSONDecodeError:
                pass
            
            console.print("[bold red]Submission Forbidden (HTTP 403).[/bold red]")
            console.print("This may be because you have not registered for the contest.")
            console.print("Please run '[bold cyan]CSOJ-cli reg <contest_id>[/bold cyan]' first.")
            sys.exit(1)
        
        if response.status_code != 200:
            handle_api_error(response)

        data = response.json().get("data", {})
        submission_id = data.get("submission_id")
        console.print(f"[bold green]✓ Submission successful![/bold green]")
        console.print(f"  Submission ID: [bold yellow]{submission_id}[/bold yellow]")
        console.print(f"  Run '[bold cyan]CSOJ-cli status {submission_id} --logs[/bold cyan]' to see the progress.")

    except requests.exceptions.RequestException as e:
        console.print(f"[bold red]Network Error:[/bold red] {e}")
        sys.exit(1)
    except Exception as e:
        console.print(f"[bold red]An unexpected error occurred:[/bold red] {e}")
        sys.exit(1)
    finally:
        # Step 5: Ensure all opened file objects are closed
        for _, (_, file_obj) in prepared_files_for_request:
            if file_obj:
                file_obj.close()


async def stream_logs(domain, jwt, submission_id, container_id, step_name, cookie=None):
    """Coroutine to connect to WebSocket and stream logs."""
    ws_protocol = "wss" if domain.startswith("https") else "ws"
    http_protocol = "https" if ws_protocol == "wss" else "http"
    domain_host = domain.replace(f"{http_protocol}://", "")
    uri = f"{ws_protocol}://{domain_host}/api/v1/ws/submissions/{submission_id}/containers/{container_id}/logs?token={jwt}"
    
    panel_title = f"[bold]Logs for Step: [blue]{step_name}[/blue] (Container: {container_id[:12]})[/bold]"
    
    connecting_text = Text(f"Connecting to {uri}", style="dim")
    connecting_text.overflow = "fold"
    console.print(Panel(connecting_text, title=panel_title, border_style="green", title_align="left"))

    ws_headers = {}
    if cookie:
        ws_headers['Cookie'] = cookie

    try:
        async with websockets.connect(uri, extra_headers=ws_headers) as websocket:
            async for message in websocket:
                try:
                    log_data = json.loads(message)
                    stream = log_data.get("stream")
                    data = log_data.get("data", "")
                    
                    if stream == "info" and "Log stream finished" in data:
                        continue

                    color = "default"
                    if stream == "stderr": color = "bright_red"
                    elif stream == "info": color = "bright_blue"
                    
                    for line in data.rstrip().splitlines():
                       console.print(f"[{color}]{line}[/{color}]")

                except json.JSONDecodeError:
                    console.print(f"[yellow]Received non-JSON message:[/yellow] {message}")
    except websockets.exceptions.ConnectionClosed as e:
        reason_text = Text(f"Log stream finished. Reason: {e.reason or 'Connection closed'} (Code: {e.code})", style="yellow")
        console.print(Panel(reason_text, title=panel_title, border_style="yellow", title_align="left"))
    except Exception as e:
        console.print(Panel(f"[bold red]WebSocket Error:[/bold red] {e}", title=panel_title, border_style="red", title_align="left"))


@cli.command(help="Get the status and optionally logs of a submission. (alias: st)")
@click.argument("submission_id")
@click.option("--logs", "-l", is_flag=True, help="Stream logs for all workflow steps.")
@click.pass_context
def status(ctx, submission_id, logs):
    config = load_config()
    domain = config["domain"]
    cookie = ctx.obj.get('cookie')
    headers = get_api_headers(cookie)
    sub_url = f"{domain}/api/v1/submissions/{submission_id}"

    try:
        with console.status("[bold green]Fetching submission status...[/]"):
            response = requests.get(sub_url, headers=headers, timeout=10)
        
        if response.status_code != 200:
            handle_api_error(response)

        sub_data = response.json().get("data", {})
        
        grid = Table.grid(expand=True, padding=(0, 1))
        grid.add_column(justify="right", style="dim cyan", width=15)
        grid.add_column(justify="left")
        
        status_val = sub_data.get("status", "Unknown")
        status_color = {"Queued": "yellow", "Running": "blue", "Success": "green", "Failed": "red"}.get(status_val, "white")
        
        grid.add_row("Submission ID:", f"[bold yellow]{sub_data.get('id')}[/bold yellow]")
        grid.add_row("Problem ID:", sub_data.get('problem_id'))
        grid.add_row("Status:", f"[bold {status_color}]{status_val}[/bold {status_color}]")
        grid.add_row("Score:", str(sub_data.get('score')))
        grid.add_row("Submitted At:", sub_data.get("CreatedAt", "").replace("T", " ").split(".")[0])
        
        info = sub_data.get('info')
        if info:
            info_str = json.dumps(info, indent=2)
            grid.add_row("Info:", Syntax(info_str, "json", theme="default", word_wrap=True))

        console.print(Panel(grid, title="[bold]Submission Status[/bold]", expand=False, border_style="magenta"))
        
        if logs:
            problem_id = sub_data.get('problem_id')
            prob_url = f"{domain}/api/v1/problems/{problem_id}"
            prob_response = requests.get(prob_url, headers=headers)
            workflow = prob_response.json().get("data", {}).get("workflow", []) if prob_response.ok else []

            containers = sub_data.get("containers", [])
            if not containers:
                console.print("[yellow]No containers found for this submission yet.[/yellow]")
                return

            sorted_containers = sorted(containers, key=lambda x: x.get('CreatedAt', ''))

            for i, container in enumerate(sorted_containers):
                step_name = workflow[i].get("name", f"Step {i+1}") if i < len(workflow) else f"Step {i+1}"
                can_show = workflow[i].get("show", False) if i < len(workflow) else False
                
                if not can_show:
                    console.print(Panel(f"[dim]Logs for step '[blue]{step_name}[/blue]' are hidden by the problem setter.[/dim]", border_style="dim"))
                    continue

                asyncio.run(stream_logs(domain, config['jwt'], submission_id, container.get("id"), step_name, cookie))

    except requests.exceptions.RequestException as e:
        console.print(f"[bold red]Network Error:[/bold red] {e}")
        sys.exit(1)


@cli.command(help="Display the leaderboard for a contest. (alias: lb)")
@click.argument("contest_id")
@click.pass_context
def leaderboard(ctx, contest_id):
    config = load_config()
    domain = config["domain"]
    cookie = ctx.obj.get('cookie')
    headers = get_api_headers(cookie)
    
    try:
        with console.status("[bold green]Fetching contest details...[/]"):
            contest_url = f"{domain}/api/v1/contests/{contest_id}"
            contest_res = requests.get(contest_url, headers=headers, timeout=10)
            if contest_res.status_code != 200:
                handle_api_error(contest_res)
            contest_data = contest_res.json().get("data", {})
            problem_ids = sorted(contest_data.get("problem_ids", []))
            contest_name = contest_data.get("name", contest_id)

        with console.status(f"[bold green]Fetching leaderboard for {contest_name}...[/]"):
            leaderboard_url = f"{domain}/api/v1/contests/{contest_id}/leaderboard"
            leaderboard_res = requests.get(leaderboard_url, headers=headers, timeout=15)
            if leaderboard_res.status_code != 200:
                handle_api_error(leaderboard_res)
            leaderboard_data = leaderboard_res.json().get("data", [])

        if not leaderboard_data:
            console.print(f"[yellow]Leaderboard for contest '{contest_name}' is currently empty.[/yellow]")
            return
            
        table = Table(title=f"[bold]Leaderboard: {contest_name}[/bold]", header_style="bold magenta")
        table.add_column("Rank", justify="right", style="bold")
        table.add_column("User", style="cyan")
        table.add_column("Total Score", justify="right", style="bold green")
        
        for prob_id in problem_ids:
            table.add_column(prob_id, justify="right")
        
        for rank, entry in enumerate(leaderboard_data, 1):
            row = [
                str(rank),
                entry.get('nickname') or entry.get('username', 'N/A'),
                str(entry.get('total_score', 0))
            ]
            problem_scores = entry.get('problem_scores', {})
            for prob_id in problem_ids:
                score = problem_scores.get(prob_id, "-")
                row.append(str(score))
            table.add_row(*row)
            
        console.print(table)

    except requests.exceptions.RequestException as e:
        console.print(f"[bold red]Network Error:[/bold red] {e}")
        sys.exit(1)


@cli.command("show-problem", help="Show the description and details of a problem. (alias: show)")
@click.argument("problem_id")
@click.pass_context
def show_problem(ctx, problem_id):
    config = load_config()
    domain = config["domain"]
    cookie = ctx.obj.get('cookie')
    headers = get_api_headers(cookie)
    url = f"{domain}/api/v1/problems/{problem_id}"
    
    try:
        with console.status(f"[bold green]Fetching problem {problem_id}...[/]"):
            response = requests.get(url, headers=headers, timeout=10)
        
        if response.status_code != 200:
            handle_api_error(response)
            
        problem = response.json().get("data", {})
        
        info_grid = Table.grid(expand=True, padding=(0, 1))
        info_grid.add_column(width=20, style="dim cyan", justify="right")
        info_grid.add_column()
        
        info_grid.add_row("Problem Name:", f"[bold]{problem.get('name')}[/bold]")
        info_grid.add_row("Problem ID:", problem.get('id'))
        info_grid.add_row("Max Submissions:", str(problem.get('max_submissions', 'Unlimited')))
        
        console.print(Panel(info_grid, title="[bold]Problem Details[/bold]", border_style="magenta", expand=False))
        console.print(Panel(Markdown(problem.get('description', 'No description available.')), title="[bold]Description[/bold]"))
        
    except requests.exceptions.RequestException as e:
        console.print(f"[bold red]Network Error:[/bold red] {e}")
        sys.exit(1)


@cli.command("ls-submissions", help="List your recent submissions. (alias: lss)")
@click.pass_context
def list_submissions(ctx):
    config = load_config()
    domain = config["domain"]
    cookie = ctx.obj.get('cookie')
    headers = get_api_headers(cookie)
    url = f"{domain}/api/v1/submissions"

    try:
        with console.status("[bold green]Fetching your submissions...[/]"):
            response = requests.get(url, headers=headers, timeout=15)
        
        if response.status_code != 200:
            handle_api_error(response)
        
        submissions = response.json().get("data", [])
        if not submissions:
            console.print("[yellow]You have not made any submissions yet.[/yellow]")
            return

        table = Table(title="[bold]Your Submissions[/bold]", header_style="bold magenta")
        table.add_column("ID", style="yellow")
        table.add_column("Problem ID", style="cyan")
        table.add_column("Status")
        table.add_column("Score", justify="right")
        table.add_column("Timestamp", style="dim")

        for sub in submissions:
            status_val = sub.get("status", "Unknown")
            status_color = {"Queued": "yellow", "Running": "blue", "Success": "green", "Failed": "red"}.get(status_val, "white")
            score = sub.get('score', 0)
            score_color = "green" if score > 0 else "default"

            table.add_row(
                sub.get('id'),
                sub.get('problem_id'),
                f"[{status_color}]{status_val}[/{status_color}]",
                f"[{score_color}]{score}[/{score_color}]",
                sub.get("CreatedAt", "").replace("T", " ").split(".")[0]
            )
        console.print(table)

    except requests.exceptions.RequestException as e:
        console.print(f"[bold red]Network Error:[/bold red] {e}")
        sys.exit(1)

# --- Main Execution ---
if __name__ == "__main__":
    cli()