import tempfile
import unittest
from contextlib import redirect_stdout
from io import StringIO
from pathlib import Path
import uuid
import os
import json
import urllib.request
import re
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import threading

from hiveos.intelligence import tracing as tracing_module
from hiveos import hive as hive_module
from hiveos.observe import (
    observe_checkpoint,
    observe_rerun_request,
    observe_run,
    observe_proxy,
    observe_step,
    run_command_observed,
    main as observe_main,
)


class ObserveHarnessTests(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_disable_browser = os.environ.get("HIVE_TRACE_DISABLE_BROWSER")

        self._trace_log_path = str(Path(tempfile.gettempdir()) / f"hive-observe-{uuid.uuid4().hex}.log")
        self._trace_home = str(Path(tempfile.gettempdir()) / f"hive-observe-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
        os.environ["HIVE_TRACE_DISABLE_BROWSER"] = "1"

    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
        if self._original_trace_home is None:
            os.environ.pop("HIVE_TRACE_HOME", None)
        else:
            os.environ["HIVE_TRACE_HOME"] = self._original_trace_home
        if self._original_disable_browser is None:
            os.environ.pop("HIVE_TRACE_DISABLE_BROWSER", None)
        else:
            os.environ["HIVE_TRACE_DISABLE_BROWSER"] = self._original_disable_browser
        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_observe_run_context_emits_start_step_finish(self):
        with observe_run("unit-test-run", metadata={"suite": "observe"}) as run_id:
            observe_step(run_id, "step-alpha", payload={"ok": True})

        started = tracing_module.read_trace_events(limit=200, filters={"event_type": "observe_run_started"})
        stepped = tracing_module.read_trace_events(limit=200, filters={"event_type": "observe_step"})
        finished = tracing_module.read_trace_events(limit=200, filters={"event_type": "observe_run_finished"})

        self.assertTrue(started, "expected observe_run_started event")
        self.assertTrue(stepped, "expected observe_step event")
        self.assertTrue(finished, "expected observe_run_finished event")
        self.assertEqual("success", finished[-1].get("outcome"))

    def test_run_command_observed_emits_cli_events(self):
        result = run_command_observed(
            ["python", "-c", "print('hello-observe')"],
            run_name="cli-test",
            open_ui_url="",
        )
        self.assertEqual(0, result["exit_code"])
        self.assertIn("run_id", result)

        started = tracing_module.read_trace_events(limit=200, filters={"event_type": "observe_run_started"})
        finished = tracing_module.read_trace_events(limit=200, filters={"event_type": "observe_run_finished"})
        self.assertTrue(started)
        self.assertTrue(finished)
        self.assertEqual("success", finished[-1].get("outcome"))

    def test_hive_module_alias_runs_trace_command(self):
        run_command_observed(["python", "-c", "pass"], run_name="hive-alias-test", open_ui_url="")

        buf = StringIO()
        with redirect_stdout(buf):
            code = hive_module.main(["trace", "ls", "--limit", "5"])
        self.assertEqual(0, code)
        self.assertIn("observe-run:", buf.getvalue())

    def test_checkpoint_and_rerun_sdk_events(self):
        with observe_run("checkpoint-test") as run_id:
            observe_checkpoint(run_id, "ckpt-1", step_name="stage.extract", state_ref="state://extract/1")
            observe_rerun_request(run_id, from_checkpoint_id="ckpt-1", reason="retry_with_override")

        checkpoint_events = tracing_module.read_trace_events(limit=200, filters={"event_type": "observe_checkpoint"})
        rerun_events = tracing_module.read_trace_events(limit=200, filters={"event_type": "observe_rerun_requested"})
        self.assertTrue(checkpoint_events)
        self.assertTrue(rerun_events)
        self.assertEqual("ckpt-1", (checkpoint_events[-1].get("payload") or {}).get("checkpoint_id"))
        self.assertEqual("ckpt-1", (rerun_events[-1].get("payload") or {}).get("from_checkpoint_id"))

    def test_trace_run_ls_and_open_flow(self):
        run_buf = StringIO()
        with redirect_stdout(run_buf):
            code = observe_main(["trace", "run", "--no-open", "--", "python", "-c", "print('trace-run')"])
        self.assertEqual(0, code)
        run_output = run_buf.getvalue()
        self.assertIn("run_id=", run_output)
        run_id = ""
        for line in run_output.splitlines():
            if line.startswith("run_id="):
                run_id = line.split()[0].split("=", 1)[1].strip()
                break
        self.assertTrue(run_id)

        ls_buf = StringIO()
        with redirect_stdout(ls_buf):
            ls_code = observe_main(["trace", "ls", "--limit", "5"])
        self.assertEqual(0, ls_code)
        self.assertIn(run_id, ls_buf.getvalue())

        open_buf = StringIO()
        with redirect_stdout(open_buf):
            open_code = observe_main(["trace", "open", run_id, "--ui-url", "http://localhost:5173"])
        self.assertEqual(0, open_code)
        self.assertIn("opened=http://localhost:5173", open_buf.getvalue())

        show_buf = StringIO()
        original_log_path = tracing_module.TRACE_LOG_PATH
        tracing_module.TRACE_LOG_PATH = str(Path(tempfile.gettempdir()) / f"hive-observe-mismatch-{uuid.uuid4().hex}.log")
        try:
            with redirect_stdout(show_buf):
                show_code = observe_main(["trace", "show", run_id, "--limit", "200"])
        finally:
            tracing_module.TRACE_LOG_PATH = original_log_path
        self.assertEqual(0, show_code)
        self.assertIn(run_id, show_buf.getvalue())

        show_json_buf = StringIO()
        with redirect_stdout(show_json_buf):
            show_json_code = observe_main(["trace", "show", run_id, "--limit", "200", "--json"])
        self.assertEqual(0, show_json_code)
        self.assertIn(f'"run_id":"{run_id}"', show_json_buf.getvalue())

    def test_quickstart_command(self):
        buf = StringIO()
        with redirect_stdout(buf):
            code = observe_main(["quickstart", "--no-open"])
        self.assertEqual(0, code)
        text = buf.getvalue()
        self.assertIn("quickstart_complete=true", text)
        self.assertIn("run_id=observe-run:", text)
        self.assertIn("next=hive trace ls", text)

    def test_doctor_command(self):
        buf = StringIO()
        with redirect_stdout(buf):
            code = observe_main(["doctor"])
        self.assertIn(code, (0, 2))
        text = buf.getvalue()
        self.assertIn("[", text)
        self.assertIn("python", text)
        self.assertIn("trace_home", text)

    def test_version_flag(self):
        buf = StringIO()
        with self.assertRaises(SystemExit) as ctx:
            with redirect_stdout(buf):
                observe_main(["--version"])
        self.assertEqual(0, int(ctx.exception.code))
        self.assertRegex(buf.getvalue().strip(), r"^hive\s+\S+$")

    def test_proxy_passthrough_emits_proxy_events(self):
        class UpstreamHandler(BaseHTTPRequestHandler):
            def do_POST(self):  # noqa: N802
                length = int(self.headers.get("Content-Length") or 0)
                body = self.rfile.read(length) if length > 0 else b""
                payload = json.dumps({"ok": True, "size": len(body)}).encode("utf-8")
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(payload)))
                self.end_headers()
                self.wfile.write(payload)

            def log_message(self, fmt, *args):  # noqa: A003
                return

        upstream = ThreadingHTTPServer(("127.0.0.1", 0), UpstreamHandler)
        thread = threading.Thread(target=upstream.serve_forever, daemon=True)
        thread.start()
        host, port = upstream.server_address
        upstream_base = f"http://{host}:{port}/v1"
        run_id = f"observe-run:{uuid.uuid4().hex[:8]}"
        try:
            with observe_proxy(run_id=run_id, upstream_base=upstream_base) as proxy:
                body = json.dumps({"hello": "proxy"}).encode("utf-8")
                req = urllib.request.Request(
                    url=f"{proxy.url}/chat/completions",
                    data=body,
                    headers={"Content-Type": "application/json"},
                    method="POST",
                )
                with urllib.request.urlopen(req, timeout=5) as resp:
                    data = json.loads(resp.read().decode("utf-8"))
                    self.assertTrue(data.get("ok"))
        finally:
            upstream.shutdown()
            upstream.server_close()

        req_events = tracing_module.read_trace_events(
            limit=200,
            filters={"run_id": run_id, "event_type": "observe_proxy_request"},
        )
        res_events = tracing_module.read_trace_events(
            limit=200,
            filters={"run_id": run_id, "event_type": "observe_proxy_response"},
        )
        self.assertTrue(req_events)
        self.assertTrue(res_events)

    def test_trace_run_proxy_mode_prints_proxy_url(self):
        class UpstreamHandler(BaseHTTPRequestHandler):
            def do_POST(self):  # noqa: N802
                payload = json.dumps({"ok": True}).encode("utf-8")
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(payload)))
                self.end_headers()
                self.wfile.write(payload)

            def log_message(self, fmt, *args):  # noqa: A003
                return

        upstream = ThreadingHTTPServer(("127.0.0.1", 0), UpstreamHandler)
        thread = threading.Thread(target=upstream.serve_forever, daemon=True)
        thread.start()
        host, port = upstream.server_address
        upstream_base = f"http://{host}:{port}/v1"
        try:
            buf = StringIO()
            with redirect_stdout(buf):
                code = observe_main(
                    [
                        "trace",
                        "run",
                        "--no-open",
                        "--proxy",
                        "--proxy-upstream",
                        upstream_base,
                        "--",
                        "python",
                        "-c",
                        "import os; print('base_set', bool(os.getenv('OPENAI_BASE_URL')))",
                    ]
                )
            self.assertEqual(0, code)
            out = buf.getvalue()
            self.assertIn("proxy=http://127.0.0.1:", out)
            self.assertIn("base_set True", out)
        finally:
            upstream.shutdown()
            upstream.server_close()


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