"""ServiceBridge middleware for Flask.

Usage::

    from flask import Flask
    from service_bridge.http.flask import init_servicebridge

    app = Flask(__name__)
    init_servicebridge(app, svc)
"""

from __future__ import annotations

import threading
from typing import TYPE_CHECKING

from flask import Flask, g, request as flask_request

if TYPE_CHECKING:
    from service_bridge.servicebridge import ServiceBridge


def init_servicebridge(
    app: Flask,
    client: "ServiceBridge",
    *,
    exclude_paths: tuple[str, ...] = (),
    propagate_trace_header: bool = True,
    auto_register: bool = True,
) -> None:
    """Register ServiceBridge before/after request hooks with a Flask app.

    After calling this function, ``g.service_bridge`` will be available inside
    every request handler, and ``g.trace_id`` / ``g.span_id`` carry the current
    trace context.

    Args:
        app: Flask application instance.
        client: ServiceBridge SDK client.
        exclude_paths: Path prefixes to skip tracing.
        propagate_trace_header: Write ``x-trace-id`` to response (default: True).
        auto_register: Register routes in the catalog on first hit (default: True).
    """
    registered: set[str] = set()
    reg_lock = threading.Lock()

    @app.before_request
    def _before() -> None:
        from service_bridge.http_catalog import extract_trace_from_headers

        path = flask_request.path or "/"
        g.service_bridge = client

        for prefix in exclude_paths:
            if path.startswith(prefix):
                return

        headers = dict(flask_request.headers)
        trace_id, parent_span_id = extract_trace_from_headers(headers)

        span = client.start_http_span(
            method=flask_request.method,
            path=path,
            trace_id=trace_id,
            parent_span_id=parent_span_id,
        )

        g.trace_id = span.trace_id
        g.span_id = span.span_id
        g._sb_span = span

    @app.after_request
    def _after(response):
        span = getattr(g, "_sb_span", None)
        if span is None:
            return response

        if propagate_trace_header:
            response.headers["x-trace-id"] = span.trace_id

        span.end(status_code=response.status_code)

        if auto_register:
            rule = flask_request.url_rule
            pattern = rule.rule if rule else flask_request.path
            key = f"{flask_request.method}:{pattern}"
            with reg_lock:
                if key not in registered:
                    registered.add(key)
                    method = flask_request.method

                    def _register(m=method, r=pattern, k=key) -> None:
                        try:
                            client.register_http_endpoint(m, r)
                        except Exception:
                            with reg_lock:
                                registered.discard(k)

                    threading.Thread(target=_register, daemon=True).start()

        return response
