import hashlib
import json
from pathlib import Path
import sys
import unittest


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from tools.tushare_lineage import (
    TushareLineageError,
    _compare_source_jobs,
    _validate_snapshot,
)


def _canonical(value):
    return json.dumps(
        value,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
        allow_nan=False,
    ).encode()


class TushareLineageTest(unittest.TestCase):
    def _snapshot(self):
        job = {
            "api_name": "daily",
            "byte_count": 123,
            "completed_at": "2026-07-31T00:00:00Z",
            "id": 7,
            "observed": {
                "row_count": 11,
                "sha256": "a" * 64,
                "size_bytes": 123,
            },
            "partition_key": "month=2024-01",
            "path": "data/parquet/daily/month=2024-01.parquet",
            "request_sha256": "b" * 64,
            "row_count": 11,
            "sha256": "a" * 64,
            "shadowed_completed_count": 0,
            "verified": True,
        }
        selection_identity = {
            key: value
            for key, value in job.items()
            if key not in {"observed", "verified"}
        }
        snapshot = {
            "format": "quant-os.tushare-scoped-snapshot/v2",
            "selected_jobs": [job],
            "selection_sha256": hashlib.sha256(
                _canonical({"selected_jobs": [selection_identity]})
            ).hexdigest(),
            "verification": {
                "double_verified_jobs": 1,
                "selected_jobs_still_latest_after_verification": True,
                "verified_jobs": 1,
            },
        }
        snapshot["manifest_sha256"] = hashlib.sha256(
            _canonical(snapshot)
        ).hexdigest()
        return snapshot

    def test_exact_source_lineage_and_tamper_detection(self):
        snapshot = self._snapshot()
        jobs = _validate_snapshot(snapshot)
        provider_job = {
            "api_name": "daily",
            "byte_count": 123,
            "job_id": 7,
            "partition_key": "month=2024-01",
            "path": "data/parquet/daily/month=2024-01.parquet",
            "request_sha256": "b" * 64,
            "row_count": 11,
            "sha256": "a" * 64,
        }
        _compare_source_jobs(jobs, [provider_job])
        provider_job["sha256"] = "c" * 64
        with self.assertRaisesRegex(
            TushareLineageError,
            "source job fields differ",
        ):
            _compare_source_jobs(jobs, [provider_job])

        snapshot["selected_jobs"][0]["path"] = "../escape.parquet"
        snapshot_without_hash = dict(snapshot)
        snapshot_without_hash.pop("manifest_sha256")
        snapshot["manifest_sha256"] = hashlib.sha256(
            _canonical(snapshot_without_hash)
        ).hexdigest()
        with self.assertRaisesRegex(
            TushareLineageError,
            "stay relative",
        ):
            _validate_snapshot(snapshot)


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