"""
Request tracking middleware for ASGI frameworks (Starlette, FastAPI).

This module provides middleware that automatically generates and tracks
request IDs and session IDs, making them available throughout the request
lifecycle via context variables.

Usage with Starlette/FastAPI:
    from a2a_llm_tracker import TrackerMiddleware

    app = FastAPI()
    app.add_middleware(TrackerMiddleware)

    # Or with options:
    app.add_middleware(
        TrackerMiddleware,
        generate_request_id=True,
        generate_session_id=True,
        request_id_header="X-Request-ID",
        session_id_header="X-Session-ID",
    )

Access the IDs anywhere in your code:
    from a2a_llm_tracker import get_request_id, get_session_id

    request_id = get_request_id()
    session_id = get_session_id()
"""
from __future__ import annotations

import re
import uuid
from contextvars import ContextVar
from typing import Callable, Optional
from ccs import LocalTransaction, GetTheConcept

# UUID regex pattern
_UUID_PATTERN = re.compile(
    r'^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$',
    re.IGNORECASE
)


def _is_uuid(value: str) -> bool:
    """Check if a string is a valid UUID."""
    return bool(_UUID_PATTERN.match(value))


def _is_concept_id(value: str) -> bool:
    """Check if a string is a concept ID (integer)."""
    try:
        int(value)
        return True
    except (ValueError, TypeError):
        return False
# Context variables for request tracking - isolated per request
_request_id_var: ContextVar[str] = ContextVar("llm_tracker_request_id", default="")
_session_id_var: ContextVar[str] = ContextVar("llm_tracker_session_id", default="")
_call_id_var: ContextVar[str] = ContextVar("llm_tracker_call_id", default="")
_parent_call_id_var: ContextVar[str] = ContextVar("llm_tracker_parent_call_id", default="")
_sequence_number_var: ContextVar[str] = ContextVar("llm_tracker_sequence_number", default="")


def get_request_id() -> str:
    """Get current request ID from context."""
    return _request_id_var.get()


def get_session_id() -> str:
    """Get current session ID from context."""
    return _session_id_var.get()


def set_request_id(request_id: str) -> None:
    """Set request ID in context (sync version - sets the raw ID)."""
    _request_id_var.set(request_id)


async def set_request_id_async(request_id: str, user_id: int = 101) -> str:
    """
    Set request ID in context, converting UUID to concept ID if needed.

    This function checks if the request_id is:
    - A UUID: Creates a new 'the_llm_request' concept and uses its ID
    - A concept ID (integer): Fetches the existing concept and uses its ID
    - Unknown format: Creates a new concept and uses its ID

    Args:
        request_id: The request ID (UUID, concept ID, or other format)
        user_id: User ID for creating concepts (default: 101)

    Returns:
        The concept ID as a string
    """
    tx = LocalTransaction()
    await tx.initialize()

    try:
        if _is_uuid(str(request_id)):
            # It's a UUID - create a new request concept
            trace_concept = await tx.MakeTheInstanceConceptLocal(
                "the_llm_request",
                str(request_id),
                False,  # isLocal=False - create new
                userId=user_id,
            )
        elif _is_concept_id(str(request_id)):
            # It's a concept ID - fetch the existing concept
            trace_concept = await GetTheConcept(int(request_id))
        else:
            # Unknown format - create as new concept
            trace_concept = await tx.MakeTheInstanceConceptLocal(
                "the_llm_request",
                str(request_id),
                False,
                userId=user_id,
            )

        await tx.commitTransaction()
        newConcept = await GetTheConcept(trace_concept.id)
        # Set the concept ID in context (not the original UUID)
        if newConcept and hasattr(newConcept, 'id'):
            concept_id = str(newConcept.id)
            _request_id_var.set(concept_id)
            return concept_id
        else:
            # Fallback to original if no concept was created
            _request_id_var.set(request_id)
            return request_id

    except Exception as e:
        await tx.rollbackTransaction()
        # On error, fall back to setting the original request_id
        _request_id_var.set(request_id)
        raise


def set_session_id(session_id: str) -> None:
    """Set session ID in context."""
    _session_id_var.set(session_id)


def get_call_id() -> str:
    """Get current call ID from context."""
    return _call_id_var.get()


def set_call_id(call_id: str) -> None:
    """Set call ID in context."""
    _call_id_var.set(call_id)


def get_parent_call_id() -> str:
    """Get current parent call ID from context."""
    return _parent_call_id_var.get()


def set_parent_call_id(parent_call_id: str) -> None:
    """Set parent call ID in context."""
    _parent_call_id_var.set(parent_call_id)


def get_sequence_number() -> str:
    """Get current sequence number from context."""
    return _sequence_number_var.get()


def set_sequence_number(sequence_number: str) -> None:
    """Set sequence number in context."""
    _sequence_number_var.set(sequence_number)


def generate_id() -> str:
    """Generate a new unique ID (UUID format)."""
    return str(uuid.uuid4())


async def generate_request_id_async(user_id: int = 101) -> str:
    """
    Generate a new request ID by creating a concept in CCS.

    This function creates a new 'the_llm_request' concept with a UUID value
    and returns the concept ID directly. Use this instead of generate_id()
    when you want to avoid the UUID-to-concept-ID conversion later.

    Args:
        user_id: User ID for creating the concept (default: 101)

    Returns:
        The concept ID as a string (not UUID)
    """
    tx = LocalTransaction()
    await tx.initialize()

    try:
        # Generate a UUID as the concept value
        request_uuid = str(uuid.uuid4())

        # Create the concept
        trace_concept = await tx.MakeTheInstanceConceptLocal(
            "the_llm_request",
            request_uuid,
            False,  # isLocal=False - create new
            userId=user_id,
        )

        await tx.commitTransaction()

        # Fetch the committed concept to get the real ID
        new_concept = await GetTheConcept(trace_concept.id)

        if new_concept and hasattr(new_concept, 'id'):
            concept_id = str(new_concept.id)
            _request_id_var.set(concept_id)
            return concept_id
        else:
            # Fallback to UUID if concept creation failed
            _request_id_var.set(request_uuid)
            return request_uuid

    except Exception:
        await tx.rollbackTransaction()
        # On error, generate a UUID as fallback
        fallback_id = str(uuid.uuid4())
        _request_id_var.set(fallback_id)
        raise




    


# Try to import Starlette - it's optional
try:
    from starlette.middleware.base import BaseHTTPMiddleware
    from starlette.requests import Request
    from starlette.responses import Response

    class TrackerMiddleware(BaseHTTPMiddleware):
        """
        Middleware that tracks request and session IDs for LLM cost tracking.

        This middleware:
        1. Extracts or generates request/session IDs from headers
        2. Sets them in context variables (accessible via get_request_id/get_session_id)
        3. Also sets them in the tracker context (for automatic inclusion in analyze_response)
        4. Adds the IDs to response headers

        Args:
            app: The ASGI application
            generate_request_id: If True, generate request ID when not in header (default: True)
            generate_session_id: If True, generate session ID when not in header (default: True)
            request_id_header: Header name for request ID (default: "X-Request-ID")
            trace_id_header: Header name for trace ID (default: "X-Trace-ID")
            session_id_header: Header name for session ID (default: "X-Session-ID")
            id_generator: Custom ID generator function (default: uuid4)
        """

        def __init__(
            self,
            app,
            generate_request_id: bool = True,
            generate_session_id: bool = True,
            generate_call_id: bool = True,
            request_id_header: str = "X-Request-ID",
            trace_id_header: str = "X-Trace-ID",
            session_id_header: str = "X-Session-ID",
            call_id_header: str = "X-Call-ID",
            parent_call_id_header: str = "X-Parent-Call-ID",
            sequence_number_header: str = "X-Sequence-Number",
            id_generator: Optional[Callable[[], str]] = None,
            user_id: int = 101,
        ):
            super().__init__(app)
            self.generate_request_id = generate_request_id
            self.generate_session_id = generate_session_id
            self.generate_call_id = generate_call_id
            self.request_id_header = request_id_header
            self.trace_id_header = trace_id_header
            self.session_id_header = session_id_header
            self.call_id_header = call_id_header
            self.parent_call_id_header = parent_call_id_header
            self.sequence_number_header = sequence_number_header
            self.id_generator = id_generator or generate_id
            self.user_id = user_id

        async def dispatch(self, request: Request, call_next) -> Response:
            # Get request ID from header
            request_id = request.headers.get(self.request_id_header, "")
            trace_id = request.headers.get(self.trace_id_header, "")
            session_id = request.headers.get(self.session_id_header, "")
            call_id = request.headers.get(self.call_id_header, "")
            parent_call_id = request.headers.get(self.parent_call_id_header, "")
            sequence_number = request.headers.get(self.sequence_number_header, "")

            incoming_tracking = {
                self.request_id_header: request_id or "",
                self.trace_id_header: trace_id or "",
                self.session_id_header: session_id or "",
                self.call_id_header: call_id or "",
                self.parent_call_id_header: parent_call_id or "",
                self.sequence_number_header: sequence_number or "",
            }
            print(f"[a2a_llm_tracker] Incoming tracking headers: {incoming_tracking}")

            if request_id:
                # Preserve inbound request ID exactly to keep end-to-end correlation stable.
                set_request_id(request_id)
            elif trace_id:
                # If request ID is missing, derive request lineage from trace ID in CCS.
                request_id = await set_request_id_async(trace_id, user_id=self.user_id)
            elif self.generate_request_id:
                # Generate a plain request ID for correlation when neither request/trace is provided.
                request_id = self.id_generator()
                set_request_id(request_id)

            # Get or generate session ID
            if not session_id and self.generate_session_id:
                session_id = self.id_generator()

            if session_id:
                set_session_id(session_id)

            # Get or generate call lineage IDs
            if not call_id and self.generate_call_id:
                call_id = self.id_generator()

            if call_id:
                set_call_id(call_id)
            if parent_call_id:
                set_parent_call_id(parent_call_id)
            if sequence_number:
                set_sequence_number(sequence_number)

            # Also set in tracker context for automatic inclusion in analyze_response
            from .context import set_context
            if not trace_id:
                trace_id = request_id
            set_context(trace_id=trace_id, session_id=session_id)

            # Add to request state for easy access in route handlers
            request.state.request_id = request_id
            request.state.trace_id = trace_id
            request.state.session_id = session_id
            request.state.call_id = call_id
            request.state.parent_call_id = parent_call_id
            request.state.sequence_number = sequence_number

            # Process request
            response = await call_next(request)

            # Add tracking IDs to response headers.
            if request_id:
                response.headers[self.request_id_header] = request_id
            # Preserve downstream trace IDs (for example, proxy-generated UUIDs).
            # Only backfill trace header when the response does not already include one.
            existing_trace_id = response.headers.get(self.trace_id_header, "")
            if not existing_trace_id and trace_id:
                response.headers[self.trace_id_header] = trace_id
            if session_id:
                response.headers[self.session_id_header] = session_id
            if call_id:
                response.headers[self.call_id_header] = call_id
            if parent_call_id:
                response.headers[self.parent_call_id_header] = parent_call_id
            if sequence_number:
                response.headers[self.sequence_number_header] = sequence_number

            outgoing_tracking = {
                self.request_id_header: response.headers.get(self.request_id_header, ""),
                self.trace_id_header: response.headers.get(self.trace_id_header, ""),
                self.session_id_header: response.headers.get(self.session_id_header, ""),
                self.call_id_header: response.headers.get(self.call_id_header, ""),
                self.parent_call_id_header: response.headers.get(self.parent_call_id_header, ""),
                self.sequence_number_header: response.headers.get(self.sequence_number_header, ""),
            }
            print(f"[a2a_llm_tracker] Outgoing tracking headers: {outgoing_tracking}")

            return response

    STARLETTE_AVAILABLE = True

except ImportError:
    STARLETTE_AVAILABLE = False
    TrackerMiddleware = None  # type: ignore


def require_starlette():
    """Raise ImportError if Starlette is not available."""
    if not STARLETTE_AVAILABLE:
        raise ImportError(
            "Starlette is required for TrackerMiddleware. "
            "Install it with: pip install starlette"
        )
