import json
import os
import tempfile
import unittest
import uuid
from contextlib import redirect_stdout
from io import StringIO
from pathlib import Path

from hiveos import config as config_module
from hiveos.intelligence import tracing as tracing_module
from hiveos.observe import ingest_tei_event, ingest_tei_events, main as observe_main
from hiveos.tei import TEI_VERSION


def _valid_tei_event(*, run_id):
    return {
        "tei_version": TEI_VERSION,
        "event_id": f"evt-{uuid.uuid4().hex[:8]}",
        "event_type": "agent.tool_call",
        "occurred_at_ms": 1772800000000,
        "run_id": run_id,
        "source": {"framework": "demo-framework", "emitter": "adapter.v1"},
        "flow_id": "flow:demo",
        "step_id": "step:search",
        "parent_step_id": "step:plan",
        "tool_call_id": "tool:web.search:1",
        "attempt": 1,
        "step_type": "tool_call",
        "outcome": "success",
        "payload": {"query": "hive trace tei"},
        "tags": ["demo"],
        "metadata": {"lane": "test"},
    }


def _tei_replay_event(
    *,
    run_id,
    event_type,
    event_id,
    occurred_at_ms,
    step_id="",
    parent_step_id="",
    tool_call_id="",
    step_type="",
    outcome="success",
    payload=None,
):
    event = {
        "tei_version": TEI_VERSION,
        "event_id": event_id,
        "event_type": event_type,
        "occurred_at_ms": int(occurred_at_ms),
        "run_id": run_id,
        "source": {"framework": "openclaw", "emitter": "hive-openclaw-shim.v1"},
        "flow_id": "flow:openclaw-live",
        "outcome": outcome,
        "payload": dict(payload or {}),
    }
    if step_id:
        event["step_id"] = step_id
    if parent_step_id:
        event["parent_step_id"] = parent_step_id
    if tool_call_id:
        event["tool_call_id"] = tool_call_id
    if step_type:
        event["step_type"] = step_type
    return event


class TeiIngestTests(unittest.TestCase):
    def setUp(self):
        self._original_trace_db_async = tracing_module.TRACE_DB_ASYNC
        self._original_trace_sampling_enabled = tracing_module.TRACE_SAMPLING_ENABLED
        self._original_trace_log_path = tracing_module.TRACE_LOG_PATH
        self._original_trace_home = os.environ.get("HIVE_TRACE_HOME")
        self._original_tei_ingest_enabled = config_module.TRACE_TEI_INGEST_ENABLED
        self._original_tei_strict_version = config_module.TRACE_TEI_STRICT_VERSION

        self._trace_log_path = str(Path(tempfile.gettempdir()) / f"hive-tei-{uuid.uuid4().hex}.log")
        self._trace_home = str(Path(tempfile.gettempdir()) / f"hive-tei-home-{uuid.uuid4().hex}")
        tracing_module.TRACE_LOG_PATH = self._trace_log_path
        tracing_module.TRACE_DB_ASYNC = False
        tracing_module.TRACE_SAMPLING_ENABLED = False
        os.environ["HIVE_TRACE_HOME"] = self._trace_home
        config_module.TRACE_TEI_INGEST_ENABLED = False
        config_module.TRACE_TEI_STRICT_VERSION = True

    def tearDown(self):
        tracing_module.TRACE_DB_ASYNC = self._original_trace_db_async
        tracing_module.TRACE_SAMPLING_ENABLED = self._original_trace_sampling_enabled
        tracing_module.TRACE_LOG_PATH = self._original_trace_log_path
        config_module.TRACE_TEI_INGEST_ENABLED = self._original_tei_ingest_enabled
        config_module.TRACE_TEI_STRICT_VERSION = self._original_tei_strict_version
        if self._original_trace_home is None:
            os.environ.pop("HIVE_TRACE_HOME", None)
        else:
            os.environ["HIVE_TRACE_HOME"] = self._original_trace_home
        self._cleanup_trace_files()
        self._cleanup_trace_home()

    def _cleanup_trace_files(self):
        base = Path(self._trace_log_path)
        for path in [base] + list(base.parent.glob(f"{base.name}.*")):
            try:
                if path.exists():
                    path.unlink()
            except Exception:
                pass

    def _cleanup_trace_home(self):
        home = Path(self._trace_home)
        if not home.exists():
            return
        for path in sorted(home.rglob("*"), reverse=True):
            try:
                if path.is_file():
                    path.unlink()
                elif path.is_dir():
                    path.rmdir()
            except Exception:
                pass

    def test_ingest_tei_event_is_gated_by_flag(self):
        run_id = f"observe-run:{uuid.uuid4().hex[:8]}"
        result = ingest_tei_event(_valid_tei_event(run_id=run_id))
        self.assertFalse(result["ok"])
        self.assertEqual("tei_ingest_disabled", result["reason"])

    def test_ingest_tei_event_writes_trace_when_enabled(self):
        config_module.TRACE_TEI_INGEST_ENABLED = True
        run_id = f"observe-run:{uuid.uuid4().hex[:8]}"
        result = ingest_tei_event(_valid_tei_event(run_id=run_id))
        self.assertTrue(result["ok"])
        self.assertEqual(run_id, result["run_id"])
        self.assertEqual("tool.result", result["event_type"])

        events = tracing_module.read_trace_events(limit=100, filters={"run_id": run_id, "event_type": "tool.result"})
        self.assertGreaterEqual(len(events), 1)
        payload = events[-1].get("payload") or {}
        self.assertEqual("tei", payload.get("source"))
        self.assertEqual(TEI_VERSION, ((payload.get("tei") or {}).get("tei_version") or ""))
        self.assertEqual("tool.result", ((payload.get("tei") or {}).get("canonical_event_type") or ""))
        self.assertEqual("success", ((payload.get("tei") or {}).get("status") or ""))
        self.assertEqual("flow:demo", events[-1].get("flow_id"))
        self.assertEqual("step:search", events[-1].get("step_id"))
        contract = payload.get("contract_v1") or {}
        self.assertEqual("1.0", contract.get("contract_version"))
        self.assertEqual("tool.result", contract.get("event_type"))
        self.assertEqual("success", contract.get("status"))
        self.assertIn("source", contract)
        self.assertIn("anchor", contract)

    def test_ingest_tei_event_rejects_version_when_strict(self):
        config_module.TRACE_TEI_INGEST_ENABLED = True
        config_module.TRACE_TEI_STRICT_VERSION = True
        run_id = f"observe-run:{uuid.uuid4().hex[:8]}"
        event = _valid_tei_event(run_id=run_id)
        event["tei_version"] = "0.0"
        result = ingest_tei_event(event)
        self.assertFalse(result["ok"])
        self.assertEqual("tei_validation_failed", result["reason"])
        self.assertTrue(any("unsupported tei_version" in err for err in (result.get("errors") or [])))

    def test_ingest_tei_event_accepts_version_when_not_strict(self):
        config_module.TRACE_TEI_INGEST_ENABLED = True
        config_module.TRACE_TEI_STRICT_VERSION = False
        run_id = f"observe-run:{uuid.uuid4().hex[:8]}"
        event = _valid_tei_event(run_id=run_id)
        event["tei_version"] = "0.0"
        result = ingest_tei_event(event)
        self.assertTrue(result["ok"])
        events = tracing_module.read_trace_events(limit=50, filters={"run_id": run_id})
        self.assertGreaterEqual(len(events), 1)

    def test_ingest_tei_events_batch_returns_itemized_results(self):
        config_module.TRACE_TEI_INGEST_ENABLED = True
        run_id = f"observe-run:{uuid.uuid4().hex[:8]}"
        valid = _valid_tei_event(run_id=run_id)
        invalid = dict(valid)
        invalid["payload"] = "bad"
        result = ingest_tei_events([valid, invalid])
        self.assertFalse(result["ok"])
        self.assertEqual(2, result["total"])
        self.assertEqual(1, result["ingested"])
        self.assertEqual(1, result["failed"])
        self.assertEqual(2, len(result["items"]))
        self.assertTrue(result["items"][0]["ok"])
        self.assertFalse(result["items"][1]["ok"])

    def test_trace_tei_validate_batch_json_reports_item_metadata(self):
        run_id = f"observe-run:{uuid.uuid4().hex[:8]}"
        valid = _valid_tei_event(run_id=run_id)
        invalid = dict(valid)
        invalid["event_id"] = f"evt-{uuid.uuid4().hex[:8]}"
        invalid["payload"] = "bad"
        batch_path = Path(tempfile.gettempdir()) / f"hive-tei-validate-{uuid.uuid4().hex}.json"
        batch_path.write_text(json.dumps([valid, invalid]), encoding="utf-8")
        out = StringIO()
        with redirect_stdout(out):
            code = observe_main(["trace", "tei", "validate", "--file", str(batch_path), "--json"])
        self.assertEqual(1, code)
        payload = json.loads(out.getvalue().strip())
        self.assertEqual("batch", payload["mode"])
        self.assertEqual(2, payload["total"])
        self.assertEqual(1, payload["valid"])
        self.assertEqual(1, payload["invalid"])
        self.assertEqual(run_id, payload["items"][0]["run_id"])
        self.assertFalse(payload["items"][1]["ok"])
        self.assertTrue(any("payload must be an object" in err for err in (payload["items"][1].get("errors") or [])))
        try:
            batch_path.unlink(missing_ok=True)
        except Exception:
            pass

    def test_trace_tei_validate_single_accepts_non_strict_version(self):
        run_id = f"observe-run:{uuid.uuid4().hex[:8]}"
        event = _valid_tei_event(run_id=run_id)
        event["tei_version"] = "0.0"
        event_path = Path(tempfile.gettempdir()) / f"hive-tei-validate-single-{uuid.uuid4().hex}.json"
        event_path.write_text(json.dumps(event), encoding="utf-8")
        out = StringIO()
        with redirect_stdout(out):
            code = observe_main(["trace", "tei", "validate", "--file", str(event_path), "--no-strict-version", "--json"])
        self.assertEqual(0, code)
        payload = json.loads(out.getvalue().strip())
        self.assertEqual("single", payload["mode"])
        self.assertTrue(payload["ok"])
        self.assertEqual(run_id, payload["items"][0]["run_id"])
        try:
            event_path.unlink(missing_ok=True)
        except Exception:
            pass

    def test_trace_tei_ingest_batch_json_writes_events(self):
        config_module.TRACE_TEI_INGEST_ENABLED = True
        run_id = f"observe-run:{uuid.uuid4().hex[:8]}"
        valid = _valid_tei_event(run_id=run_id)
        invalid = dict(valid)
        invalid["event_id"] = f"evt-{uuid.uuid4().hex[:8]}"
        invalid["payload"] = "bad"
        batch_path = Path(tempfile.gettempdir()) / f"hive-tei-ingest-{uuid.uuid4().hex}.json"
        batch_path.write_text(json.dumps([valid, invalid]), encoding="utf-8")
        out = StringIO()
        with redirect_stdout(out):
            code = observe_main(["trace", "tei", "ingest", "--file", str(batch_path), "--json"])
        self.assertEqual(1, code)
        payload = json.loads(out.getvalue().strip())
        self.assertFalse(payload["ok"])
        self.assertEqual(2, payload["total"])
        self.assertEqual(1, payload["ingested"])
        self.assertEqual(1, payload["failed"])
        events = tracing_module.read_trace_events(limit=200, filters={"run_id": run_id})
        self.assertGreaterEqual(len(events), 1)
        try:
            batch_path.unlink(missing_ok=True)
        except Exception:
            pass

    def test_trace_tei_ingest_respects_ingest_disabled_flag(self):
        config_module.TRACE_TEI_INGEST_ENABLED = False
        run_id = f"observe-run:{uuid.uuid4().hex[:8]}"
        event = _valid_tei_event(run_id=run_id)
        event_path = Path(tempfile.gettempdir()) / f"hive-tei-ingest-disabled-{uuid.uuid4().hex}.json"
        event_path.write_text(json.dumps(event), encoding="utf-8")
        out = StringIO()
        with redirect_stdout(out):
            code = observe_main(["trace", "tei", "ingest", "--file", str(event_path), "--json"])
        self.assertEqual(1, code)
        payload = json.loads(out.getvalue().strip())
        self.assertFalse(payload["ok"])
        self.assertEqual("tei_ingest_disabled", payload["reason"])
        try:
            event_path.unlink(missing_ok=True)
        except Exception:
            pass

    def test_tei_run_supports_anchors_and_replay_plan_recommended(self):
        config_module.TRACE_TEI_INGEST_ENABLED = True
        run_id = f"observe-run:tei-live-{uuid.uuid4().hex[:8]}"
        command = "python -c \"print('tei replay source')\""
        events = [
            _tei_replay_event(
                run_id=run_id,
                event_type="turn.start",
                event_id=f"evt-{uuid.uuid4().hex[:8]}",
                occurred_at_ms=1772800001000,
                step_id="step:turn.start",
                step_type="agent_step",
                payload={"step_name": "turn.start", "command": command},
            ),
            _tei_replay_event(
                run_id=run_id,
                event_type="tool.request",
                event_id=f"evt-{uuid.uuid4().hex[:8]}",
                occurred_at_ms=1772800001010,
                step_id="step:openclaw.command",
                parent_step_id="step:turn.start",
                tool_call_id="tool:openclaw.command:1",
                step_type="tool_call",
                payload={"step_name": "tool.request", "command": command, "idempotency_hint": "high"},
            ),
            _tei_replay_event(
                run_id=run_id,
                event_type="tool.result",
                event_id=f"evt-{uuid.uuid4().hex[:8]}",
                occurred_at_ms=1772800001030,
                step_id="step:openclaw.command",
                parent_step_id="step:turn.start",
                tool_call_id="tool:openclaw.command:1",
                step_type="tool_call",
                payload={"step_name": "tool.result", "exit_code": 0, "idempotency_hint": "medium"},
            ),
            _tei_replay_event(
                run_id=run_id,
                event_type="turn.finish",
                event_id=f"evt-{uuid.uuid4().hex[:8]}",
                occurred_at_ms=1772800001040,
                step_id="step:turn.finish",
                parent_step_id="step:openclaw.command",
                step_type="agent_step",
                payload={"step_name": "turn.finish", "exit_code": 0},
            ),
        ]
        result = ingest_tei_events(events)
        self.assertTrue(result["ok"])

        anchors_out = StringIO()
        with redirect_stdout(anchors_out):
            anchors_code = observe_main(["trace", "anchors", run_id, "--json"])
        self.assertEqual(0, anchors_code)
        anchors_payload = json.loads(anchors_out.getvalue().strip())
        self.assertGreaterEqual(int(anchors_payload.get("anchor_count") or 0), 1)
        self.assertGreaterEqual(len(list(anchors_payload.get("recommended_replay_anchors") or [])), 1)

        plan_out = StringIO()
        with redirect_stdout(plan_out):
            plan_code = observe_main(["trace", "replay-plan", run_id, "--recommended", "--json"])
        self.assertEqual(0, plan_code)
        plan_payload = json.loads(plan_out.getvalue().strip())
        replay_command = plan_payload.get("replay_command")
        if isinstance(replay_command, list):
            replay_text = " ".join([str(part) for part in replay_command if str(part)])
        else:
            replay_text = str(replay_command or "")
        self.assertIn("python -c", replay_text)
        run_file = Path(self._trace_home) / "runs" / f"{run_id.replace(':', '_')}.json"
        self.assertTrue(run_file.exists())
        record = json.loads(run_file.read_text(encoding="utf-8"))
        self.assertEqual("success", str(record.get("status") or ""))
        self.assertEqual(0, int(record.get("exit_code") or -1))
        self.assertGreaterEqual(int(record.get("finished_at_ms") or 0), int(record.get("started_at_ms") or 0))

    def test_tei_error_event_finalizes_run_as_failed(self):
        config_module.TRACE_TEI_INGEST_ENABLED = True
        run_id = f"observe-run:tei-live-{uuid.uuid4().hex[:8]}"
        events = [
            _tei_replay_event(
                run_id=run_id,
                event_type="turn.start",
                event_id=f"evt-{uuid.uuid4().hex[:8]}",
                occurred_at_ms=1772800002000,
                step_id="step:turn.start",
                step_type="agent_step",
                payload={"step_name": "turn.start", "command": "openclaw agent --message fail"},
            ),
            _tei_replay_event(
                run_id=run_id,
                event_type="error",
                event_id=f"evt-{uuid.uuid4().hex[:8]}",
                occurred_at_ms=1772800002010,
                step_id="step:openclaw.command",
                parent_step_id="step:turn.start",
                step_type="tool_call",
                outcome="failed",
                payload={"step_name": "error", "error": "boom"},
            ),
        ]
        result = ingest_tei_events(events)
        self.assertTrue(result["ok"])
        run_file = Path(self._trace_home) / "runs" / f"{run_id.replace(':', '_')}.json"
        self.assertTrue(run_file.exists())
        record = json.loads(run_file.read_text(encoding="utf-8"))
        self.assertEqual("failed", str(record.get("status") or ""))
        self.assertNotEqual(0, int(record.get("exit_code") or 0))


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