test_measurement.py

"""Tests for the measurement / spend-avoided layer (tanglebrain/measurement.py).

Fully hermetic: the usage log is a temp path and pricing is injected, so nothing touches the real
~/.cache or the packaged config. Covers the estimation/cost math, the fault-tolerant log I/O, and
the rollup/format.
"""
from __future__ import annotations

import json
import tempfile
import unittest
from dataclasses import dataclass
from pathlib import Path
from unittest.mock import patch

from tanglebrain.measurement import (
    PLACEHOLDER_PRICING,
    PRICING_HEADER,
    Pricing,
    cloud_equiv_usd,
    default_log_path,
    estimate_tokens,
    format_rollup,
    load_pricing,
    record_task,
    read_records,
    rollup,
    save_pricing,
    validate_pricing,
)

# A fixed, non-placeholder pricing so cost assertions are exact and the caveat is off.
FIXED = Pricing(reference_model="test-frontier", input_per_mtok=2.0, output_per_mtok=10.0, is_placeholder=False)


@dataclass
class FakeEntry:
    """Stand-in for a RosterEntry (record_task only reads .tier / .id)."""

    id: str
    tier: str


class EstimateTokensTest(unittest.TestCase):
    def test_empty_is_zero(self):
        self.assertEqual(estimate_tokens(""), 0)
        self.assertEqual(estimate_tokens(None), 0)  # falsy guard

    def test_non_empty_is_at_least_one(self):
        self.assertEqual(estimate_tokens("ab"), 1)  # 2 // 4 == 0 -> clamped to 1

    def test_chars_over_four(self):
        self.assertEqual(estimate_tokens("a" * 40), 10)


class CloudEquivTest(unittest.TestCase):
    def test_known_math(self):
        # 1M in @ $2 + 1M out @ $10 = $12
        self.assertAlmostEqual(cloud_equiv_usd(1_000_000, 1_000_000, FIXED), 12.0)

    def test_zero_tokens_zero_cost(self):
        self.assertEqual(cloud_equiv_usd(0, 0, FIXED), 0.0)


class RecordTaskTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.mkdtemp()
        self.log = Path(self.tmp) / "sub" / "usage.jsonl"  # nested: parent must be created

    def test_appends_well_formed_record(self):
        record_task(
            path="local",
            entry=FakeEntry("gpt-oss-120b", "local"),
            prompt="a" * 40,
            response="b" * 80,
            log_path=self.log,
            pricing=FIXED,
        )
        records = read_records(self.log)
        self.assertEqual(len(records), 1)
        rec = records[0]
        self.assertEqual(rec["path"], "local")
        self.assertEqual(rec["tier"], "local")
        self.assertEqual(rec["model"], "gpt-oss-120b")
        self.assertEqual(rec["in_tokens_est"], 10)
        self.assertEqual(rec["out_tokens_est"], 20)
        self.assertAlmostEqual(rec["spend_avoided_usd"], cloud_equiv_usd(10, 20, FIXED), places=6)
        self.assertEqual(rec["pricing_ref"], "test-frontier")

    def test_appends_accumulate(self):
        for _ in range(3):
            record_task(path="router", entry=FakeEntry("claude", "sub"),
                        prompt="hi", response="there", log_path=self.log, pricing=FIXED)
        self.assertEqual(len(read_records(self.log)), 3)

    def test_api_tier_avoids_nothing(self):
        record_task(path="model", entry=FakeEntry("gpt-paid", "api"),
                    prompt="x" * 40, response="y" * 40, log_path=self.log, pricing=FIXED)
        rec = read_records(self.log)[0]
        self.assertGreater(rec["cloud_equiv_usd"], 0.0)
        self.assertEqual(rec["spend_avoided_usd"], 0.0)

    def test_none_entry_is_unknown(self):
        record_task(path="router", entry=None, prompt="x", response="y",
                    log_path=self.log, pricing=FIXED)
        rec = read_records(self.log)[0]
        self.assertEqual(rec["tier"], "unknown")
        self.assertEqual(rec["model"], "unknown")

    def test_logging_failure_never_raises(self):
        # Point the log at a path whose parent cannot be created (a file used as a directory).
        blocker = Path(self.tmp) / "blocker"
        blocker.write_text("i am a file")
        bad = blocker / "nested" / "usage.jsonl"
        # Must not raise despite the unwritable path.
        record_task(path="local", entry=FakeEntry("x", "local"), prompt="p", response="r",
                    log_path=bad, pricing=FIXED)


class ReadRecordsTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.mkdtemp()
        self.log = Path(self.tmp) / "usage.jsonl"

    def test_missing_file_is_empty(self):
        self.assertEqual(read_records(self.log), [])

    def test_skips_corrupt_and_blank_lines(self):
        self.log.write_text(
            json.dumps({"tier": "sub", "spend_avoided_usd": 1.0}) + "\n"
            "{ not valid json\n"
            "\n"
            + json.dumps({"tier": "local", "spend_avoided_usd": 2.0}) + "\n"
        )
        records = read_records(self.log)
        self.assertEqual(len(records), 2)  # the garbage + blank lines dropped


class RollupTest(unittest.TestCase):
    def test_totals_and_by_tier(self):
        records = [
            {"tier": "local", "in_tokens_est": 10, "out_tokens_est": 20,
             "cloud_equiv_usd": 1.0, "spend_avoided_usd": 1.0},
            {"tier": "sub", "in_tokens_est": 5, "out_tokens_est": 5,
             "cloud_equiv_usd": 0.5, "spend_avoided_usd": 0.5},
            {"tier": "local", "in_tokens_est": 1, "out_tokens_est": 1,
             "cloud_equiv_usd": 0.1, "spend_avoided_usd": 0.1},
        ]
        s = rollup(records)
        self.assertEqual(s["tasks"], 3)
        self.assertEqual(s["by_tier"], {"local": 2, "sub": 1})
        self.assertEqual(s["in_tokens_est"], 16)
        self.assertEqual(s["out_tokens_est"], 26)
        self.assertAlmostEqual(s["spend_avoided_usd"], 1.6)

    def test_empty_records(self):
        s = rollup([])
        self.assertEqual(s["tasks"], 0)
        self.assertEqual(s["by_tier"], {})
        self.assertEqual(s["spend_avoided_usd"], 0.0)

    def test_tolerates_bad_numeric_fields(self):
        s = rollup([{"tier": "local", "in_tokens_est": "oops", "spend_avoided_usd": None}])
        self.assertEqual(s["tasks"], 1)
        self.assertEqual(s["in_tokens_est"], 0)
        self.assertEqual(s["spend_avoided_usd"], 0.0)


class LoadPricingTest(unittest.TestCase):
    def test_loads_packaged_pricing(self):
        # The packaged config/pricing.yaml carries the default Claude Sonnet anchor ($3/$15).
        p = load_pricing()
        self.assertIsInstance(p, Pricing)
        self.assertFalse(p.is_placeholder)
        self.assertEqual(p.input_per_mtok, 3.0)
        self.assertEqual(p.output_per_mtok, 15.0)

    def test_missing_file_falls_back_to_placeholder(self):
        self.assertIs(load_pricing("/nonexistent/pricing.yaml"), PLACEHOLDER_PRICING)

    def test_corrupt_file_falls_back(self):
        tmp = tempfile.mkdtemp()
        bad = Path(tmp) / "pricing.yaml"
        bad.write_text("input_per_mtok: not-a-number\n")
        self.assertIs(load_pricing(bad), PLACEHOLDER_PRICING)


class ValidatePricingTest(unittest.TestCase):
    def _ok(self, **over):
        d = {"reference_model": "M", "input_per_mtok": 3.0, "output_per_mtok": 15.0, "placeholder": False}
        d.update(over)
        return d

    def test_valid(self):
        p = validate_pricing(self._ok())
        self.assertEqual((p.reference_model, p.input_per_mtok, p.output_per_mtok, p.is_placeholder),
                         ("M", 3.0, 15.0, False))

    def test_strips_reference_model(self):
        self.assertEqual(validate_pricing(self._ok(reference_model="  M  ")).reference_model, "M")

    def test_empty_model_rejected(self):
        with self.assertRaises(ValueError):
            validate_pricing(self._ok(reference_model="   "))

    def test_negative_rate_rejected(self):
        with self.assertRaises(ValueError):
            validate_pricing(self._ok(output_per_mtok=-5))

    def test_non_numeric_rate_rejected(self):
        with self.assertRaises(ValueError):
            validate_pricing(self._ok(input_per_mtok="lots"))

    def test_bool_rate_rejected(self):
        # bool is a subclass of int — must not slip through as a rate.
        with self.assertRaises(ValueError):
            validate_pricing(self._ok(input_per_mtok=True))

    def test_nan_rate_rejected(self):
        with self.assertRaises(ValueError):
            validate_pricing(self._ok(input_per_mtok=float("nan")))

    def test_non_bool_placeholder_rejected(self):
        with self.assertRaises(ValueError):
            validate_pricing(self._ok(placeholder="yes"))


class SavePricingTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.mkdtemp()
        self.path = Path(self.tmp) / "config" / "pricing.yaml"
        # Route backups into a temp state dir, not the real ~/.cache.
        self._env = patch.dict("os.environ", {"TANGLEBRAIN_STATE_DIR": str(Path(self.tmp) / "state")}, clear=False)
        self._env.start()
        self.addCleanup(self._env.stop)

    def test_roundtrips_and_preserves_header(self):
        p = Pricing("Claude Sonnet", 3.0, 15.0, False)
        save_pricing(p, self.path)
        text = self.path.read_text()
        self.assertIn(PRICING_HEADER.splitlines()[0], text)  # header survived
        back = load_pricing(self.path)
        self.assertEqual(back.reference_model, "Claude Sonnet")
        self.assertEqual(back.input_per_mtok, 3.0)
        self.assertFalse(back.is_placeholder)

    def test_no_tmp_left_behind(self):
        save_pricing(Pricing("M", 1.0, 2.0, True), self.path)
        leftovers = list(self.path.parent.glob("*.tmp"))
        self.assertEqual(leftovers, [])

    def test_backup_created_on_overwrite(self):
        save_pricing(Pricing("First", 1.0, 2.0, False), self.path)   # creates the file (no prior → no backup)
        save_pricing(Pricing("Second", 9.0, 9.0, False), self.path)  # overwrites → backs up "First"
        backups = list((Path(self.tmp) / "state" / "backups").glob("pricing-*.yaml"))
        self.assertTrue(backups)
        self.assertIn("First", backups[0].read_text())
        self.assertEqual(load_pricing(self.path).reference_model, "Second")

    def test_placeholder_flag_roundtrips(self):
        save_pricing(Pricing("M", 1.0, 2.0, True), self.path)
        self.assertTrue(load_pricing(self.path).is_placeholder)

    def test_preserves_existing_file_header_verbatim(self):
        # A save over an existing file must keep that file's curated header (no drift/doc loss).
        self.path.parent.mkdir(parents=True, exist_ok=True)
        custom = "# CUSTOM HEADER\n# COST_BASIS provenance line\n"
        self.path.write_text(custom + 'placeholder: false\nreference_model: "Old"\n'
                                      'input_per_mtok: 1.0\noutput_per_mtok: 2.0\n')
        save_pricing(Pricing("New", 9.0, 9.0, False), self.path)
        text = self.path.read_text()
        self.assertIn("# CUSTOM HEADER", text)
        self.assertIn("COST_BASIS provenance line", text)  # specific provenance survives the save
        self.assertEqual(load_pricing(self.path).reference_model, "New")

    def test_adversarial_values_roundtrip(self):
        # _render_pricing must produce YAML that load_pricing reads back identically.
        for ref in ['has: a colon', 'has "double" quotes', "has 'single'", "back\\slash",
                    "unicode ێ", "line\nbreak", "tab\there"]:
            with self.subTest(ref=ref):
                save_pricing(Pricing(ref, 0.0, 1e20, False), self.path)
                back = load_pricing(self.path)
                self.assertEqual(back.reference_model, ref)
                self.assertEqual(back.input_per_mtok, 0.0)
                self.assertEqual(back.output_per_mtok, 1e20)


class OriginAttributionTest(unittest.TestCase):
    """#74: the origin field on records, its rollup bucket, and the --stats line."""

    def setUp(self):
        self.tmp = tempfile.mkdtemp()
        self.log = Path(self.tmp) / "usage.jsonl"

    def _record(self, **kwargs):
        record_task(
            path="model", entry=FakeEntry("m", "local"), prompt="p", response="r",
            log_path=self.log, pricing=FIXED, **kwargs,
        )

    def test_record_writes_origin_when_given_and_omits_when_absent(self):
        self._record(origin="serve")
        self._record()
        tagged, untagged = read_records(self.log)
        self.assertEqual(tagged["origin"], "serve")
        self.assertNotIn("origin", untagged)

    def test_record_writes_parent_task_id_on_task_records(self):
        # #74: an external caller's identity (the serve header) rides on a kind="task" record.
        self._record(parent_task_id="tc-session-42")
        record = read_records(self.log)[0]
        self.assertEqual(record["kind"], "task")
        self.assertEqual(record["parent_task_id"], "tc-session-42")

    def test_rollup_buckets_by_origin_with_untagged_sentinel(self):
        summary = rollup([
            {"tier": "local", "origin": "serve"},
            {"tier": "local", "origin": "serve"},
            {"tier": "sub", "origin": "cli"},
            {"tier": "sub"},  # pre-#74 record — never guessed at
            {"kind": "delegate", "model": "m", "origin": "serve"},  # delegates stay out
        ])
        self.assertEqual(summary["by_origin"], {"serve": 2, "cli": 1, "untagged": 1})

    def test_format_rollup_shows_origin_split_only_when_tagged(self):
        tagged = format_rollup(
            rollup([{"tier": "local", "origin": "serve"}, {"tier": "local"}]), FIXED
        )
        self.assertIn("By origin:", tagged)
        self.assertIn("serve 1", tagged)
        self.assertIn("untagged 1", tagged)
        # All-untagged history says nothing — the line stays hidden.
        untagged_only = format_rollup(rollup([{"tier": "local"}]), FIXED)
        self.assertNotIn("By origin:", untagged_only)


class FormatRollupTest(unittest.TestCase):
    def test_renders_figures(self):
        s = rollup([{"tier": "local", "in_tokens_est": 10, "out_tokens_est": 20,
                     "cloud_equiv_usd": 1.5, "spend_avoided_usd": 1.5}])
        out = format_rollup(s, FIXED)
        self.assertIn("Tasks routed:", out)
        self.assertIn("$1.50", out)
        self.assertIn("test-frontier", out)
        self.assertNotIn("PLACEHOLDER", out)

    def test_placeholder_caveat_shown(self):
        out = format_rollup(rollup([]), PLACEHOLDER_PRICING)
        self.assertIn("PLACEHOLDER", out)


class DefaultLogPathTest(unittest.TestCase):
    def test_honors_state_dir_env(self):
        import os
        from unittest.mock import patch

        with patch.dict(os.environ, {"TANGLEBRAIN_STATE_DIR": "/tmp/tb-test"}, clear=False):
            self.assertEqual(default_log_path(), Path("/tmp/tb-test/usage.jsonl"))


class DelegateObservabilityTest(unittest.TestCase):
    """kind='delegate' records: written, kept out of the headline, rolled up separately, thread-safe."""

    def setUp(self):
        self.tmp = tempfile.mkdtemp()
        self.log = str(Path(self.tmp) / "usage.jsonl")

    def _read(self):
        return read_records(self.log)

    def test_record_defaults_to_task_kind(self):
        record_task(path="router", entry=FakeEntry("claude", "sub"),
                    prompt="hi", response="yo", log_path=self.log, pricing=FIXED)
        self.assertEqual(self._read()[0]["kind"], "task")

    def test_record_delegate_kind(self):
        record_task(path="delegate", entry=FakeEntry("local-x", "local"),
                    prompt="hi", response="yo", kind="delegate", log_path=self.log, pricing=FIXED)
        self.assertEqual(self._read()[0]["kind"], "delegate")

    def test_rollup_excludes_delegates_from_headline(self):
        records = [
            {"kind": "task", "tier": "sub", "in_tokens_est": 10, "out_tokens_est": 10,
             "cloud_equiv_usd": 1.0, "spend_avoided_usd": 1.0},
            {"kind": "delegate", "model": "local-x", "in_tokens_est": 100, "out_tokens_est": 200,
             "cloud_equiv_usd": 5.0, "spend_avoided_usd": 5.0},
        ]
        s = rollup(records)
        # Headline counts the one task only — delegate tokens/spend must NOT inflate it.
        self.assertEqual(s["tasks"], 1)
        self.assertEqual(s["by_tier"], {"sub": 1})
        self.assertEqual(s["in_tokens_est"], 10)
        self.assertEqual(s["out_tokens_est"], 10)
        self.assertAlmostEqual(s["spend_avoided_usd"], 1.0)
        # Delegate sub-rollup is separate + informational.
        d = s["delegates"]
        self.assertEqual(d["count"], 1)
        self.assertEqual(
            d["by_backend"],
            {"local-x": {"count": 1, "in_tokens_est": 100, "out_tokens_est": 200}},
        )
        self.assertEqual(d["in_tokens_est"], 100)
        self.assertEqual(d["out_tokens_est"], 200)
        self.assertAlmostEqual(d["cloud_equiv_usd"], 5.0)

    def test_kindless_record_counts_as_task(self):
        s = rollup([{"tier": "local", "in_tokens_est": 4, "spend_avoided_usd": 0.2}])
        self.assertEqual(s["tasks"], 1)
        self.assertEqual(s["delegates"]["count"], 0)

    def test_by_backend_aggregates_multiple(self):
        records = [
            {"kind": "delegate", "model": "local-x", "in_tokens_est": 10, "out_tokens_est": 5},
            {"kind": "delegate", "model": "local-x", "in_tokens_est": 20, "out_tokens_est": 5},
            {"kind": "delegate", "model": "cheap-sub", "in_tokens_est": 1, "out_tokens_est": 1},
        ]
        d = rollup(records)["delegates"]
        self.assertEqual(d["count"], 3)
        self.assertEqual(d["by_backend"]["local-x"]["count"], 2)
        self.assertEqual(d["by_backend"]["local-x"]["in_tokens_est"], 30)
        self.assertEqual(d["by_backend"]["cheap-sub"]["count"], 1)

    def test_record_writes_task_id_when_given(self):
        record_task(path="router", entry=FakeEntry("claude", "sub"), prompt="hi", response="yo",
                    task_id="task-abc", log_path=self.log, pricing=FIXED)
        rec = self._read()[0]
        self.assertEqual(rec["task_id"], "task-abc")
        self.assertNotIn("parent_task_id", rec)

    def test_record_writes_parent_task_id_for_delegate(self):
        record_task(path="delegate", entry=FakeEntry("local-x", "local"), prompt="hi", response="yo",
                    kind="delegate", parent_task_id="task-abc", log_path=self.log, pricing=FIXED)
        rec = self._read()[0]
        self.assertEqual(rec["parent_task_id"], "task-abc")
        self.assertNotIn("task_id", rec)

    def test_record_omits_linkage_fields_when_absent(self):
        record_task(path="local", entry=FakeEntry("local-x", "local"), prompt="hi", response="yo",
                    log_path=self.log, pricing=FIXED)
        rec = self._read()[0]
        self.assertNotIn("task_id", rec)
        self.assertNotIn("parent_task_id", rec)

    def test_rollup_groups_delegates_by_parent(self):
        records = [
            {"kind": "delegate", "model": "local-x", "parent_task_id": "p1"},
            {"kind": "delegate", "model": "cheap-sub", "parent_task_id": "p1"},
            {"kind": "delegate", "model": "local-x", "parent_task_id": "p2"},
        ]
        by_parent = rollup(records)["delegates"]["by_parent"]
        self.assertEqual(by_parent["p1"]["count"], 2)
        self.assertEqual(by_parent["p1"]["by_backend"], {"local-x": 1, "cheap-sub": 1})
        self.assertEqual(by_parent["p2"]["count"], 1)
        self.assertNotIn("unlinked", by_parent)

    def test_rollup_unlinked_delegate_grouped_under_sentinel(self):
        # A delegate with no parent_task_id (run outside a propagated task) groups under "unlinked".
        by_parent = rollup([{"kind": "delegate", "model": "local-x"}])["delegates"]["by_parent"]
        self.assertEqual(by_parent["unlinked"]["count"], 1)

    def test_format_shows_linked_parents(self):
        s = rollup([
            {"kind": "delegate", "model": "local-x", "parent_task_id": "p1"},
            {"kind": "delegate", "model": "local-x", "parent_task_id": "p2"},
            {"kind": "delegate", "model": "local-x"},
        ])
        out = format_rollup(s, FIXED)
        self.assertIn("Linked to:", out)
        self.assertIn("2 parent task(s)", out)
        self.assertIn("1 unlinked", out)

    def test_format_all_unlinked_reads_cleanly(self):
        # When no delegate is linked, the line should read "N unlinked", not "0 parent task(s), ...".
        s = rollup([{"kind": "delegate", "model": "local-x"},
                    {"kind": "delegate", "model": "local-x"}])
        out = format_rollup(s, FIXED)
        self.assertIn("Linked to:    2 unlinked", out)
        self.assertNotIn("parent task(s)", out)

    def test_format_shows_delegate_section_when_present(self):
        s = rollup([
            {"kind": "task", "tier": "sub", "in_tokens_est": 1, "out_tokens_est": 1,
             "cloud_equiv_usd": 0.1, "spend_avoided_usd": 0.1},
            {"kind": "delegate", "model": "local-x", "in_tokens_est": 50, "out_tokens_est": 50,
             "cloud_equiv_usd": 2.0},
        ])
        out = format_rollup(s, FIXED)
        self.assertIn("Delegated sub-tasks", out)
        self.assertIn("local-x", out)

    def test_format_omits_delegate_section_when_absent(self):
        s = rollup([{"kind": "task", "tier": "sub", "spend_avoided_usd": 0.1}])
        self.assertNotIn("Delegated sub-tasks", format_rollup(s, FIXED))

    def test_concurrent_appends_are_serialized(self):
        import threading

        def worker(n):
            record_task(path="delegate", entry=FakeEntry(f"m{n}", "local"),
                        prompt="p", response="r", kind="delegate", log_path=self.log, pricing=FIXED)

        threads = [threading.Thread(target=worker, args=(i,)) for i in range(20)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        recs = self._read()
        self.assertEqual(len(recs), 20)  # 20 well-formed lines — no interleaved/corrupted writes
        self.assertTrue(all(r.get("kind") == "delegate" for r in recs))


if __name__ == "__main__":
    unittest.main()