"""
If snapshots need to be (re)generated, set GENERATE_SNAPSHOTS=1
and run the tests as usual.
"""

import asyncio
import os
import platform
import subprocess
from typing import Callable
from typing import Dict
from typing import Generator

import aiohttp
from ddtrace._trace.sampler import DatadogSampler
from ddtrace._trace.sampling_rule import SamplingRule
from ddtrace.profiling import Profiler
from ddtrace.propagation.http import HTTPPropagator
from ddtrace.trace import Tracer
import pytest


@pytest.fixture
def tracer(testagent_port: str, testagent: aiohttp.ClientSession) -> Tracer:
    tracer = Tracer()
    tracer._agent_url = f"http://localhost:{testagent_port}"
    return tracer


@pytest.fixture
def trace_sample_rate() -> float:
    return 1.0


@pytest.fixture
def stats_tracer(tracer: Tracer, trace_sample_rate: float) -> Generator[Tracer, None, None]:
    tracer.configure(compute_stats_enabled=True)
    tracer._sampler = DatadogSampler(rules=[SamplingRule(sample_rate=trace_sample_rate)])
    for processor in tracer._span_processors:
        if processor.__class__.__name__ == "SpanStatsProcessorV06":
            processor._agent_url = tracer._agent_url
    yield tracer


@pytest.mark.parametrize(
    "operation_name,service,resource,error,span_type,meta,metrics,response_code",
    [
        # First value is the reference data (also stored in the snapshot)
        ("root", "custom_service", "/url/endpoint", 0, "web", {}, {}, 200),
        ("root2", "custom_service", "/url/endpoint", 0, "web", {}, {}, 400),
        ("root", "custom_service2", "/url/endpoint", 0, "web", {}, {}, 400),
        ("root", "custom_service", "/url/endpoint/2", 0, "web", {}, {}, 400),
        ("root", "custom_service", "/url/endpoint", 1, "web", {}, {}, 400),
        ("root", "custom_service", "/url/endpoint", 0, "http", {}, {}, 400),
        (
            "root",
            "custom_service",
            "/url/endpoint",
            0,
            "http",
            {"meta": "value"},
            {},
            400,
        ),
        (
            "root",
            "custom_service",
            "/url/endpoint",
            0,
            "http",
            {},
            {"metrics": 2.3},
            400,
        ),
        ("root", "custom_service", "/url/endpoint", 0, "web", {}, {}, 200),
    ],
)
async def test_single_trace(
    testagent,
    testagent_url,
    tracer,
    operation_name,
    service,
    resource,
    error,
    span_type,
    meta,
    metrics,
    response_code,
):
    await testagent.get(f"{testagent_url}/test/session/start?test_session_token=test_single_trace")
    tracer = Tracer()
    tracer._agent_url = testagent_url
    with tracer.trace(operation_name, service=service, resource=resource, span_type=span_type) as span:
        if error is not None:
            span.error = error
        for k, v in meta.items():
            span.set_tag(k, v)
        for k, v in metrics.items():
            span.set_metric(k, v)
    tracer.shutdown()
    resp = await testagent.get(f"{testagent_url}/test/session/snapshot?test_session_token=test_single_trace")

    assert resp.status == response_code


async def test_multi_trace(testagent_url, testagent, tracer):
    await testagent.get(f"{testagent_url}/test/session/start?test_session_token=test_multi_trace")
    with tracer.trace("root0"):
        with tracer.trace("child0"):
            pass
    with tracer.trace("root1"):
        with tracer.trace("child1"):
            pass
    tracer.flush()
    resp = await testagent.get(f"{testagent_url}/test/session/snapshot?test_session_token=test_multi_trace")
    assert resp.status == 200

    # Run the snapshot test again.
    await testagent.get(f"{testagent_url}/test/session/start?test_session_token=test_multi_trace")
    with tracer.trace("root0"):
        with tracer.trace("child0"):
            pass
    with tracer.trace("root1"):
        with tracer.trace("child1"):
            pass
    tracer.flush()
    resp = await testagent.get(f"{testagent_url}/test/session/snapshot?test_session_token=test_multi_trace")
    assert resp.status == 200

    # Simulate a failed snapshot with a missing trace.
    await testagent.get(f"{testagent_url}/test/session/start?test_session_token=test_multi_trace")
    with tracer.trace("root0"):
        with tracer.trace("child0"):
            pass
    tracer.flush()
    resp = await testagent.get(f"{testagent_url}/test/session/snapshot?test_session_token=test_multi_trace")
    assert resp.status == 400
    tracer.shutdown()


async def test_trace_distributed_same_payload(testagent_url, testagent, tracer):
    await testagent.get(f"{testagent_url}/test/session/start?test_session_token=test_trace_distributed_same_payload")
    with tracer.trace("root0"):
        with tracer.trace("child0") as span:
            ctx = span.context

    tracer.context_provider.activate(ctx)
    with tracer.trace("root1"):
        with tracer.trace("child1"):
            pass
    tracer.flush()
    resp = await testagent.get(
        f"{testagent_url}/test/session/snapshot?test_session_token=test_trace_distributed_same_payload"
    )
    assert resp.status == 200


async def test_trace_distributed_propagated(testagent_url, testagent, tracer):
    await testagent.get(f"{testagent_url}/test/session/start?test_session_token=test_trace_distributed_propagated")
    headers = {
        "x-datadog-trace-id": "1234",
        "x-datadog-parent-id": "5678",
    }
    context = HTTPPropagator.extract(headers)
    tracer.context_provider.activate(context)

    with tracer.trace("root"):
        with tracer.trace("child"):
            pass
    tracer.flush()
    resp = await testagent.get(
        f"{testagent_url}/test/session/snapshot?test_session_token=test_trace_distributed_propagated"
    )

    assert resp.status == 200, await resp.text()


async def test_trace_missing_received(testagent_url, testagent, tracer):
    resp = await testagent.get(f"{testagent_url}/test/session/start?test_session_token=test_trace_missing_received")
    assert resp.status == 200, await resp.text()

    with tracer.trace("root0"):
        with tracer.trace("child0"):
            pass
    tracer.flush()
    resp = await testagent.get(f"{testagent_url}/test/session/snapshot?test_session_token=test_trace_missing_received")
    assert resp.status == 200

    # Do another snapshot without sending any traces.
    resp = await testagent.get(f"{testagent_url}/test/session/start?test_session_token=test_trace_missing_received")
    assert resp.status == 200, await resp.text()
    resp = await testagent.get(f"{testagent_url}/test/session/snapshot?test_session_token=test_trace_missing_received")
    assert resp.status == 400


def _tracestats_traces(tracer: Tracer) -> None:
    for i in range(5):
        with tracer.trace("http.request", resource="/users/view") as span:
            if i == 4:
                span.error = 1


def _tracestats_traces_no_error(tracer: Tracer) -> None:
    for i in range(5):
        with tracer.trace("http.request", resource="/users/view"):
            pass


def _tracestats_traces_missing_trace(tracer: Tracer) -> None:
    for i in range(4):
        with tracer.trace("http.request", resource="/users/view") as span:
            if i == 3:
                span.error = 1


def _tracestats_traces_extra_trace(tracer: Tracer) -> None:
    _tracestats_traces(tracer)
    with tracer.trace("http.request", resource="/users/list"):
        pass


@pytest.mark.parametrize("trace_sample_rate", [0.0])  # Don't send any traces
@pytest.mark.parametrize(
    "do_traces,fail",
    [
        (
            _tracestats_traces,
            False,
        ),  # Keep this parametrization first as the next ones depend on it.
        (_tracestats_traces_no_error, True),
        (_tracestats_traces_missing_trace, True),
        (_tracestats_traces_extra_trace, True),
    ],
)
async def test_tracestats(
    testagent_url: str,
    testagent: aiohttp.ClientSession,
    stats_tracer: Tracer,
    testagent_snapshot_ci_mode: bool,
    trace_sample_rate: float,
    do_traces: Callable[[Tracer], None],
    fail: bool,
) -> None:
    await testagent.get(f"{testagent_url}/test/session/start?test_session_token=test_trace_stats")
    do_traces(stats_tracer)
    stats_tracer.shutdown()  # force out the stats
    resp = await testagent.get(f"{testagent_url}/test/session/snapshot?test_session_token=test_trace_stats")
    if fail:
        assert resp.status == 400
    else:
        assert resp.status == 200, await resp.text()


async def test_cmd(testagent_url: str, testagent: aiohttp.ClientSession, tracer: Tracer) -> None:
    """Test the commands provided with the library.

    Note that this test reuses the trace/snapshot from test_single_trace above.
    """
    env = os.environ.copy()
    env["DD_TRACE_AGENT_URL"] = testagent_url
    p = subprocess.run(
        ["ddapm-test-agent-session-start", "--test-session-token=test_single_trace"],
        env=env,
    )
    assert p.returncode == 0
    with tracer.trace("root", service="custom_service", resource="/url/endpoint", span_type="web"):
        pass
    tracer.flush()
    p = subprocess.run(["ddapm-test-agent-snapshot", "--test-session-token=test_single_trace"], env=env)
    assert p.returncode == 0

    # Ensure failing snapshots work.
    p = subprocess.run(
        ["ddapm-test-agent-session-start", "--test-session-token=test_single_trace"],
        env=env,
    )
    assert p.returncode == 0
    with tracer.trace("root1234", service="custom_service", resource="/url/endpoint", span_type="web"):
        pass
    tracer.shutdown()
    p = subprocess.run(["ddapm-test-agent-snapshot", "--test-session-token=test_single_trace"], env=env)
    assert p.returncode == 1


@pytest.mark.skipif(platform.system() == "Windows", reason="Profiling broken on windows")
async def test_profiling_endpoint(testagent_url: str, testagent: aiohttp.ClientSession, tracer: Tracer) -> None:
    p = Profiler(tracer=tracer)
    p.start()
    p.stop(flush=True)
    resp = await testagent.get(f"{testagent_url}/test/session/requests")
    assert resp.status == 200
    data = await resp.json()
    assert len(data) >= 1
    assert data[-1]["url"].endswith("/profiling/v1/input")


async def test_race_condition(
    testagent: aiohttp.ClientSession,
    testagent_port: int,
    v04_reference_http_trace_payload_headers: Dict[str, str],
    v04_reference_http_trace_payload_data: bytes,
    testagent_url: str,
) -> None:
    """
    Reproduction of: "RuntimeError: readany() called while another coroutine is already waiting for incoming data"
    when trace requests are made and being read simultaneously
    """
    reqs = []
    for i in range(50):
        reqs.append(
            testagent.put(
                testagent_url + "/v0.4/traces",
                headers=v04_reference_http_trace_payload_headers,
                data=v04_reference_http_trace_payload_data,
            )
        )
        reqs.append(testagent.get(testagent_url + "/test/session/traces"))

    resps = await asyncio.gather(*reqs)
    for r in resps:
        assert r.status == 200
