from datetime import datetime, timedelta, timezone
from uuid import uuid4

import unittest

from agenttrace import Run, Span, build_ingest_payload


class ContractsTest(unittest.TestCase):
    def test_build_ingest_payload_happy_path(self) -> None:
        run_id = uuid4()
        root_span_id = uuid4()
        child_span_id = uuid4()
        now = datetime.now(timezone.utc)

        run = Run(run_id=run_id, name="demo", created_at=now, metadata_json={"env": "test"})
        root_span = Span(
            span_id=root_span_id,
            run_id=run_id,
            parent_span_id=None,
            span_type="agent",
            name="root",
            start_time=now,
            end_time=now + timedelta(milliseconds=10),
            duration_ms=10,
            input_json={"q": "hello"},
            output_json={"a": "world"},
            error=None,
            metadata_json={},
        )
        child_span = Span(
            span_id=child_span_id,
            run_id=run_id,
            parent_span_id=root_span_id,
            span_type="tool",
            name="tool-call",
            start_time=now,
            end_time=now + timedelta(milliseconds=5),
            duration_ms=5,
            input_json={"x": 1},
            output_json={"y": 2},
            error=None,
            metadata_json={},
        )

        payload = build_ingest_payload([run], [root_span, child_span])

        self.assertEqual(len(payload["runs"]), 1)
        self.assertEqual(len(payload["spans"]), 2)

    def test_reject_unknown_parent_span(self) -> None:
        run_id = uuid4()
        now = datetime.now(timezone.utc)
        run = Run(run_id=run_id, name="demo", created_at=now, metadata_json={})
        span = Span(
            span_id=uuid4(),
            run_id=run_id,
            parent_span_id=uuid4(),
            span_type="chain",
            name="bad-parent",
            start_time=now,
            end_time=now,
            duration_ms=0,
            input_json={},
            output_json={},
            error=None,
            metadata_json={},
        )

        with self.assertRaises(ValueError):
            build_ingest_payload([run], [span])

    def test_reject_span_with_unknown_run(self) -> None:
        now = datetime.now(timezone.utc)
        run = Run(run_id=uuid4(), name="demo", created_at=now, metadata_json={})
        span = Span(
            span_id=uuid4(),
            run_id=uuid4(),
            parent_span_id=None,
            span_type="llm",
            name="orphan-run",
            start_time=now,
            end_time=now,
            duration_ms=0,
            input_json={},
            output_json={},
            error=None,
            metadata_json={},
        )

        with self.assertRaises(ValueError):
            build_ingest_payload([run], [span])


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