"""AutoPIL ASGI Middleware. See packages/core/autopil/middleware.py for full docs."""

from __future__ import annotations

import dataclasses
import json
import re
import uuid
from typing import Any, Callable
from urllib.parse import parse_qs

from .guard import ContextGuard
from .models import Decision, SensitivityLevel


@dataclasses.dataclass
class RouteRule:
    """
    Governance rule for a URL path pattern.

    See packages/core/autopil/middleware.py for full documentation.
    """
    path_pattern:      str
    agent_role:        str
    source_id:         str
    sensitivity_level: SensitivityLevel = SensitivityLevel.MEDIUM
    methods:           frozenset | None = None
    user_id_header:    str              = "X-User-ID"
    session_id_header: str              = "X-Session-ID"
    query_param:       str              = "q"
    query_body_field:  str | None       = None


class AutoPILMiddleware:
    """ASGI middleware that enforces AutoPIL governance at the API layer."""

    def __init__(self, app: Any, *, guard: ContextGuard,
                 rules: list[RouteRule], on_deny: str = "block") -> None:
        self.app     = app
        self.guard   = guard
        self.on_deny = on_deny
        self._compiled = [(_compile_pattern(r.path_pattern), r) for r in rules]

    async def __call__(self, scope: dict, receive: Callable, send: Callable) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        path   = scope.get("path", "")
        method = scope.get("method", "GET").upper()
        rule, path_params = self._match(path, method)
        if rule is None:
            await self.app(scope, receive, send)
            return

        raw_headers = {
            k.decode("latin-1").lower(): v.decode("latin-1")
            for k, v in scope.get("headers", [])
        }
        qs_params  = parse_qs(scope.get("query_string", b"").decode())
        user_id    = raw_headers.get(rule.user_id_header.lower(), "anonymous")
        session_id = raw_headers.get(rule.session_id_header.lower(), "") or str(uuid.uuid4())
        query      = qs_params.get(rule.query_param, [""])[0]

        replay_receive = receive
        if rule.query_body_field:
            body_bytes, replay_receive = await _buffer_body(receive)
            try:
                body_json = json.loads(body_bytes)
                query = body_json.get(rule.query_body_field, query)
            except (json.JSONDecodeError, AttributeError, UnicodeDecodeError):
                pass

        source_id       = _resolve_source_id(rule.source_id, path_params)
        effective_query = query or path

        decision, policy_name, reason = self.guard._evaluate_request(
            agent_role=rule.agent_role, user_id=user_id, source_id=source_id,
            sensitivity_level=rule.sensitivity_level, sid=session_id,
            query=effective_query,
        )

        event_id = str(uuid.uuid4())
        self.guard._record_event(
            event_id=event_id, agent_role=rule.agent_role, user_id=user_id,
            source_id=source_id, query=effective_query, decision=decision,
            policy_name=policy_name, reason=reason,
            sensitivity_level=rule.sensitivity_level, sid=session_id,
            context_hash=None, source_type="api",
        )

        if decision == Decision.DENY and self.on_deny == "block":
            await _send_deny_response(send, {
                "error": "access_denied", "decision": "DENY",
                "source_id": source_id, "agent_role": rule.agent_role,
                "policy_name": policy_name, "reason": reason,
                "event_id": event_id, "session_id": session_id,
            })
            return

        autopil_headers = [
            (b"x-autopil-decision",   decision.value.encode()),
            (b"x-autopil-event-id",   event_id.encode()),
            (b"x-autopil-policy",     policy_name.encode()),
            (b"x-autopil-session-id", session_id.encode()),
        ]

        async def send_with_headers(message: dict) -> None:
            if message["type"] == "http.response.start":
                message = {**message, "headers": list(message.get("headers", [])) + autopil_headers}
            await send(message)

        await self.app(scope, replay_receive, send_with_headers)

    def _match(self, path: str, method: str) -> tuple[RouteRule | None, dict]:
        for pattern, rule in self._compiled:
            if rule.methods and method not in rule.methods:
                continue
            m = pattern.match(path)
            if m:
                return rule, m.groupdict()
        return None, {}


def _compile_pattern(path_pattern: str) -> re.Pattern:
    tokens = re.split(r'(\{[^}]+\}|\*\*|\*)', path_pattern)
    parts  = []
    for tok in tokens:
        if tok.startswith('{') and tok.endswith('}'):
            parts.append(f'(?P<{tok[1:-1]}>[^/]+)')
        elif tok == '**':
            parts.append('.*')
        elif tok == '*':
            parts.append('[^/]*')
        else:
            parts.append(re.escape(tok))
    return re.compile('^' + ''.join(parts) + '$')


def _resolve_source_id(template: str, path_params: dict) -> str:
    return re.sub(r'\{(\w+)\}', lambda m: path_params.get(m.group(1), m.group(0)), template)


async def _buffer_body(receive: Callable) -> tuple[bytes, Callable]:
    messages: list[dict] = []
    while True:
        msg = await receive()
        messages.append(msg)
        if not msg.get("more_body", False):
            break
    body  = b"".join(m.get("body", b"") for m in messages)
    state = {"index": 0}
    async def replay() -> dict:
        i = state["index"]; state["index"] += 1; return messages[i]
    return body, replay


async def _send_deny_response(send: Callable, body_dict: dict) -> None:
    body = json.dumps(body_dict).encode()
    await send({"type": "http.response.start", "status": 403, "headers": [
        (b"content-type", b"application/json"),
        (b"content-length", str(len(body)).encode()),
        (b"x-autopil-decision", b"DENY"),
    ]})
    await send({"type": "http.response.body", "body": body, "more_body": False})
