"""
Build the ASGI callable (scope, receive, send) for Knative func-python.
Reads scope["event"], parses event.data into the user's Pydantic model, builds StaxContext, calls handler.
"""

from __future__ import annotations

import asyncio
import json
import os
from collections.abc import Callable
from typing import Any, MutableMapping, TypeVar

from pydantic import BaseModel

from stax_sdk.automation.context import StaxContext
from stax_sdk.automation.decorator import _event_data_as_dict, _raw_event_to_dict
from stax_sdk.cx import AsyncClient as CxAsyncClient
from stax_sdk.cx import Client as CxClient
from stax_sdk.ta import AsyncClient as TaAsyncClient
from stax_sdk.ta import Client as TaClient

T = TypeVar("T", bound=BaseModel)


def _build_context_sync(
    raw_event: dict[str, Any],
    inject_headers: Callable[[MutableMapping[str, str]], None] | None,
) -> StaxContext:
    module_id = os.environ.get("STAX_MODULE_ID", "")
    cx = CxClient(inject_headers=inject_headers)
    ta = TaClient(inject_headers=inject_headers)
    return StaxContext(raw_event=raw_event, module_id=module_id, cx=cx, ta=ta)


def _build_context_async(
    raw_event: dict[str, Any],
    inject_headers: Callable[[MutableMapping[str, str]], None] | None,
) -> StaxContext:
    module_id = os.environ.get("STAX_MODULE_ID", "")
    cx = CxAsyncClient(inject_headers=inject_headers)
    ta = TaAsyncClient(inject_headers=inject_headers)
    return StaxContext(raw_event=raw_event, module_id=module_id, cx=cx, ta=ta)


def asgi_app_for_sync(
    handler: Callable[[T, StaxContext], Any],
    event_data_model: type[T],
    inject_headers: Callable[[MutableMapping[str, str]], None] | None = None,
) -> Callable[..., Any]:
    """Return an ASGI app that parses the event, builds StaxContext (sync cx/ta), runs handler in executor."""

    async def asgi_app(scope: dict, receive: Callable, send: Callable) -> None:
        event = scope.get("event")
        if event is None:
            await send({"type": "http.response.start", "status": 400})
            await send({"type": "http.response.body", "body": b"No event in scope"})
            return
        raw_event = _raw_event_to_dict(event)
        try:
            data_dict = _event_data_as_dict(event)
            event_data = event_data_model.model_validate(data_dict)
        except Exception as e:
            await send({"type": "http.response.start", "status": 400})
            await send(
                {"type": "http.response.body", "body": json.dumps({"error": str(e)}).encode()}
            )
            return
        context = _build_context_sync(raw_event, inject_headers)
        try:
            result = await asyncio.to_thread(handler, event_data, context)
        except Exception as e:
            await send({"type": "http.response.start", "status": 500})
            await send(
                {"type": "http.response.body", "body": json.dumps({"error": str(e)}).encode()}
            )
            return
        body = b""
        if result is not None:
            if isinstance(result, (dict, list)):
                body = json.dumps(result).encode()
            elif isinstance(result, bytes):
                body = result
            else:
                body = str(result).encode()
        await send(
            {
                "type": "http.response.start",
                "status": 200,
                "headers": [[b"content-type", b"application/json"]],
            }
        )
        await send({"type": "http.response.body", "body": body})

    return asgi_app


def asgi_app_for_async(
    handler: Callable[[T, StaxContext], Any],
    event_data_model: type[T],
    inject_headers: Callable[[MutableMapping[str, str]], None] | None = None,
) -> Callable[..., Any]:
    """Return an ASGI app that uses async cx/ta and awaits the handler."""

    async def asgi_app(scope: dict, receive: Callable, send: Callable) -> None:
        event = scope.get("event")
        if event is None:
            await send({"type": "http.response.start", "status": 400})
            await send({"type": "http.response.body", "body": b"No event in scope"})
            return
        raw_event = _raw_event_to_dict(event)
        try:
            data_dict = _event_data_as_dict(event)
            event_data = event_data_model.model_validate(data_dict)
        except Exception as e:
            await send({"type": "http.response.start", "status": 400})
            await send(
                {"type": "http.response.body", "body": json.dumps({"error": str(e)}).encode()}
            )
            return
        context = _build_context_async(raw_event, inject_headers)
        try:
            result = await handler(event_data, context)
        except Exception as e:
            await send({"type": "http.response.start", "status": 500})
            await send(
                {"type": "http.response.body", "body": json.dumps({"error": str(e)}).encode()}
            )
            return
        body = b""
        if result is not None:
            if isinstance(result, (dict, list)):
                body = json.dumps(result).encode()
            elif isinstance(result, bytes):
                body = result
            else:
                body = str(result).encode()
        await send(
            {
                "type": "http.response.start",
                "status": 200,
                "headers": [[b"content-type", b"application/json"]],
            }
        )
        await send({"type": "http.response.body", "body": body})

    return asgi_app
