import json
import threading

import requests
from flask import Flask, Response, request
from opentelemetry import trace
from opentelemetry.instrumentation.flask import FlaskInstrumentor
from opentelemetry.instrumentation.requests import RequestsInstrumentor
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor
from werkzeug.serving import make_server


resource = Resource.create({"service.name": "http-propagation-demo"})
provider = TracerProvider(resource=resource)
provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))
trace.set_tracer_provider(provider)

app_a = Flask("service-a")
app_b = Flask("service-b")
FlaskInstrumentor().instrument_app(app_a)
FlaskInstrumentor().instrument_app(app_b)
RequestsInstrumentor().instrument()


def current_span_ids():
    span = trace.get_current_span()
    context = span.get_span_context()
    parent = getattr(span, "parent", None)
    return {
        "trace_id": f"{context.trace_id:032x}",
        "span_id": f"{context.span_id:016x}",
        "parent_span_id": f"{parent.span_id:016x}" if parent else None,
    }


def json_response(payload):
    return Response(
        json.dumps(payload, indent=2, sort_keys=True) + "\n",
        mimetype="application/json",
    )


@app_b.get("/work")
def service_b_work():
    span = current_span_ids()
    return json_response(
        {
            "service": "service-b",
            "received_traceparent": request.headers.get("traceparent", ""),
            "trace_id": span["trace_id"],
            "span_id": span["span_id"],
            "parent_span_id": span["parent_span_id"],
        }
    )


@app_a.get("/call-b")
def service_a_call_b():
    service_a = current_span_ids()
    downstream = requests.get("http://127.0.0.1:8082/work", timeout=5).json()
    traceparent_parts = downstream["received_traceparent"].split("-")
    traceparent_parent_span_id = traceparent_parts[2] if len(traceparent_parts) == 4 else ""
    return json_response(
        {
            "same_trace": service_a["trace_id"] == downstream["trace_id"],
            "service_a_trace_id": service_a["trace_id"],
            "service_b_trace_id": downstream["trace_id"],
            "traceparent_header": downstream["received_traceparent"],
            "traceparent_parent_span_id": traceparent_parent_span_id,
            "service_b_parent_span_id": downstream["parent_span_id"],
            "parent_matches_traceparent": downstream["parent_span_id"]
            == traceparent_parent_span_id,
        }
    )


def serve(app, port):
    server = make_server("127.0.0.1", port, app)
    server.serve_forever()


if __name__ == "__main__":
    threading.Thread(target=serve, args=(app_b, 8082), daemon=True).start()
    threading.Thread(target=serve, args=(app_a, 8081), daemon=True).start()
    print("service-a listening on http://127.0.0.1:8081/call-b", flush=True)
    print("service-b listening on http://127.0.0.1:8082/work", flush=True)
    threading.Event().wait()