#!/usr/bin/env python3
"""
Hyperspell (Cascade) end-to-end smoke test.

This file is intentionally "new user" friendly:
- load `.env` from repo root (if python-dotenv is installed)
- configure tracing to send spans to Hyperspell
- generate a simple trace with a tool span
- verify via the HTTP API (list traces / fetch trace tree)
- exercise the OpenAI wrapper (runs 1 short completion)

Required env:
- CASCADE_API_KEY        (get from Hyperspell dashboard settings)
- OPENAI_API_KEY         (for the OpenAI wrapper test)

Recommended env:
- CASCADE_ENDPOINT=https://api-hyperspell.runcascade.com/v1/traces

Run as a script:
    python tests/test_hyperspell.py

Run with pytest:
    python -m pytest -q tests/test_hyperspell.py
"""

from __future__ import annotations

import json
import os
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Optional

import pytest  # type: ignore[import-not-found]

from cascade import (
    get_trace_tree,
    init_tracing,
    list_traces,
    tool,
    trace_session,
    trace_run,
    wrap_llm_client,
)

HYPERSPELL_DEFAULT_ENDPOINT = "https://api-hyperspell.runcascade.com/v1/traces"


def _load_dotenv_from_repo_root() -> None:
    try:
        from dotenv import load_dotenv  # type: ignore
    except ImportError:
        return

    env_path = Path(__file__).resolve().parents[1] / ".env"
    if env_path.exists():
        load_dotenv(env_path)


def _truthy_env(name: str) -> bool:
    return os.getenv(name, "").strip().lower() in {"1", "true", "yes", "y", "on"}


def _force_flush_best_effort() -> None:
    # The SDK uses BatchSpanProcessor; force_flush reduces "I ran it but see nothing" confusion.
    try:
        from opentelemetry import trace as otel_trace  # type: ignore

        provider = otel_trace.get_tracer_provider()
        force_flush = getattr(provider, "force_flush", None)
        if callable(force_flush):
            force_flush()
    except Exception:
        pass


def _base_url_from_traces_endpoint(endpoint: str) -> str:
    endpoint = (endpoint or "").strip()
    if endpoint.endswith("/v1/traces"):
        endpoint = endpoint[: -len("/v1/traces")]
    return endpoint.rstrip("/")


def _get_spans_since_raw(
    *,
    endpoint: str,
    api_key: str,
    trace_id: str,
    since_time: int,
) -> dict:
    """
    Call the backend live-poll endpoint directly.

    This is what the dashboard uses to decide whether a trace is still "live".
    """
    base = _base_url_from_traces_endpoint(endpoint)
    url = f"{base}/api/traces/{trace_id}/spans/since?{urllib.parse.urlencode({'since_time': since_time})}"
    req = urllib.request.Request(url, headers={"X-API-Key": api_key}, method="GET")
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            raw = resp.read().decode("utf-8")
            return json.loads(raw) if raw else {}
    except urllib.error.HTTPError as e:
        detail = ""
        try:
            detail = e.read().decode("utf-8")
        except Exception:
            detail = str(e)
        raise RuntimeError(f"get_spans_since failed: {e.code} {detail}") from e


def _wait_for_trace_completed(
    *,
    endpoint: str,
    api_key: str,
    trace_id: str,
    timeout_s: float = 30.0,
) -> None:
    deadline = time.time() + timeout_s
    last = None
    while time.time() < deadline:
        last = _get_spans_since_raw(endpoint=endpoint, api_key=api_key, trace_id=trace_id, since_time=0)
        if last.get("trace_completed") is True:
            return
        time.sleep(1.0)

    # Include a bit of context for debugging (root end_time is the key signal).
    root_end_time = None
    try:
        tree = get_trace_tree(trace_id)
        root_end_time = (tree or {}).get("root", {}).get("end_time")
    except Exception:
        pass

    raise AssertionError(
        "Trace never became completed for live polling "
        f"within {timeout_s:.0f}s. trace_completed={last.get('trace_completed') if isinstance(last, dict) else last}. "
        f"root.end_time={root_end_time}. "
        "This is why the dashboard keeps showing 'Loading new spans...'."
    )


@tool
def hello(name: str) -> str:
    return f"hello {name}"


@tool
def add(a: int, b: int) -> int:
    return a + b


def _init_hyperspell_tracing(project: str) -> tuple[str, str]:
    endpoint = os.getenv("CASCADE_ENDPOINT", HYPERSPELL_DEFAULT_ENDPOINT).strip()
    api_key = os.getenv("CASCADE_API_KEY", "").strip()
    if not api_key:
        pytest.fail("CASCADE_API_KEY not set; set it to run Hyperspell smoke test")

    init_tracing(
        project=project,
        endpoint=endpoint,
        api_key=api_key,
        metadata={"source": "tests/test_hyperspell.py"},
    )
    return endpoint, api_key


def _send_basic_trace(project: str) -> str:
    with trace_run("QuickstartRun", metadata={"project": project}) as root:
        assert root is not None, "init_tracing() failed; root span is None"
        trace_id = format(root.get_span_context().trace_id, "032x")

        msg = hello("world")
        assert msg == "hello world"
        assert add(2, 3) == 5

        return trace_id


def _send_session_traces(*, session_id: str) -> list[str]:
    """
    Mirrors the docs snippet:

        with trace_session("chat-123"):
            with trace_run("Turn", metadata={"turn": 1}): ...
            with trace_run("Turn", metadata={"turn": 2}): ...

    Returns:
        List of trace_ids for each Turn.
    """
    trace_ids: list[str] = []
    with trace_session(session_id):
        with trace_run("Turn", metadata={"turn": 1}) as run1:
            assert run1 is not None
            trace_ids.append(format(run1.get_span_context().trace_id, "032x"))

        with trace_run("Turn", metadata={"turn": 2}) as run2:
            assert run2 is not None
            trace_ids.append(format(run2.get_span_context().trace_id, "032x"))

    return trace_ids


def _verify_trace_via_api(project: str, trace_id: str, timeout_s: float = 15.0) -> None:
    """
    Best-effort: poll list_traces until the trace shows up, then fetch its tree.
    This can be eventually consistent; keep it short for a smoke test.
    """
    deadline = time.time() + timeout_s
    last_total: Optional[int] = None

    while time.time() < deadline:
        resp = list_traces(project, days=1, limit=50)
        traces = resp.get("traces", []) or []
        last_total = resp.get("total")
        if any(t.get("trace_id") == trace_id for t in traces):
            tree = get_trace_tree(trace_id)
            assert isinstance(tree, dict)
            return
        time.sleep(1.0)

    pytest.fail(
        f"Trace not visible via API within {timeout_s:.0f}s (total={last_total}). "
        "Check CASCADE_ENDPOINT/CASCADE_API_KEY and dashboard project."
    )


def _run_openai_wrapper_one_call() -> None:
    openai_key = os.getenv("OPENAI_API_KEY", "").strip()
    if not openai_key:
        pytest.fail("OPENAI_API_KEY not set (required for OpenAI wrapper smoke call)")

    try:
        from openai import OpenAI  # type: ignore
    except ImportError:
        pytest.fail("openai package not installed. Run: pip install openai")

    client = wrap_llm_client(OpenAI(api_key=openai_key))
    with trace_run("LLMRun", metadata={"llm_provider": "openai"}):
        resp = client.chat.completions.create(
            model=os.getenv("OPENAI_MODEL", "gpt-4o-mini"),
            messages=[{"role": "user", "content": "Say hi in one sentence."}],
        )
        text = resp.choices[0].message.content
        assert isinstance(text, str) and text.strip()


def test_hyperspell_tracing_smoke() -> None:
    _load_dotenv_from_repo_root()

    project = os.getenv("CASCADE_PROJECT", "my_first_hyperspell_project").strip() or "my_first_hyperspell_project"
    endpoint, api_key = _init_hyperspell_tracing(project=project)

    trace_id = _send_basic_trace(project=project)
    _force_flush_best_effort()

    _verify_trace_via_api(project=project, trace_id=trace_id, timeout_s=30.0)
    _wait_for_trace_completed(endpoint=endpoint, api_key=api_key, trace_id=trace_id, timeout_s=30.0)


def test_hyperspell_openai_wrapper() -> None:
    _load_dotenv_from_repo_root()

    project = os.getenv("CASCADE_PROJECT", "my_first_hyperspell_project").strip() or "my_first_hyperspell_project"
    _init_hyperspell_tracing(project=project)

    _run_openai_wrapper_one_call()
    _force_flush_best_effort()


def test_hyperspell_trace_session_groups_turns() -> None:
    _load_dotenv_from_repo_root()

    project = os.getenv("CASCADE_PROJECT", "my_first_hyperspell_project").strip() or "my_first_hyperspell_project"
    endpoint, api_key = _init_hyperspell_tracing(project=project)

    # Keep the "chat-123" shape from docs, but add a suffix so the test is isolated.
    session_id = f"chat-123-{int(time.time())}"

    trace_ids = _send_session_traces(session_id=session_id)
    _force_flush_best_effort()

    # Make sure the dashboard polling would consider each trace finished.
    for tid in trace_ids:
        _wait_for_trace_completed(endpoint=endpoint, api_key=api_key, trace_id=tid, timeout_s=30.0)

    # Verify the session_id filter works end-to-end.
    deadline = time.time() + 30.0
    while time.time() < deadline:
        resp = list_traces(project, days=1, session_id=session_id, limit=100)
        seen = {t.get("trace_id") for t in (resp.get("traces", []) or [])}
        if all(tid in seen for tid in trace_ids):
            return
        time.sleep(1.0)

    pytest.fail("Session grouping traces not visible via list_traces(session_id=...) within 30s")


def main() -> int:
    # Script entrypoint: run the same logic as the pytest tests, but with simple prints.
    _load_dotenv_from_repo_root()

    endpoint = os.getenv("CASCADE_ENDPOINT", HYPERSPELL_DEFAULT_ENDPOINT).strip()
    api_key = os.getenv("CASCADE_API_KEY", "").strip()
    project = os.getenv("CASCADE_PROJECT", "my_first_hyperspell_project").strip() or "my_first_hyperspell_project"

    print("=" * 64)
    print("Hyperspell (Cascade) Smoke Test")
    print("=" * 64)
    print(f"Endpoint: {endpoint}")
    print(f"Project : {project}")
    if not api_key:
        print("❌ CASCADE_API_KEY is not set")
        return 1

    init_tracing(project=project, endpoint=endpoint, api_key=api_key, metadata={"source": "tests/test_hyperspell.py"})
    trace_id = _send_basic_trace(project=project)
    _force_flush_best_effort()

    print(f"✅ Sent trace: {trace_id}")

    _verify_trace_via_api(project=project, trace_id=trace_id, timeout_s=30.0)
    print("✅ Verified trace is visible via API")

    _wait_for_trace_completed(endpoint=endpoint, api_key=api_key, trace_id=trace_id, timeout_s=30.0)
    print("✅ Verified trace is marked completed (dashboard should stop polling)")

    _run_openai_wrapper_one_call()
    _force_flush_best_effort()
    print("✅ OpenAI wrapper smoke call completed")

    session_id = f"chat-123-{int(time.time())}"
    trace_ids = _send_session_traces(session_id=session_id)
    _force_flush_best_effort()
    print(f"✅ Sent session traces: session_id={session_id} traces={len(trace_ids)}")

    for tid in trace_ids:
        _wait_for_trace_completed(endpoint=endpoint, api_key=api_key, trace_id=tid, timeout_s=30.0)
    print("✅ Session traces marked completed")

    deadline = time.time() + 30.0
    while time.time() < deadline:
        resp = list_traces(project, days=1, session_id=session_id, limit=100)
        seen = {t.get("trace_id") for t in (resp.get("traces", []) or [])}
        if all(tid in seen for tid in trace_ids):
            print("✅ Verified session grouping via list_traces(session_id=...)")
            break
        time.sleep(1.0)
    else:
        print("⚠️  Session traces not visible via session_id filter within 30s (may be eventual consistency)")

    print("Open dashboard at: https://hyperspell-dashboard.runcascade.com/")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())