"""Shared fixtures for integration tests."""

from __future__ import annotations

import contextlib
import gzip
import json
import os
import threading
import time
from http.server import BaseHTTPRequestHandler, HTTPServer
from typing import TYPE_CHECKING

import pytest
from google.protobuf.json_format import MessageToDict
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
    ExportTraceServiceRequest,
)

if TYPE_CHECKING:
    from collections.abc import Callable, Iterator

import flowlines

# ---------------------------------------------------------------------------
# Skip helpers for API keys
# ---------------------------------------------------------------------------


def _skip_unless_openai() -> None:
    """Skip the current test unless the OpenAI API key and package are available."""
    if not os.environ.get("OPENAI_API_KEY"):
        pytest.skip("OPENAI_API_KEY not set")
    try:
        import openai  # noqa: F401
    except ImportError:
        pytest.skip("openai not installed")


def _skip_unless_anthropic() -> None:
    """Skip the current test unless the Anthropic API key and package are available."""
    if not os.environ.get("ANTHROPIC_API_KEY"):
        pytest.skip("ANTHROPIC_API_KEY not set")
    try:
        import anthropic  # noqa: F401
    except ImportError:
        pytest.skip("anthropic not installed")


def _skip_unless_gemini() -> None:
    """Skip the current test unless the Gemini API key and package are available."""
    if not os.environ.get("GEMINI_API_KEY"):
        pytest.skip("GEMINI_API_KEY not set")
    try:
        import google.genai  # noqa: F401
    except ImportError:
        pytest.skip("google-genai not installed")


def _skip_unless_cohere() -> None:
    """Skip the current test unless the Cohere API key and package are available."""
    if not os.environ.get("CO_API_KEY"):
        pytest.skip("CO_API_KEY not set")
    try:
        import cohere  # noqa: F401
    except ImportError:
        pytest.skip("cohere not installed")


# ---------------------------------------------------------------------------
# In-process OTLP capture server
# ---------------------------------------------------------------------------


class OTLPCaptureServer:
    """Threaded HTTP server that captures OTLP trace export requests."""

    def __init__(self) -> None:
        self._records: list[dict] = []
        self._lock = threading.Lock()

        # Create a per-instance handler class to avoid shared state
        server_ref = self

        class _Handler(BaseHTTPRequestHandler):
            def do_POST(self) -> None:  # noqa: N802
                if self.path != "/v1/traces":
                    self.send_response(404)
                    self.end_headers()
                    return

                content_length = int(self.headers.get("Content-Length", 0))
                raw_body = self.rfile.read(content_length)

                if self.headers.get("Content-Encoding") == "gzip":
                    raw_body = gzip.decompress(raw_body)

                content_type = self.headers.get("Content-Type", "")
                if "application/json" in content_type:
                    body: dict = json.loads(raw_body)
                else:
                    req = ExportTraceServiceRequest()
                    req.ParseFromString(raw_body)
                    body = MessageToDict(req, preserving_proto_field_name=True)

                headers = {k: v for k, v in self.headers.items()}
                record = {
                    "headers": headers,
                    "body": body,
                }

                with server_ref._lock:
                    server_ref._records.append(record)

                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(b"{}")

            def log_message(self, format: str, *args: object) -> None:  # noqa: A002
                pass

        self._server = HTTPServer(("127.0.0.1", 0), _Handler)
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()

    @property
    def port(self) -> int:
        return self._server.server_address[1]

    @property
    def endpoint(self) -> str:
        return f"http://127.0.0.1:{self.port}"

    def get_records(self) -> list[dict]:
        with self._lock:
            return list(self._records)

    def get_spans(self) -> list[dict]:
        """Extract a flat list of spans from all received records."""
        spans: list[dict] = []
        for record in self.get_records():
            for rs in record["body"].get("resource_spans", []):
                for ss in rs.get("scope_spans", []):
                    spans.extend(ss.get("spans", []))
        return spans

    def wait_for_spans(self, min_count: int = 1, timeout: float = 30.0) -> list[dict]:
        """Poll until at least *min_count* spans have been received."""
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            spans = self.get_spans()
            if len(spans) >= min_count:
                return spans
            time.sleep(0.25)
        return self.get_spans()

    def shutdown(self) -> None:
        self._server.shutdown()
        self._thread.join(timeout=5)


# ---------------------------------------------------------------------------
# capture_server factory fixture
# ---------------------------------------------------------------------------


@pytest.fixture()
def capture_server() -> Iterator[Callable[[], OTLPCaptureServer]]:
    """Factory fixture: call it to create a new capture server.

    All servers created during the test are shut down automatically.
    """
    servers: list[OTLPCaptureServer] = []

    def _factory() -> OTLPCaptureServer:
        srv = OTLPCaptureServer()
        servers.append(srv)
        return srv

    yield _factory

    for srv in servers:
        srv.shutdown()


# ---------------------------------------------------------------------------
# OTEL state reset (runs after each test)
# ---------------------------------------------------------------------------


@pytest.fixture(autouse=True)
def _reset_otel_state() -> Iterator[None]:
    """Reset all OpenTelemetry / Traceloop / Flowlines global state after each test."""
    yield

    # 1. Reset Flowlines singleton
    from flowlines._init import _reset as _flowlines_reset

    _flowlines_reset()

    # 2. Shut down the active TracerProvider if it's an SDK provider
    from opentelemetry import trace
    from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider

    provider = trace.get_tracer_provider()
    if isinstance(provider, SDKTracerProvider):
        with contextlib.suppress(Exception):
            provider.shutdown()

    # 3. Reset OTel global tracer provider
    from opentelemetry.util._once import Once

    trace._TRACER_PROVIDER = None  # type: ignore[attr-defined]
    trace._TRACER_PROVIDER_SET_ONCE = Once()  # type: ignore[attr-defined]

    # 4. Reset Traceloop's TracerWrapper singleton
    try:
        from traceloop.sdk.tracing.tracing import TracerWrapper

        if hasattr(TracerWrapper, "instance"):
            del TracerWrapper.instance
        TracerWrapper.endpoint = None  # type: ignore[assignment]
    except ImportError:
        pass

    # 5. Uninstrument OpenAI instrumentor
    try:
        from opentelemetry.instrumentation.openai import OpenAIInstrumentor

        OpenAIInstrumentor().uninstrument()
    except Exception:
        pass

    # 6. Uninstrument Anthropic instrumentor
    try:
        from opentelemetry.instrumentation.anthropic import AnthropicInstrumentor

        AnthropicInstrumentor().uninstrument()
    except Exception:
        pass

    # 7. Uninstrument Google Generative AI instrumentor
    try:
        from opentelemetry.instrumentation.google_generativeai import (
            GoogleGenerativeAiInstrumentor,
        )

        GoogleGenerativeAiInstrumentor().uninstrument()
    except Exception:
        pass

    # 8. Uninstrument Cohere instrumentor
    try:
        from opentelemetry.instrumentation.cohere import CohereInstrumentor

        CohereInstrumentor().uninstrument()
    except Exception:
        pass


# ---------------------------------------------------------------------------
# Helper: simple OpenAI call
# ---------------------------------------------------------------------------


def make_openai_call(
    *,
    user_id: str | None = None,
    session_id: str | None = None,
    agent_id: str | None = None,
) -> str:
    """Make a minimal OpenAI chat completion call and return the response text.

    When *user_id* or *session_id* are provided, the call is wrapped in
    ``flowlines.context()`` so that the resulting spans carry those attributes.
    """
    from openai import OpenAI

    client = OpenAI()

    def _call() -> str:
        response = client.chat.completions.create(
            model="gpt-4.1-nano",
            messages=[{"role": "user", "content": "Say 'hello' and nothing else."}],
        )
        return response.choices[0].message.content or ""

    if user_id is not None or session_id is not None:
        with flowlines.context(
            user_id=user_id or "",
            session_id=session_id or "",
            agent_id=agent_id,
        ):
            return _call()
    return _call()


def make_openai_responses_call(
    *,
    user_id: str | None = None,
    session_id: str | None = None,
    agent_id: str | None = None,
) -> str:
    """Make a minimal OpenAI Responses API call and return the response text.

    When *user_id* or *session_id* are provided, the call is wrapped in
    ``flowlines.context()`` so that the resulting spans carry those attributes.
    """
    from openai import OpenAI

    client = OpenAI()

    def _call() -> str:
        response = client.responses.create(
            model="gpt-4.1-nano",
            input=[{"role": "user", "content": "Say 'hello' and nothing else."}],
        )
        return response.output_text or ""

    if user_id is not None or session_id is not None:
        with flowlines.context(
            user_id=user_id or "",
            session_id=session_id or "",
            agent_id=agent_id,
        ):
            return _call()
    return _call()


# ---------------------------------------------------------------------------
# Helper: simple Anthropic call
# ---------------------------------------------------------------------------


def make_anthropic_call(
    *,
    user_id: str | None = None,
    session_id: str | None = None,
    agent_id: str | None = None,
) -> str:
    """Make a minimal Anthropic message call and return the response text.

    When *user_id* or *session_id* are provided, the call is wrapped in
    ``flowlines.context()`` so that the resulting spans carry those attributes.
    """
    from anthropic import Anthropic

    client = Anthropic()

    def _call() -> str:
        response = client.messages.create(
            model="claude-haiku-4-5",
            max_tokens=100,
            messages=[{"role": "user", "content": "Say 'hello' and nothing else."}],
        )
        return response.content[0].text if response.content else ""

    if user_id is not None or session_id is not None:
        with flowlines.context(
            user_id=user_id or "",
            session_id=session_id or "",
            agent_id=agent_id,
        ):
            return _call()
    return _call()


# ---------------------------------------------------------------------------
# Helper: simple Gemini call
# ---------------------------------------------------------------------------


def make_gemini_call(
    *,
    user_id: str | None = None,
    session_id: str | None = None,
    agent_id: str | None = None,
) -> str:
    """Make a minimal Google Gemini call and return the response text.

    When *user_id* or *session_id* are provided, the call is wrapped in
    ``flowlines.context()`` so that the resulting spans carry those attributes.
    """
    from google import genai

    client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

    def _call() -> str:
        response = client.models.generate_content(
            model="gemini-2.5-flash-lite",
            contents="Say 'hello' and nothing else.",
        )
        return response.text or ""

    if user_id is not None or session_id is not None:
        with flowlines.context(
            user_id=user_id or "",
            session_id=session_id or "",
            agent_id=agent_id,
        ):
            return _call()
    return _call()


def make_gemini_chat_call(
    *,
    user_id: str | None = None,
    session_id: str | None = None,
    agent_id: str | None = None,
) -> str:
    """Make a Gemini call via the Chat API and return the response text.

    Unlike ``make_gemini_call`` (which passes ``contents`` as a plain string),
    the Chat API always passes ``contents`` as a list of ``Content`` objects.
    This exercises the list-of-contents code path in the instrumentor.

    When *user_id* or *session_id* are provided, the call is wrapped in
    ``flowlines.context()`` so that the resulting spans carry those attributes.
    """
    from google import genai

    client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
    chat = client.chats.create(model="gemini-2.5-flash-lite")

    def _call() -> str:
        response = chat.send_message("Say 'hello' and nothing else.")
        return response.text or ""

    if user_id is not None or session_id is not None:
        with flowlines.context(
            user_id=user_id or "",
            session_id=session_id or "",
            agent_id=agent_id,
        ):
            return _call()
    return _call()


def make_gemini_chat_stream_call(
    *,
    user_id: str | None = None,
    session_id: str | None = None,
    agent_id: str | None = None,
) -> str:
    """Make a streaming Gemini call via the Chat API and return the response.

    Uses ``chat.send_message_stream()`` and consumes every chunk, which
    mirrors the typical real-world usage pattern.  Like
    ``make_gemini_chat_call`` this exercises the list-of-contents code path.

    When *user_id* or *session_id* are provided, the call is wrapped in
    ``flowlines.context()`` so that the resulting spans carry those attributes.
    """
    from google import genai

    client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
    chat = client.chats.create(model="gemini-2.5-flash-lite")

    def _call() -> str:
        text = ""
        for chunk in chat.send_message_stream("Say 'hello' and nothing else."):
            text += chunk.text or ""
        return text

    if user_id is not None or session_id is not None:
        with flowlines.context(
            user_id=user_id or "",
            session_id=session_id or "",
            agent_id=agent_id,
        ):
            return _call()
    return _call()


# ---------------------------------------------------------------------------
# Helper: simple Cohere rerank call
# ---------------------------------------------------------------------------


def make_cohere_rerank_call(
    *,
    user_id: str | None = None,
    session_id: str | None = None,
    agent_id: str | None = None,
) -> str:
    """Make a minimal Cohere rerank call and return the top-ranked document text.

    When *user_id* or *session_id* are provided, the call is wrapped in
    ``flowlines.context()`` so that the resulting spans carry those attributes.
    """
    import cohere

    client = cohere.ClientV2()

    documents = [
        "Carson City is the capital of Nevada.",
        "Washington, D.C. is the capital of the United States.",
        "The Commonwealth of the Northern Mariana Islands is a U.S. territory.",
        "Sacramento is the capital of California.",
    ]

    def _call() -> str:
        response = client.rerank(
            model="rerank-v3.5",
            query="What is the capital of the United States?",
            documents=documents,
        )
        return documents[response.results[0].index] if response.results else ""

    if user_id is not None or session_id is not None:
        with flowlines.context(
            user_id=user_id or "",
            session_id=session_id or "",
            agent_id=agent_id,
        ):
            return _call()
    return _call()
