#!/usr/bin/env python3
"""
Send prompts / code-review requests to your local OpenAI-compatible gateway
(vLLM behind FastAPI gateway). No curl needed.

Examples:
   python3 review.py --base-url http://10.21.18.46:1234/big --key pg116-SECRET     --review ein6-prog3-27813/link/set_link.c --review ein6-prog3-27813/link/set_link.h

  # code review of a file
  python3 review.py --base-url http://10.21.18.46:1234/fast --key pg116-SECRET \
    --review main.c --lang c

  # review multiple files (repeat --review)
  python3 review.py --base-url http://10.21.18.46:1234/fast --key pg116-SECRET \
    --review main.c --review include/util.h

Notes:
- This targets /v1/chat/completions.
- It prints only the assistant message by default; use --raw to print full JSON.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional

import requests


def read_text(path: Path, max_bytes: Optional[int] = None) -> str:
    data = path.read_bytes()
    if max_bytes is not None and len(data) > max_bytes:
        data = data[:max_bytes]
    text= data.decode("utf-8", errors="replace")
    lines = text.splitlines(keepends=True)
    numbered_lines = [
        f"{i+1}: {line}" for i, line in enumerate(lines)
    ]

    return "".join(numbered_lines)

def build_review_prompt(files: List[Path], lang: str) -> str:
    parts: List[str] = []
    parts.append(
        f"You are a senior C code reviewer.\n"
    "Be concise and do it step by step.\nData structures must not be modified.\n"
    "Find real problems only:\n"
    "1- non-linear complexity\n"
    "2- memory errors (buffer overflow, memory leaks)\n"
    "3- undefined behavior\n"
    "4- style issue, error-prone formulations\n"
    "\n"
    "OUTPUT FORMAT (plain text only, no markdown):\n"
    "For each issue:\n"
    "FILE:function:line — problem description.\n"
    "Fix: description of the fix.\n"
    "\n"
    "Rules:\n"
    "- No blank lines.\n"
    "- No code blocks.\n"
    )
    parts.append("\n--- CODE ---\n")
    for p in files:
        parts.append(f"\n### File: {p.name}\n")
        parts.append("```" + lang + "\n")
        parts.append(read_text(p))
        parts.append("\n```\n")
    return "".join(parts)


def post_chat_completion(
    base_url: str,
    api_key: str,
    model: str,
    messages: List[Dict[str, str]],
    max_tokens: int,
    timeout: int,
) -> Dict[str, Any]:
    url = base_url.rstrip("/") + "/v1/chat/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }
    payload = {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
    }
    r = requests.post(url, headers=headers, data=json.dumps(payload), timeout=timeout)
    # helpful error printing
    if r.status_code != 200:
        raise RuntimeError(f"HTTP {r.status_code}: {r.text}")
    return r.json()


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--base-url",
        required=True,
        help="Gateway base URL, e.g. http://10.21.18.78:1234/fast (no /v1)",
    )
    ap.add_argument("--key", default=os.environ.get("LAB_API_KEY", ""), help="Bearer token (or set LAB_API_KEY).")
    ap.add_argument("--model", default="", help='Model name for vLLM (default: "inferred from base-url").')
    ap.add_argument("--max-tokens", type=int, default=512, help="Requested max_tokens (gateway may clamp).")
    ap.add_argument("--timeout", type=int, default=600, help="HTTP timeout in seconds.")
    ap.add_argument("--raw", action="store_true", help="Print full JSON response.")

    mode = ap.add_mutually_exclusive_group(required=True)
    mode.add_argument("--prompt", help="Send a basic prompt.")
    mode.add_argument("--review", action="append", default=[], help="Path to a file to review (repeatable).")

    ap.add_argument("--lang", default="c", help="Language label for code fences (default: c).")
    args = ap.parse_args()

    if not args.key:
        print("Error: missing --key (or set LAB_API_KEY).", file=sys.stderr)
        return 2

    if args.prompt:
        messages = [{"role": "user", "content": args.prompt}]
    else:
        files = [Path(p) for p in args.review]
        for p in files:
            if not p.exists() or not p.is_file():
                print(f"Error: file not found: {p}", file=sys.stderr)
                return 2
        review_text = build_review_prompt(files, args.lang)
        messages = [
            {
                "role": "system",
                "content": "You are a strict and helpful code reviewer. Be technical and actionable.",
            },
            {"role": "user", "content": review_text},
        ]
    if not args.model:
        # infer from base-url suffix (/fast or /big)
        if args.base_url.rstrip("/").endswith("/big"):
           args.model = "big"
        else:
           args.model = "fast"
    resp = post_chat_completion(
        base_url=args.base_url,
        api_key=args.key,
        model=args.model,
        messages=messages,
        max_tokens=args.max_tokens,
        timeout=args.timeout,
    )

    if args.raw:
        print(json.dumps(resp, indent=2, ensure_ascii=False))
        return 0

    try:
        text = resp["choices"][0]["message"]["content"]
    except Exception:
        print(json.dumps(resp, indent=2, ensure_ascii=False))
        return 0

    print(text)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())