#!/usr/bin/env python3
"""JustEvals client: upload eval suites, runs and outputs; read back feedback.

Standard library only, Python 3.9+. Copy this file into your project or run it
directly as a CLI.

    export JUSTEVALS_API_KEY=je_...            # project API key
    export JUSTEVALS_URL=https://justevals.com  # optional

    python justevals.py check
    python justevals.py push manifest.json
    python justevals.py feedback --suite core-v2 [--runs RUN,RUN] [--batch B] [--since ISO]

File references anywhere in a manifest (case inputs, output files) may be
{"path": "local/file.jpg"}, {"url": "https://..."} or {"file_id": "fil_..."},
plus optional "role" and "label". Paths and URLs are uploaded once per project
(deduplicated by SHA-256) and replaced by file IDs.

Manifest format (all top-level keys optional except "suite"):

    {
      "suite": {"slug": "core-v2", "name": "Core v2", "review_config": "score_1_5",
                "cases": [{"key": "O01", "title": "...", "instructions": "...",
                           "expected": "...", "failure_conditions": "...",
                           "inputs": [{"path": "in/O01.jpg", "role": "input"}],
                           "params": {...}, "tags": [...]}]},
      "runs": [{"name": "Sunburst baseline", "model": "...", "batch": "evr_x3m54f",
                "is_baseline": true, "external_id": "evr_x3m54f:baseline",
                "config": {...},
                "outputs": [{"case": "O01", "sample": 0, "status": "ok",
                             "files": [{"path": "out/O01.jpg"}],
                             "cost_usd": 0.09, "latency_ms": 41000,
                             "scores": {"judge_score": 4}}]}],
      "judgments": {"comparisons": [...], "reviews": [...]}
    }
"""

from __future__ import annotations

import argparse
import concurrent.futures
import hashlib
import json
import mimetypes
import os
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any, Iterable

__version__ = "0.2.1"

DEFAULT_URL = "https://justevals.com"
OUTPUT_BATCH = 200
JUDGMENT_BATCH = 500
EXTRA_TYPES = {
    ".webp": "image/webp",
    ".avif": "image/avif",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".png": "image/png",
    ".mp4": "video/mp4",
    ".webm": "video/webm",
    ".json": "application/json",
    ".txt": "text/plain",
}


class JustEvalsError(Exception):
    def __init__(self, status: int, message: str, details: Any = None):
        super().__init__(f"{status}: {message}" + (f" {details}" if details else ""))
        self.status = status
        self.details = details


def _content_type(name: str, fallback: str | None = None) -> str:
    ext = os.path.splitext(urllib.parse.urlparse(name).path)[1].lower()
    return EXTRA_TYPES.get(ext) or mimetypes.guess_type(name)[0] or fallback or "application/octet-stream"


class Client:
    def __init__(
        self,
        api_key: str | None = None,
        base_url: str | None = None,
        timeout: float = 120,
        workers: int = 8,
        base_dir: str | os.PathLike | None = None,
    ):
        self.api_key = api_key or os.environ.get("JUSTEVALS_API_KEY", "")
        if not self.api_key:
            raise ValueError("set JUSTEVALS_API_KEY or pass api_key")
        self.base_url = (base_url or os.environ.get("JUSTEVALS_URL") or DEFAULT_URL).rstrip("/")
        self.timeout = timeout
        self.workers = workers
        self.base_dir = Path(base_dir) if base_dir else Path.cwd()
        self._file_ids: dict[str, str] = {}  # cache: path/url -> file id
        self._lock = threading.Lock()
        self._project_id: str | None = None
        # URLs already uploaded in earlier pushes, so they aren't downloaded again.
        self._url_cache_path = Path(
            os.environ.get("JUSTEVALS_CACHE_DIR") or Path.home() / ".cache" / "justevals"
        ) / "url-files.json"
        self._url_cache: dict[str, str] | None = None

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

    def _request(self, method: str, url: str, body: bytes | None = None,
                 headers: dict[str, str] | None = None, auth: bool = True) -> bytes:
        headers = dict(headers or {})
        headers.setdefault("User-Agent", f"justevals-python/{__version__}")
        if auth:
            headers["Authorization"] = f"Bearer {self.api_key}"
        for attempt in range(4):
            req = urllib.request.Request(url, data=body, method=method, headers=headers)
            try:
                with urllib.request.urlopen(req, timeout=self.timeout) as res:
                    return res.read()
            except urllib.error.HTTPError as err:
                raw = err.read()
                if err.code >= 500 and attempt < 3:
                    time.sleep(1.5 * 2**attempt)
                    continue
                try:
                    data = json.loads(raw)
                    raise JustEvalsError(err.code, data.get("error", ""), data.get("details")) from None
                except ValueError:
                    raise JustEvalsError(err.code, raw.decode(errors="replace")[:500]) from None
            except (urllib.error.URLError, TimeoutError, ConnectionError):
                if attempt < 3:
                    time.sleep(1.5 * 2**attempt)
                    continue
                raise
        raise RuntimeError("unreachable")

    def _api(self, method: str, path: str, payload: Any = None, query: dict | None = None) -> Any:
        url = f"{self.base_url}/api/v1{path}"
        if query:
            q = {k: v for k, v in query.items() if v is not None}
            if q:
                url += "?" + urllib.parse.urlencode(q)
        body = None
        headers = {}
        if payload is not None:
            body = json.dumps(payload, default=str).encode()
            headers["Content-Type"] = "application/json"
        raw = self._request(method, url, body, headers)
        return json.loads(raw) if raw else None

    # -- Files --------------------------------------------------------------

    def upload_bytes(self, data: bytes, content_type: str, filename: str | None = None) -> str:
        """Uploads bytes (once per project) and returns the file ID."""
        sha = hashlib.sha256(data).hexdigest()
        created = self._api("POST", "/files", {
            "sha256": sha, "size": len(data), "content_type": content_type,
            **({"filename": filename[-255:]} if filename else {}),
        })
        upload = created.get("upload")
        if upload:
            self._request(upload["method"], upload["url"], data, upload["headers"], auth=False)
            self._api("POST", f"/files/{created['id']}/complete")
        return created["id"]

    def upload_file(self, path: str | os.PathLike) -> str:
        p = Path(path)
        if not p.is_absolute():
            p = self.base_dir / p
        key = f"path:{p.resolve()}"
        if key not in self._file_ids:
            self._file_ids[key] = self.upload_bytes(p.read_bytes(), _content_type(p.name), p.name)
        return self._file_ids[key]

    def _url_cache_key(self, url: str) -> str:
        with self._lock:
            if self._project_id is None:
                self._project_id = self.project()["id"]
        return f"{self.base_url}|{self._project_id}|{url}"

    def _load_url_cache(self) -> dict[str, str]:
        with self._lock:
            if self._url_cache is None:
                try:
                    self._url_cache = json.loads(self._url_cache_path.read_text())
                except (OSError, ValueError):
                    self._url_cache = {}
            return self._url_cache

    def _remember_url(self, key: str, file_id: str) -> None:
        cache = self._load_url_cache()
        with self._lock:
            cache[key] = file_id
            try:
                self._url_cache_path.parent.mkdir(parents=True, exist_ok=True)
                tmp = self._url_cache_path.with_suffix(f".{os.getpid()}.tmp")
                tmp.write_text(json.dumps(cache))
                tmp.replace(self._url_cache_path)
            except OSError:
                pass  # the cache is an optimisation only

    def upload_url(self, url: str) -> str:
        """Downloads a URL (e.g. an expiring model-provider link) and uploads it.

        URLs uploaded before (to this project) are reused without downloading."""
        key = f"url:{url}"
        if key in self._file_ids:
            return self._file_ids[key]
        cache_key = self._url_cache_key(url)
        cached = self._load_url_cache().get(cache_key)
        if cached:
            try:
                if self._api("GET", f"/files/{cached}").get("status") == "ready":
                    self._file_ids[key] = cached
                    return cached
            except JustEvalsError:
                pass
        req = urllib.request.Request(url, headers={"User-Agent": f"justevals-python/{__version__}"})
        with urllib.request.urlopen(req, timeout=self.timeout) as res:
            data = res.read()
            ctype = (res.headers.get_content_type() or "").lower()
        if ctype in ("", "application/octet-stream", "binary/octet-stream"):
            ctype = _content_type(url)
        name = os.path.basename(urllib.parse.urlparse(url).path) or None
        file_id = self.upload_bytes(data, ctype, name)
        self._file_ids[key] = file_id
        self._remember_url(cache_key, file_id)
        return file_id

    def _resolve_ref(self, ref: dict) -> dict:
        ref = dict(ref)
        if "path" in ref:
            ref["file_id"] = self.upload_file(ref.pop("path"))
        elif "url" in ref:
            ref["file_id"] = self.upload_url(ref.pop("url"))
        return ref

    def _resolve_refs(self, items: Iterable[dict], field: str) -> list[dict]:
        """Uploads all path/url refs in items[field] in parallel."""
        items = [dict(i) for i in items]
        refs = [(i, j) for i, item in enumerate(items) for j, _ in enumerate(item.get(field) or [])]
        with concurrent.futures.ThreadPoolExecutor(self.workers) as pool:
            resolved = list(pool.map(lambda ij: self._resolve_ref(items[ij[0]][field][ij[1]]), refs))
        for item in items:
            if field in item:
                item[field] = list(item[field])
        for (i, j), ref in zip(refs, resolved):
            items[i][field][j] = ref
        return items

    # -- Suites, runs, outputs ----------------------------------------------

    def project(self) -> dict:
        return self._api("GET", "/project")

    def put_suite(self, slug: str, cases: list[dict], **fields: Any) -> dict:
        """Creates or updates a suite from its full case list."""
        cases = self._resolve_refs(cases, "inputs")
        return self._api("PUT", f"/suites/{urllib.parse.quote(slug)}", {**fields, "cases": cases})

    def get_suite(self, slug: str, version: int | None = None) -> dict:
        return self._api("GET", f"/suites/{urllib.parse.quote(slug)}", query={"version": version})

    def create_run(self, suite: str, name: str, **fields: Any) -> dict:
        return self._api("POST", "/runs", {"suite": suite, "name": name, **fields})

    def add_outputs(self, run_id: str, outputs: list[dict]) -> int:
        saved = 0
        for start in range(0, len(outputs), OUTPUT_BATCH):
            batch = self._resolve_refs(outputs[start:start + OUTPUT_BATCH], "files")
            saved += self._api("POST", f"/runs/{urllib.parse.quote(run_id)}/outputs", {"outputs": batch})["saved"]
        return saved

    def complete_run(self, run_id: str, status: str = "complete", notes: str | None = None) -> dict:
        body: dict[str, Any] = {"status": status}
        if notes is not None:
            body["notes"] = notes
        return self._api("POST", f"/runs/{urllib.parse.quote(run_id)}/complete", body)

    def get_run(self, run_id: str) -> dict:
        return self._api("GET", f"/runs/{urllib.parse.quote(run_id)}")

    def list_runs(self, suite: str | None = None, batch: str | None = None) -> list[dict]:
        return self._api("GET", "/runs", query={"suite": suite, "batch": batch})["runs"]

    def import_judgments(self, comparisons: list[dict] | None = None, reviews: list[dict] | None = None) -> dict:
        """Records imported or automatic judgments, in batches under the API's limits."""
        comparisons, reviews = comparisons or [], reviews or []
        total = {"comparisons": 0, "reviews": 0}
        for start in range(0, len(comparisons), JUDGMENT_BATCH):
            n = self._api("POST", "/judgments", {"comparisons": comparisons[start:start + JUDGMENT_BATCH]})
            total["comparisons"] += n["comparisons"]
        for start in range(0, len(reviews), JUDGMENT_BATCH):
            n = self._api("POST", "/judgments", {"reviews": reviews[start:start + JUDGMENT_BATCH]})
            total["reviews"] += n["reviews"]
        return total

    def delete_judgments(self, comparisons: list[dict] | None = None, reviews: list[dict] | None = None) -> dict:
        """Removes imported/automatic judgments: comparisons by outputs + source + reviewer,
        reviews by output + source + reviewer."""
        return self._api("DELETE", "/judgments", {"comparisons": comparisons or [], "reviews": reviews or []})

    def feedback(self, suite: str, runs: list[str] | None = None, batch: str | None = None,
                 since: str | None = None) -> dict:
        return self._api("GET", "/feedback", query={
            "suite": suite, "runs": ",".join(runs) if runs else None, "batch": batch, "since": since,
        })

    def push(self, manifest: dict, log=print) -> dict:
        """Uploads a whole manifest: suite, then runs with outputs, then judgments."""
        suite = dict(manifest["suite"])
        slug = suite.pop("slug")
        cases = suite.pop("cases")
        result = self.put_suite(slug, cases, **suite)
        log(f"suite {slug} v{result['version']} ({result['cases']} cases, {'new version' if result['created'] else 'unchanged'})")
        pushed = []
        for run_spec in manifest.get("runs", []):
            run_spec = dict(run_spec)
            outputs = run_spec.pop("outputs", [])
            status = run_spec.pop("status", "complete")
            run_spec.setdefault("suite_version", result["version"])
            run = self.create_run(slug, **run_spec)
            n = self.add_outputs(run["id"], outputs) if outputs else 0
            if status in ("complete", "failed"):
                self.complete_run(run["id"], status)
            log(f"run {run['name']}: {n} outputs -> {run['url']}")
            pushed.append(run)
        if manifest.get("judgments"):
            counts = self.import_judgments(**manifest["judgments"])
            log(f"judgments: {counts['comparisons']} comparisons, {counts['reviews']} reviews")
        return {"suite": result, "runs": pushed}


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(prog="justevals", description=__doc__.split("\n\n")[0])
    sub = parser.add_subparsers(dest="cmd", required=True)
    sub.add_parser("check", help="verify the API key and show the project")
    p_push = sub.add_parser("push", help="upload a manifest JSON file")
    p_push.add_argument("manifest")
    p_fb = sub.add_parser("feedback", help="print labels and notes as JSON")
    p_fb.add_argument("--suite", required=True)
    p_fb.add_argument("--runs")
    p_fb.add_argument("--batch")
    p_fb.add_argument("--since")
    p_suite = sub.add_parser("suite", help="print a suite version with case labels")
    p_suite.add_argument("slug")
    p_suite.add_argument("--version", type=int)
    args = parser.parse_args(argv)

    try:
        if args.cmd == "check":
            client = Client()
            print(json.dumps({"url": client.base_url, **client.project()}, indent=2))
        elif args.cmd == "push":
            path = Path(args.manifest)
            client = Client(base_dir=path.resolve().parent)
            client.push(json.loads(path.read_text()), log=lambda m: print(m, file=sys.stderr))
        elif args.cmd == "feedback":
            client = Client()
            runs = args.runs.split(",") if args.runs else None
            print(json.dumps(client.feedback(args.suite, runs, args.batch, args.since), indent=2, ensure_ascii=False))
        elif args.cmd == "suite":
            print(json.dumps(Client().get_suite(args.slug, args.version), indent=2, ensure_ascii=False))
    except JustEvalsError as err:
        print(f"error {err}", file=sys.stderr)
        return 1
    return 0


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