#!/usr/bin/env python3
"""Regression tests for rep.docstring_of -- the fabricated-DOC defect.

The DOC line is part of the model's INPUT and is presented as recovered ground truth, so a wrong
one instructs the model to invent a docstring that was never in the source. The published rule
(`docstring_of_published`) asked only "is co_consts[0] a string", which emitted a DOC line for 98
code objects across the two published benchmarks that have no implicit docstring at all.

Ground truth here is `ast.get_docstring` on the source, not a hand-written expectation, so these
tests measure the rule against Python's own definition of a docstring.

    python3 -m unittest test_rep_docstring -v          # unit cases only
    python3 test_rep_docstring.py --benchmarks         # + exhaustive sweep of both benchmarks
"""
from __future__ import annotations

import ast
import json
import marshal
import sys
import types
import unittest
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from pybytecode_core.rep import (CO_OPTIMIZED, disassemble_v2, docstring_of,  # noqa: E402
                                docstring_of_published)

ROOT = Path(__file__).resolve().parent.parent


def compile_src(src: str) -> types.CodeType:
    return compile(src, "<t>", "exec", dont_inherit=True, optimize=0)


def walk(co: types.CodeType):
    yield co
    for c in co.co_consts:
        if isinstance(c, types.CodeType):
            yield from walk(c)


def by_qualname(src: str) -> dict[str, types.CodeType]:
    return {c.co_qualname: c for c in walk(compile_src(src))}


class FabricatedDoc(unittest.TestCase):
    """Each case is a construct where the published rule emits a DOC line and must not."""

    def test_function_with_docstring_is_still_emitted(self):
        co = by_qualname("def f():\n    'the doc'\n    return 1\n")["f"]
        self.assertEqual(docstring_of(co), "the doc")

    def test_function_without_docstring_reserves_slot_zero(self):
        """CPython fills slot 0 with None for a docstring-less function, so BOTH rules agree.

        This is why the defect is invisible at function level and only shows up in genexps,
        modules and class bodies. Asserted so the claim is checked, not assumed.
        """
        co = by_qualname("def f():\n    return 'hello'\n")["f"]
        self.assertIsNone(co.co_consts[0])
        self.assertIsNone(docstring_of_published(co))
        self.assertIsNone(docstring_of(co))

    def test_generator_expression_slot_zero_is_a_literal(self):
        """A genexp is CO_OPTIMIZED but does NOT reserve slot 0 -- the real function-level defect."""
        src = "def f(xs):\n    return sum(1 for x in xs if x != '=')\n"
        gen = [c for c in by_qualname(src).values() if c.co_name == "<genexpr>"]
        self.assertTrue(gen, "no genexpr code object on this interpreter")
        co = gen[0]
        self.assertEqual(co.co_consts[0], "=")
        self.assertEqual(docstring_of_published(co), "=")       # the defect
        self.assertIsNone(docstring_of(co))                     # fixed

    def test_class_body_docstring_is_explicit_not_lost(self):
        # a class stores __doc__ explicitly, so DOC would duplicate the instruction stream
        co = by_qualname("class C:\n    'cdoc'\n    x = 1\n")["C"]
        self.assertIsNotNone(docstring_of_published(co))
        self.assertIsNone(docstring_of(co))

    def test_class_body_without_docstring_emits_nothing(self):
        co = by_qualname("class C:\n    x = 1\n")["C"]
        self.assertIsNone(docstring_of(co))

    def test_module_docstring_is_explicit_not_lost(self):
        co = compile_src("'mdoc'\nx = 1\n")
        self.assertIsNotNone(docstring_of_published(co))
        self.assertIsNone(docstring_of(co))

    def test_module_first_const_is_a_default_argument(self):
        # the case that makes an "is slot 0 referenced" test insufficient: slot 0 is never loaded
        # directly, only inside the defaults tuple
        co = compile_src("def f(c='WMAP5'):\n    return c\n")
        self.assertEqual(co.co_consts[0], "WMAP5")
        self.assertIsNotNone(docstring_of_published(co))
        self.assertIsNone(docstring_of(co))

    def test_generator_expression_has_no_docstring(self):
        src = "def f(xs):\n    return tuple('a' for _ in xs)\n"
        for q, co in by_qualname(src).items():
            if co.co_name == "<genexpr>":
                self.assertIsNone(docstring_of(co), q)
                break
        else:
            self.skipTest("no genexpr code object on this interpreter")

    def test_docstring_equal_to_a_body_literal_is_not_lost(self):
        """CPython de-duplicates consts, so one slot serves both the docstring and the literal.

        A rule that reasoned from "is slot 0 loaded" would drop this docstring -- reintroducing the
        v1 information loss this module exists to undo. Scope kind gets it right.
        """
        src = "def f():\n    'pass'\n    x = 'pass'\n    return x\n"
        co = by_qualname(src)["f"]
        self.assertEqual(ast.get_docstring(ast.parse(src).body[0]), "pass")
        self.assertEqual(co.co_consts, ("pass",))
        self.assertEqual(docstring_of(co), "pass")

    def test_lambda_has_no_docstring(self):
        co = by_qualname("f = lambda: 'a'\n")["<lambda>"]
        self.assertIsNone(docstring_of(co))


class RepWiring(unittest.TestCase):
    def test_default_rule_is_the_published_one(self):
        """Changing this default changes the public benchmark's input. It must be deliberate."""
        co = compile_src("def f(xs):\n    return sum(1 for x in xs if x != '=')\n")
        self.assertIn("DOC '='", disassemble_v2(co))
        self.assertIn("DOC '='", disassemble_v2(co, doc_rule="published"))
        self.assertNotIn("DOC '='", disassemble_v2(co, doc_rule="fixed"))

    def test_bad_rule_rejected(self):
        with self.assertRaises(ValueError):
            disassemble_v2(compile_src("x = 1\n"), doc_rule="v2")

    def test_fixed_rule_propagates_into_nested_code_objects(self):
        src = "class C:\n    'cdoc'\n    def m(self):\n        return 'lit'\n"
        co = compile_src(src)
        self.assertNotIn("DOC", disassemble_v2(co, doc_rule="fixed"))
        self.assertIn("DOC", disassemble_v2(co, doc_rule="published"))


def ast_truth(src: str) -> dict[str, bool]:
    """qualname -> has an implicit docstring, mirroring co_qualname. Ground truth."""
    out: dict[str, bool] = {}

    def rec(node, prefix):
        for ch in ast.iter_child_nodes(node):
            if isinstance(ch, (ast.FunctionDef, ast.AsyncFunctionDef)):
                q = f"{prefix}{ch.name}"
                out[q] = ast.get_docstring(ch) is not None
                rec(ch, f"{q}.<locals>.")
            elif isinstance(ch, ast.ClassDef):
                q = f"{prefix}{ch.name}"
                out[q] = ast.get_docstring(ch) is not None
                rec(ch, f"{q}.")
            else:
                rec(ch, prefix)

    tree = ast.parse(src)
    out["<module>"] = ast.get_docstring(tree) is not None
    rec(tree, "")
    return out


def sweep_benchmarks() -> int:
    """Exhaustive: every code object of both published benchmarks against ast.get_docstring."""
    fails = 0
    for name, d in (("csn-3.12-licensed", ROOT / "benchmarks" / "csn-3.12-licensed"),
                    ("mbpp-ood", ROOT / "benchmarks" / "mbpp-ood")):
        if not (d / "bench.jsonl").exists():
            print(f"  SKIP {name} (not present)")
            continue
        rows = [json.loads(l) for l in (d / "bench.jsonl").read_text().splitlines() if l.strip()]
        n = spurious_old = spurious_new = missed_new = 0
        for r in rows:
            src = (d / r["src_path"]).read_text()
            truth = ast_truth(src)
            co = marshal.loads((d / r["pyc_path"]).read_bytes()[16:])
            for c in walk(co):
                n += 1
                real = truth.get(c.co_qualname) is True and bool(c.co_flags & CO_OPTIMIZED)
                spurious_old += (docstring_of_published(c) is not None) and not real
                got = docstring_of(c) is not None
                spurious_new += got and not real
                missed_new += real and not got
        ok = spurious_new == 0 and missed_new == 0
        print(f"  {'OK  ' if ok else 'FAIL'} {name}: {n} code objects | spurious DOC "
              f"published={spurious_old} fixed={spurious_new} | real docstrings missed="
              f"{missed_new}")
        fails += 0 if ok else 1
    return fails


if __name__ == "__main__":
    if "--benchmarks" in sys.argv:
        sys.argv.remove("--benchmarks")
        print("=== exhaustive sweep of both published benchmarks ===")
        rc = sweep_benchmarks()
        print()
        r = unittest.main(exit=False, verbosity=2).result
        raise SystemExit(1 if rc or not r.wasSuccessful() else 0)
    unittest.main()