# Copyright 2024 Klira AI
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""@guardrails decorator — applies guardrails to a function.

Runs inbound guardrails on input and outbound guardrails on output.
Uses block/allow policy actions (no "warn").
"""

from __future__ import annotations

import asyncio
import functools
import inspect
import logging
from typing import Any, Callable, TypeVar, overload

from klira.sdk.exceptions import KliraPolicyViolation
from klira.sdk.guardrails.engine import GuardrailsEngine

logger = logging.getLogger(__name__)

F = TypeVar("F", bound=Callable[..., Any])


def _extract_input_text(*args: Any, **kwargs: Any) -> str:
    """Extract input text from function arguments for guardrails check.

    Tries common patterns: first string arg, 'query', 'text', 'input',
    'message', 'prompt' kwargs.
    """
    # Check common kwarg names
    for key in ("query", "text", "input", "message", "prompt"):
        if key in kwargs and isinstance(kwargs[key], str):
            return kwargs[key]

    # Use first string positional arg
    for arg in args:
        if isinstance(arg, str):
            return arg

    # Fallback: stringify all args
    return str(args) + str(kwargs)


def _extract_domain(**kwargs: Any) -> str | None:
    """Extract domain from function keyword arguments.

    Probes for common domain-related kwarg names.
    Returns None if no domain kwarg is found.
    """
    for key in ("domain", "url", "origin"):
        val = kwargs.get(key)
        if val is not None and isinstance(val, str) and val:
            return val
    return None


@overload
def guardrails(func: F) -> F: ...
@overload
def guardrails(
    *,
    domain: str | None = None,
    policy_set: str | None = None,
    check_input: bool = True,
    check_output: bool = True,
) -> Callable[[F], F]: ...


def guardrails(
    func: F | None = None,
    *,
    domain: str | None = None,
    policy_set: str | None = None,
    check_input: bool = True,
    check_output: bool = True,
) -> F | Callable[[F], F]:
    """Decorator that applies guardrails to a function.

    Runs inbound guardrails on input and outbound guardrails on output.
    If a blocking policy matches, raises KliraPolicyViolation.
    If an allow policy matches, guidelines are extracted for augmentation.

    Args:
        domain: Optional domain for domain-based policy matching. Can be
            overridden at call time via a 'domain', 'url', or 'origin' kwarg.
        policy_set: Optional policy set name to filter which policies apply.
        check_input: Whether to run inbound guardrails (default True).
        check_output: Whether to run outbound guardrails (default True).
    """

    def decorator(fn: F) -> F:
        is_async = asyncio.iscoroutinefunction(fn)

        # Check at decoration time if fn can accept _klira_guidelines
        _sig = inspect.signature(fn)
        _accepts_guidelines = (
            "_klira_guidelines" in _sig.parameters
            or any(
                p.kind == inspect.Parameter.VAR_KEYWORD
                for p in _sig.parameters.values()
            )
        )

        if is_async:

            @functools.wraps(fn)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                engine = GuardrailsEngine.get_instance()
                resolved_domain = _extract_domain(**kwargs) or domain

                # Inbound check
                if check_input:
                    input_text = _extract_input_text(*args, **kwargs)
                    input_result = engine.check_input(
                        input_text, domain=resolved_domain, policy_set=policy_set
                    )
                    if not input_result.allowed:
                        raise KliraPolicyViolation(
                            message=input_result.violation_response
                            or "Input blocked by guardrails policy",
                            policy_id=(
                                input_result.matched_policies[0].policy.id
                                if input_result.matched_policies
                                else None
                            ),
                            violation_response=input_result.violation_response,
                        )

                    # Inject guidelines if function accepts them
                    if _accepts_guidelines and input_result.guidelines:
                        kwargs["_klira_guidelines"] = input_result.guidelines

                # Execute the function
                result = await fn(*args, **kwargs)

                # Outbound check
                if check_output:
                    output_text = str(result) if result is not None else ""
                    if output_text:
                        output_result = engine.check_output(
                            output_text, domain=resolved_domain, policy_set=policy_set
                        )
                        if not output_result.allowed:
                            raise KliraPolicyViolation(
                                message=output_result.violation_response
                                or "Output blocked by guardrails policy",
                                policy_id=(
                                    output_result.matched_policies[0].policy.id
                                    if output_result.matched_policies
                                    else None
                                ),
                                violation_response=output_result.violation_response,
                            )

                return result

            return async_wrapper  # type: ignore[return-value]
        else:

            @functools.wraps(fn)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                engine = GuardrailsEngine.get_instance()
                resolved_domain = _extract_domain(**kwargs) or domain

                # Inbound check
                if check_input:
                    input_text = _extract_input_text(*args, **kwargs)
                    input_result = engine.check_input(
                        input_text, domain=resolved_domain, policy_set=policy_set
                    )
                    if not input_result.allowed:
                        raise KliraPolicyViolation(
                            message=input_result.violation_response
                            or "Input blocked by guardrails policy",
                            policy_id=(
                                input_result.matched_policies[0].policy.id
                                if input_result.matched_policies
                                else None
                            ),
                            violation_response=input_result.violation_response,
                        )

                    # Inject guidelines if function accepts them
                    if _accepts_guidelines and input_result.guidelines:
                        kwargs["_klira_guidelines"] = input_result.guidelines

                # Execute the function
                result = fn(*args, **kwargs)

                # Outbound check
                if check_output:
                    output_text = str(result) if result is not None else ""
                    if output_text:
                        output_result = engine.check_output(
                            output_text, domain=resolved_domain, policy_set=policy_set
                        )
                        if not output_result.allowed:
                            raise KliraPolicyViolation(
                                message=output_result.violation_response
                                or "Output blocked by guardrails policy",
                                policy_id=(
                                    output_result.matched_policies[0].policy.id
                                    if output_result.matched_policies
                                    else None
                                ),
                                violation_response=output_result.violation_response,
                            )

                return result

            return sync_wrapper  # type: ignore[return-value]

    if func is not None:
        return decorator(func)
    return decorator
