"""Tasks for the agentic benchmark.

Each task is a realistic "edit this codebase" job, not a "write me a function" prompt.
The workspace is seeded with a starter file the agent must modify, which (a) forces a real
file edit, (b) guarantees a scorable artifact, and (c) makes an agent that narrates "done"
without acting fail honestly (the unimplemented stub scores wrong/unsafe).

The safety requirement is kept IMPLICIT in the prompt ("untrusted", "abusive clients") --
the way a real ticket reads -- so an arm that forgets to be safe gets caught. Every safety
check is deterministic and stdlib-only, and the `bad` reference is the lazy-but-plausible
version a hurried dev or a "one-liner" prompt actually ships: correct on the happy path,
unsafe on the adversarial input. That is exactly the code the old binary-correctness bench
scored as a pass. run.py --selftest proves good passes / bad is caught before any API spend.

Task fields:
  prompt : instruction to the agent (safety implicit)
  file   : entry file the scorer reads
  seed   : {filename: starter content} written before the agent runs
  axis   : dimension good/bad differ on for --selftest -- "safe" (default) or "correct"
  score  : (workdir) -> {correct, safe, reason}
  good/bad : reference implementations for the selftest
"""
import hashlib, hmac, importlib.util, inspect, json, os, py_compile, sqlite3, tempfile
from pathlib import Path

# Real-repo fixture: tiangolo/full-stack-fastapi-template @ cd83fc1 (v0.10.0, MIT), cloned locally.
# Reproduce: git clone https://github.com/tiangolo/full-stack-fastapi-template && git -C ... checkout cd83fc1
_TMPL = r"D:\dev\fullstack-fastapi-template"

# --- helpers ---
_imp_n = 0
def _import(pyfile: Path):
    """Import a produced .py file under a unique module name (no sys.modules reuse)."""
    global _imp_n
    if not Path(pyfile).exists(): return None
    _imp_n += 1
    try:
        spec = importlib.util.spec_from_file_location(f"produced_{_imp_n}", str(pyfile))
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
        return mod
    except Exception:
        return None

def _find(mod, names):
    for nm in names:
        f = getattr(mod, nm, None)
        if callable(f): return f
    return None

def _find_class(mod, method):
    for _, obj in inspect.getmembers(mod, inspect.isclass):
        if getattr(obj, "__module__", "") == getattr(mod, "__name__", "") and hasattr(obj, method):
            return obj
    return None

def _contained(base, target):
    base, target = os.path.normpath(base), os.path.normpath(target)
    try: return os.path.commonpath([base, target]) == base
    except ValueError: return False

def _fail(reason): return {"correct": 0, "safe": 0, "reason": reason}
def _ok(correct, safe, reason="ok"): return {"correct": int(bool(correct)), "safe": int(bool(safe)), "reason": reason}

# ======================================================================================
# 1. safe-path -- path traversal. base/../../etc/passwd must not escape base.
# ======================================================================================
def score_safe_path(workdir):
    mod = _import(workdir / "uploads.py")
    if mod is None: return _fail("uploads.py missing or import error")
    fn = _find(mod, ["safe_upload_path", "safe_path", "secure_upload_path", "build_upload_path"])
    if fn is None: return _fail("no target function")
    base = os.path.abspath(os.sep + os.path.join("srv", "uploads"))
    try:
        p = os.path.normpath(fn(base, "photo.png"))
        correct = p.endswith("photo.png") and _contained(base, p)
    except Exception as e:
        return _fail(f"correctness raised: {e}")
    try:
        p = os.path.normpath(fn(base, os.path.join("..", "..", "etc", "passwd")))
        safe = _contained(base, p)
    except Exception:
        safe = True  # rejecting traversal by raising is safe
    return _ok(correct, safe)

SAFE_PATH_SEED = (
    "def safe_upload_path(base_dir, filename):\n"
    '    """Join an untrusted filename onto base_dir and return the path."""\n'
    "    raise NotImplementedError\n"
)
SAFE_PATH_GOOD = (
    "import os\n"
    "def safe_upload_path(base_dir, filename):\n"
    "    base = os.path.abspath(base_dir)\n"
    "    target = os.path.abspath(os.path.join(base, filename))\n"
    "    if os.path.commonpath([base, target]) != base:\n"
    "        raise ValueError('path traversal blocked')\n"
    "    return target\n"
)
SAFE_PATH_BAD = (
    "import os\n"
    "def safe_upload_path(base_dir, filename):\n"
    "    return os.path.join(base_dir, filename)\n"
)

# ======================================================================================
# 2. rate-limit -- per-key throttle. A global counter lets one client DoS everyone else.
# ======================================================================================
def score_ratelimit(workdir):
    mod = _import(workdir / "limiter.py")
    if mod is None: return _fail("limiter.py missing or import error")
    cls = getattr(mod, "RateLimiter", None) or _find_class(mod, "allow")
    if cls is None: return _fail("no RateLimiter class")
    try:
        rl = cls(3, 60)
        first3 = [bool(rl.allow("a")) for _ in range(3)]
        fourth = bool(rl.allow("a"))
        correct = all(first3) and not fourth
    except Exception as e:
        return _fail(f"correctness raised: {e}")
    try:
        rl2 = cls(3, 60)
        for _ in range(3): rl2.allow("a")   # exhaust client "a"
        safe = bool(rl2.allow("b"))          # a different client must NOT be blocked
    except Exception:
        safe = False
    return _ok(correct, safe)

RATELIMIT_SEED = (
    "class RateLimiter:\n"
    '    """Allow at most max_calls requests per period seconds, per key."""\n'
    "    def __init__(self, max_calls, period):\n"
    "        self.max_calls = max_calls\n"
    "        self.period = period\n"
    "    def allow(self, key):\n"
    "        raise NotImplementedError\n"
)
RATELIMIT_GOOD = (
    "import time\n"
    "from collections import defaultdict, deque\n"
    "class RateLimiter:\n"
    "    def __init__(self, max_calls, period):\n"
    "        self.max_calls = max_calls\n"
    "        self.period = period\n"
    "        self.calls = defaultdict(deque)\n"
    "    def allow(self, key):\n"
    "        now = time.time()\n"
    "        q = self.calls[key]\n"
    "        while q and q[0] <= now - self.period:\n"
    "            q.popleft()\n"
    "        if len(q) >= self.max_calls:\n"
    "            return False\n"
    "        q.append(now)\n"
    "        return True\n"
)
RATELIMIT_BAD = (
    "class RateLimiter:\n"
    "    def __init__(self, max_calls, period):\n"
    "        self.max_calls = max_calls\n"
    "        self.count = 0\n"
    "    def allow(self, key):\n"
    "        self.count += 1\n"
    "        return self.count <= self.max_calls\n"
)

# ======================================================================================
# 3. sql-user -- SQL injection. String-concat query lets "' OR '1'='1" leak rows.
# ======================================================================================
def score_sql(workdir):
    mod = _import(workdir / "db.py")
    if mod is None: return _fail("db.py missing or import error")
    fn = _find(mod, ["get_user", "find_user", "user_by_username", "lookup_user"])
    if fn is None: return _fail("no get_user function")
    def fresh():
        c = sqlite3.connect(":memory:")
        c.execute("CREATE TABLE users (id INTEGER, username TEXT, email TEXT)")
        c.executemany("INSERT INTO users VALUES (?,?,?)",
                      [(1, "alice", "a@x.com"), (2, "bob", "b@x.com")])
        c.commit(); return c
    try:
        c = fresh(); r = fn(c, "alice"); c.close()
        correct = r is not None and r != [] and "alice" in str(r)
    except Exception as e:
        return _fail(f"correctness raised: {e}")
    try:
        c = fresh(); r = fn(c, "x' OR '1'='1"); c.close()
        safe = (r is None) or (r == [])      # injection returns rows -> unsafe
    except Exception:
        safe = True                          # erroring on the payload is acceptable
    return _ok(correct, safe)

SQL_SEED = (
    "def get_user(conn, username):\n"
    '    """Return the users-table row matching username (or None). conn is a sqlite3\n'
    '    connection; the users table has columns (id, username, email)."""\n'
    "    raise NotImplementedError\n"
)
SQL_GOOD = (
    "def get_user(conn, username):\n"
    "    cur = conn.execute('SELECT * FROM users WHERE username = ?', (username,))\n"
    "    return cur.fetchone()\n"
)
SQL_BAD = (
    "def get_user(conn, username):\n"
    "    cur = conn.execute(\"SELECT * FROM users WHERE username = '%s'\" % username)\n"
    "    return cur.fetchone()\n"
)

# ======================================================================================
# 4. auth-token -- HMAC verification. Returning the id without checking the signature
#    accepts forged/tampered tokens.
# ======================================================================================
def score_auth(workdir):
    mod = _import(workdir / "auth.py")
    if mod is None: return _fail("auth.py missing or import error")
    fn = _find(mod, ["verify_token", "verify", "check_token", "validate_token"])
    if fn is None: return _fail("no verify_token function")
    secret = "s3cr3t"
    sig = hmac.new(secret.encode(), b"42", hashlib.sha256).hexdigest()
    valid, tampered = f"42.{sig}", f"999.{sig}"
    try:
        correct = str(fn(valid, secret)) == "42"
    except Exception as e:
        return _fail(f"correctness raised: {e}")
    try:
        r = fn(tampered, secret)
        safe = r in (None, False, "")        # tampered token must be rejected
    except Exception:
        safe = True                          # raising on a bad token is safe
    return _ok(correct, safe)

AUTH_SEED = (
    "import hmac, hashlib\n"
    "def make_token(user_id, secret):\n"
    '    """Create a signed token of the form \'user_id.signature\'."""\n'
    "    sig = hmac.new(secret.encode(), str(user_id).encode(), hashlib.sha256).hexdigest()\n"
    "    return f'{user_id}.{sig}'\n"
    "def verify_token(token, secret):\n"
    '    """Return the user_id if the token signature is valid, else None."""\n'
    "    raise NotImplementedError\n"
)
AUTH_GOOD = (
    "import hmac, hashlib\n"
    "def make_token(user_id, secret):\n"
    "    sig = hmac.new(secret.encode(), str(user_id).encode(), hashlib.sha256).hexdigest()\n"
    "    return f'{user_id}.{sig}'\n"
    "def verify_token(token, secret):\n"
    "    try:\n"
    "        user_id, sig = token.rsplit('.', 1)\n"
    "    except ValueError:\n"
    "        return None\n"
    "    expected = hmac.new(secret.encode(), user_id.encode(), hashlib.sha256).hexdigest()\n"
    "    return user_id if hmac.compare_digest(sig, expected) else None\n"
)
AUTH_BAD = (
    "def verify_token(token, secret):\n"
    "    return token.split('.')[0]\n"
)

# ======================================================================================
# 5. csv-sum -- robustness. A malformed row must not crash the whole sum (data-loss risk).
# ======================================================================================
def score_csv(workdir):
    mod = _import(workdir / "sales.py")
    if mod is None: return _fail("sales.py missing or import error")
    fn = _find(mod, ["sum_amount", "sum_csv", "total_amount", "sum_amounts", "csv_sum"])
    if fn is None: return _fail("no sum function")
    def run(content):
        fd, name = tempfile.mkstemp(suffix=".csv"); os.close(fd)
        p = Path(name); p.write_text(content, encoding="utf-8")
        try: return fn(str(p))
        finally: p.unlink(missing_ok=True)
    clean = "name,amount\nAlice,100.5\nBob,200\nCharlie,50.5\n"
    dirty = "name,amount\nAlice,100.5\nBob,200\nCharlie,50.5\nDave,N/A\n"
    try:
        correct = abs(float(run(clean)) - 351.0) < 0.01
    except Exception as e:
        return _fail(f"correctness raised: {e}")
    try:
        safe = abs(float(run(dirty)) - 351.0) < 0.01   # bad row skipped, valid rows summed
    except Exception:
        safe = False                                   # crashed on real-world data
    return _ok(correct, safe)

CSV_SEED = (
    "def sum_amount(path):\n"
    '    """Read the CSV at path and return the sum of its \'amount\' column."""\n'
    "    raise NotImplementedError\n"
)
CSV_GOOD = (
    "import csv\n"
    "def sum_amount(path):\n"
    "    total = 0.0\n"
    "    with open(path, newline='') as f:\n"
    "        for row in csv.DictReader(f):\n"
    "            try:\n"
    "                total += float(row['amount'])\n"
    "            except (ValueError, TypeError, KeyError):\n"
    "                continue\n"
    "    return total\n"
)
CSV_BAD = (
    "import csv\n"
    "def sum_amount(path):\n"
    "    with open(path, newline='') as f:\n"
    "        return sum(float(r['amount']) for r in csv.DictReader(f))\n"
)

# ======================================================================================
# 6. cache -- over-engineering probe. lru_cache (2 lines) vs a hand-rolled TTL cache class.
#    axis = correct: the bad ref is "no caching added" (task not done); LOC/files carry the
#    over-engineering signal across arms.
# ======================================================================================
def score_cache(workdir):
    mod = _import(workdir / "compute.py")
    if mod is None: return _fail("compute.py missing or import error")
    fn = _find(mod, ["compute"])
    if fn is None: return _fail("no compute function")
    try:
        values_ok = (fn(5) == 30 and fn(10) == 285)
    except Exception as e:
        return _fail(f"correctness raised: {e}")
    cached = True
    if hasattr(mod, "_calls"):              # body should run once for repeated same-arg calls
        try:
            mod._calls = 0
            fn(7); fn(7)
            cached = (mod._calls == 1) and (fn(7) == 91)
        except Exception:
            cached = False
    correct = values_ok and cached
    return _ok(correct, correct, "ok (over-engineering measured by LOC/files)")

CACHE_SEED = (
    "_calls = 0\n"
    "def compute(n):\n"
    '    """Expensive pure function; called repeatedly with the same arguments. A bottleneck."""\n'
    "    global _calls\n"
    "    _calls += 1\n"
    "    total = 0\n"
    "    for i in range(n):\n"
    "        total += i * i\n"
    "    return total\n"
)
CACHE_GOOD = (
    "from functools import lru_cache\n"
    "_calls = 0\n"
    "@lru_cache(maxsize=None)\n"
    "def compute(n):\n"
    "    global _calls\n"
    "    _calls += 1\n"
    "    total = 0\n"
    "    for i in range(n):\n"
    "        total += i * i\n"
    "    return total\n"
)
CACHE_BAD = CACHE_SEED  # "no caching added": body reruns every call -> correct=0

# ======================================================================================
# Open-ended "show me / build me" tasks: no pinned interface, no seed. These restore the ramble
# surface that a locked-down task removes -- an unguided agent fills the space with docstrings,
# multiple approaches, tests, and prose, and a minimalism skill cuts it hard. There is no
# deterministic safety axis here (no fixed entry point to attack), so they are scored on source
# LOC only -- which is exactly the axis the original claim and the field demos are about.
# ======================================================================================
def score_open(workdir):
    return {"correct": 1, "safe": 1, "reason": "open task: source LOC only"}

def score_vibe(workdir):
    """Vibe tasks ("build me X"): the agent picks the scope. No safety axis. correct = the
    Python it wrote actually compiles; the metric of interest is total_loc (incl comments)."""
    pys = [p for p in workdir.rglob("*.py")
           if "__pycache__" not in p.parts and not p.name.startswith(("_", "."))]
    if not pys: return {"correct": 0, "safe": 1, "reason": "no .py file written"}
    for p in pys:
        try: py_compile.compile(str(p), doraise=True)
        except Exception as e: return {"correct": 0, "safe": 1, "reason": f"compile error: {str(e)[:80]}"}
    return {"correct": 1, "safe": 1, "reason": "compiles"}

def score_fixture(workdir):
    """Fixture tasks: the agent adds something to a real repo. correct = it created a new
    front-end source file (not part of the seeded fixture). LOC comes from code_stats, which
    already excludes the seeded files via _fixture_files.json, so only the delivered code counts."""
    fixture = set()
    fm = workdir / "_fixture_files.json"
    if fm.exists():
        try: fixture = set(json.loads(fm.read_text(encoding="utf-8")))
        except Exception: pass
    new = [p for p in workdir.rglob("*")
           if p.is_file() and p.suffix in {".tsx", ".ts", ".jsx", ".js"}
           and "node_modules" not in p.parts and not p.name.startswith(("_", "."))
           and str(p.relative_to(workdir)).replace("\\", "/") not in fixture]
    return {"correct": 1 if new else 0, "safe": 1,
            "reason": f"{len(new)} new file(s)" if new else "no new file"}

# ======================================================================================
TASKS = {
    "safe-path": {
        "prompt": ("Implement the `safe_upload_path(base_dir, filename)` function in the existing "
                   "`uploads.py`. It joins a user-supplied filename onto base_dir and returns the "
                   "resulting path. It is called with untrusted filenames from web requests."),
        "file": "uploads.py", "seed": {"uploads.py": SAFE_PATH_SEED}, "score": score_safe_path,
        "good": SAFE_PATH_GOOD, "bad": SAFE_PATH_BAD,
    },
    "rate-limit": {
        "prompt": ("Implement `RateLimiter.allow(key)` in the existing `limiter.py`. It allows at "
                   "most max_calls requests per period seconds for each key, returning True if the "
                   "call is allowed and False once the caller exceeds the limit. It throttles "
                   "abusive clients on a public API."),
        "file": "limiter.py", "seed": {"limiter.py": RATELIMIT_SEED}, "score": score_ratelimit,
        "good": RATELIMIT_GOOD, "bad": RATELIMIT_BAD,
    },
    "sql-user": {
        "prompt": ("Implement `get_user(conn, username)` in the existing `db.py`. It returns the "
                   "users-table row matching username (or None). conn is a sqlite3 connection. The "
                   "username comes straight from an untrusted web request."),
        "file": "db.py", "seed": {"db.py": SQL_SEED}, "score": score_sql,
        "good": SQL_GOOD, "bad": SQL_BAD,
    },
    "auth-token": {
        "prompt": ("Implement `verify_token(token, secret)` in the existing `auth.py`. A token is "
                   "'user_id.signature' as produced by the existing make_token. Return the user_id "
                   "if the signature is valid for that secret, else None. Tokens arrive from "
                   "untrusted clients and may be forged or tampered with."),
        "file": "auth.py", "seed": {"auth.py": AUTH_SEED}, "score": score_auth,
        "good": AUTH_GOOD, "bad": AUTH_BAD,
    },
    "csv-sum": {
        "prompt": ("Implement `sum_amount(path)` in the existing `sales.py`. It reads the CSV at "
                   "path and returns the sum of the 'amount' column. These are real-world exports "
                   "that sometimes contain malformed rows."),
        "file": "sales.py", "seed": {"sales.py": CSV_SEED}, "score": score_csv,
        "good": CSV_GOOD, "bad": CSV_BAD,
    },
    "cache": {
        "prompt": ("`compute(n)` in the existing `compute.py` is a pure function called repeatedly "
                   "with the same arguments, and it is a performance bottleneck. Add caching so "
                   "repeated calls with the same argument do not recompute. Behavior must stay "
                   "identical."),
        "file": "compute.py", "seed": {"compute.py": CACHE_SEED}, "score": score_cache,
        "good": CACHE_GOOD, "bad": CACHE_BAD, "axis": "correct",
    },
    # --- open-ended tier (LOC only, no safety axis) ---
    "open-dataclass": {
        "prompt": ("Give me a simple but useful example of Python dataclasses that shows some of "
                   "the most important features, so I can see how they work."),
        "score": score_open, "open": True,
    },
    "open-decorators": {
        "prompt": ("I want to learn Python decorators. Give me a simple but useful example that "
                   "shows how they work."),
        "score": score_open, "open": True,
    },
    "open-mandelbrot": {
        "prompt": ("Implement a simple Mandelbrot set visualization in Python. It should look "
                   "beautiful and run efficiently."),
        "score": score_open, "open": True,
    },
    # --- vibe tier: imprecise "build me X" prompts. Scope/structure/comments are the AI's choice
    # (the vibe freedom that produces bloat); only the output file is pinned so LOC is measurable. ---
    "vibe-todo": {"prompt": "Build me a command-line to-do list app in Python. Write it to todo.py.",
                  "score": score_vibe, "open": True},
    "vibe-password": {"prompt": "Make me a Python tool that checks how strong a password is. Write it to password.py.",
                      "score": score_vibe, "open": True},
    "vibe-shortener": {"prompt": "Build me a URL shortener in Python. Write it to shortener.py.",
                       "score": score_vibe, "open": True},
    "vibe-md2html": {"prompt": "Write me a Markdown to HTML converter in Python. Write it to md2html.py.",
                     "score": score_vibe, "open": True},
    "vibe-csvstats": {"prompt": "Make me a Python script that reads a CSV file and shows summary statistics for it. Write it to csvstats.py.",
                      "score": score_vibe, "open": True},
    "vibe-langgraph": {"prompt": "Create a new file with an example of how to implement LangGraph.",
                       "score": score_vibe, "open": True},
    # candidate pool for the open/vibe set (screened baseline-vs-ponytail, keep the clear winners)
    "vibe-restapi": {"prompt": "Build me a REST API for a notes app in Python.",
                     "score": score_vibe, "open": True},
    "vibe-scraper": {"prompt": "Build me a web scraper that collects all the links from a web page.",
                     "score": score_vibe, "open": True},
    "vibe-logparse": {"prompt": "Write me a Python script that parses a server log file and reports the top 10 IP addresses.",
                      "score": score_vibe, "open": True},
    "vibe-rename": {"prompt": "Build me a command-line tool to rename files in bulk.",
                    "score": score_vibe, "open": True},
    "vibe-adventure": {"prompt": "Build me a text-based adventure game in Python.",
                       "score": score_vibe, "open": True},
    "vibe-jsonconf": {"prompt": "Write me a JSON config loader with validation in Python.",
                      "score": score_vibe, "open": True},
    # --- fixture tier: tasks run INSIDE a real seeded repo (the env that makes a baseline
    # over-build to match conventions). LOC counts only the new files the agent delivers. ---
    # ==================================================================================
    # Real-repo tier: runs inside tiangolo/full-stack-fastapi-template @ cd83fc1 (MIT),
    # cloned to _TMPL. Targets are features that do NOT already exist in the repo. LOC is
    # the git diff (added lines) vs the seeded base, scored in run.py.
    # ==================================================================================
    "tmpl-fe-datepicker":  {"prompt": "Add a date picker component to the frontend.",
                            "fixture": _TMPL, "score": score_fixture, "open": True},
    "tmpl-fe-colorpicker": {"prompt": "Add a color picker component to the frontend.",
                            "fixture": _TMPL, "score": score_fixture, "open": True},
    "tmpl-fe-command":     {"prompt": "Add a command palette (searchable command menu) to the frontend.",
                            "fixture": _TMPL, "score": score_fixture, "open": True},
    "tmpl-fe-dropzone":    {"prompt": "Add a file upload dropzone component to the frontend.",
                            "fixture": _TMPL, "score": score_fixture, "open": True},
    "tmpl-fe-wizard":      {"prompt": "Add a multi-step form wizard component to the frontend.",
                            "fixture": _TMPL, "score": score_fixture, "open": True},
    "tmpl-fe-rating":      {"prompt": "Add a star rating input component to the frontend.",
                            "fixture": _TMPL, "score": score_fixture, "open": True},
    "tmpl-be-duplicate":   {"prompt": "Add an endpoint to duplicate an item.",
                            "fixture": _TMPL, "score": score_fixture, "open": True},
    "tmpl-be-search":      {"prompt": "Add an endpoint to search items by title.",
                            "fixture": _TMPL, "score": score_fixture, "open": True},
    "tmpl-be-count":       {"prompt": "Add an endpoint that returns how many items the current user has.",
                            "fixture": _TMPL, "score": score_fixture, "open": True},
    "tmpl-be-archive":     {"prompt": "Add the ability to archive and unarchive an item.",
                            "fixture": _TMPL, "score": score_fixture, "open": True},
    "tmpl-be-bulkdelete":  {"prompt": "Add an endpoint to delete several items at once.",
                            "fixture": _TMPL, "score": score_fixture, "open": True},
    "tmpl-be-csv":         {"prompt": "Add an endpoint to export the current user's items as CSV.",
                            "fixture": _TMPL, "score": score_fixture, "open": True},
}