# 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.

"""Core decorators: @workflow, @agent, @task, @tool.

All decorators use tracer.start_as_current_span() directly.
No framework detection in decorators — decorators create spans,
adapters handle framework-specific behavior.

Key invariants:
- Every decorator that produces a result captures output on the span,
  truncated to 500 chars. (Learning #4)
- Span attributes are consistent across ALL frameworks. (Learning #1)
- No lazy imports or service locators.
"""

from __future__ import annotations

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

from opentelemetry import context, trace

from klira.sdk.tracing.propagation import (
    is_in_klira_trace,
    set_klira_context,
)

logger = logging.getLogger(__name__)

OUTPUT_TRUNCATION_LIMIT = 5000
INPUT_TRUNCATION_LIMIT = 10000

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


def _extract_input_text(*args: Any, **kwargs: Any) -> str | None:
    """Extract user input text from function arguments.

    Checks common keyword argument names first, then falls back to the first
    positional string argument.
    """
    for key in ("query", "text", "input", "message", "prompt"):
        if key in kwargs and isinstance(kwargs[key], str):
            return kwargs[key]
    for arg in args:
        if isinstance(arg, str):
            return arg
    return None


def _capture_input(span: trace.Span, *args: Any, **kwargs: Any) -> None:
    """Capture user input text on the span, truncated to INPUT_TRUNCATION_LIMIT."""
    input_text = _extract_input_text(*args, **kwargs)
    if input_text is not None:
        if len(input_text) > INPUT_TRUNCATION_LIMIT:
            input_text = input_text[:INPUT_TRUNCATION_LIMIT]
        span.set_attribute("klira.input", input_text)


def _capture_output(span: trace.Span, result: Any) -> None:
    """Capture function output on the span, truncated to OUTPUT_TRUNCATION_LIMIT.

    Never truncate decorator output capture. (Learning #4)
    Actually: truncate to 5000 chars to keep traces manageable but always capture.
    """
    if result is not None:
        output_str = str(result)
        if len(output_str) > OUTPUT_TRUNCATION_LIMIT:
            output_str = output_str[:OUTPUT_TRUNCATION_LIMIT]
        span.set_attribute("klira.output", output_str)


def _make_span_wrapper(
    func: Callable[..., Any],
    span_name: str,
    entity_type: str,
    entity_name: str,
    extra_attributes: dict[str, Any] | None = None,
    auto_trace: bool = False,
    user_id: str | None = None,
    conversation_id: str | None = None,
) -> Callable[..., Any]:
    """Create a wrapper that creates a span around the function call.

    For @workflow with auto_trace=True, also creates a root
    klira.user.message span if not already in a trace context.

    Args:
        user_id: User identifier. Set on the span and used for
            auto-trace root span when not in a trace context.
        conversation_id: Conversation identifier. Set on the span and
            used for auto-trace root span.
    """
    is_async = asyncio.iscoroutinefunction(func)

    attributes: dict[str, Any] = {
        "klira.entity_type": entity_type,
        "klira.entity_name": entity_name,
    }
    if user_id is not None:
        attributes["klira.user_id"] = user_id
    if conversation_id is not None:
        attributes["klira.conversation_id"] = conversation_id
    if extra_attributes:
        attributes.update(extra_attributes)

    def _resolve_user_id(**kwargs: Any) -> str:
        """Resolve user_id from decorator param, kwargs, or context."""
        if user_id is not None:
            return user_id
        if "user_id" in kwargs:
            return str(kwargs["user_id"])
        from klira.sdk.tracing.propagation import get_klira_user_id

        ctx_uid = get_klira_user_id()
        if ctx_uid is not None:
            return ctx_uid
        return "anonymous"

    def _resolve_conversation_id(**kwargs: Any) -> str:
        """Resolve conversation_id from decorator param, kwargs, or context."""
        if conversation_id is not None:
            return conversation_id
        if "conversation_id" in kwargs:
            return str(kwargs["conversation_id"])
        from klira.sdk.tracing.propagation import get_klira_conversation_id

        ctx_cid = get_klira_conversation_id()
        if ctx_cid is not None:
            return ctx_cid
        return str(uuid.uuid4())

    if is_async:

        @functools.wraps(func)
        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
            tracer = trace.get_tracer("klira")
            token = None
            root_span = None
            root_span_token = None

            try:
                # Auto-trace: create root span if not in trace context
                if auto_trace and not is_in_klira_trace():
                    resolved_uid = _resolve_user_id(**kwargs)
                    resolved_cid = _resolve_conversation_id(**kwargs)
                    resolved_mid = str(uuid.uuid4())
                    ctx = set_klira_context(
                        user_id=resolved_uid,
                        conversation_id=resolved_cid,
                        message_id=resolved_mid,
                    )
                    token = context.attach(ctx)
                    root_attrs: dict[str, Any] = {
                        "klira.entity_type": "user_message",
                        "klira.user_id": resolved_uid,
                        "klira.conversation_id": resolved_cid,
                        "klira.message_id": resolved_mid,
                    }
                    try:
                        from klira.sdk import Klira

                        if Klira.is_initialized():
                            config = Klira.get_config()
                            if config.evals_run is not None:
                                root_attrs["klira.evals.evals_run"] = config.evals_run
                    except Exception:
                        logger.debug("Failed to inject klira.evals.evals_run", exc_info=True)
                    root_span = tracer.start_span(
                        "klira.user.message",
                        attributes=root_attrs,
                    )
                    _capture_input(root_span, *args, **kwargs)
                    ctx_with_span = trace.set_span_in_context(
                        root_span, context.get_current()
                    )
                    root_span_token = context.attach(ctx_with_span)

                # Enrich span attrs with context-resolved values if not set
                runtime_attrs = dict(attributes)
                if "klira.user_id" not in runtime_attrs:
                    resolved = _resolve_user_id(**kwargs)
                    if resolved != "anonymous":
                        runtime_attrs["klira.user_id"] = resolved
                if "klira.conversation_id" not in runtime_attrs:
                    from klira.sdk.tracing.propagation import (
                        get_klira_conversation_id,
                    )

                    ctx_cid = get_klira_conversation_id()
                    if ctx_cid is not None:
                        runtime_attrs["klira.conversation_id"] = ctx_cid
                if "klira.framework" not in runtime_attrs:
                    try:
                        from klira.sdk import Klira

                        if Klira.is_initialized():
                            fw = Klira.get_config().framework
                            if fw is not None:
                                runtime_attrs["klira.framework"] = fw
                    except Exception:
                        logger.debug("Failed to inject klira.framework", exc_info=True)

                with tracer.start_as_current_span(
                    span_name, attributes=runtime_attrs
                ) as span:
                    if entity_type == "workflow":
                        _capture_input(span, *args, **kwargs)
                    result = await func(*args, **kwargs)
                    _capture_output(span, result)
                    return result
            finally:
                if root_span is not None:
                    root_span.end()
                if root_span_token is not None:
                    context.detach(root_span_token)
                if token is not None:
                    context.detach(token)

        return async_wrapper
    else:

        @functools.wraps(func)
        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
            tracer = trace.get_tracer("klira")
            token = None
            root_span = None
            root_span_token = None

            try:
                # Auto-trace: create root span if not in trace context
                if auto_trace and not is_in_klira_trace():
                    resolved_uid = _resolve_user_id(**kwargs)
                    resolved_cid = _resolve_conversation_id(**kwargs)
                    resolved_mid = str(uuid.uuid4())
                    ctx = set_klira_context(
                        user_id=resolved_uid,
                        conversation_id=resolved_cid,
                        message_id=resolved_mid,
                    )
                    token = context.attach(ctx)
                    root_attrs: dict[str, Any] = {
                        "klira.entity_type": "user_message",
                        "klira.user_id": resolved_uid,
                        "klira.conversation_id": resolved_cid,
                        "klira.message_id": resolved_mid,
                    }
                    try:
                        from klira.sdk import Klira

                        if Klira.is_initialized():
                            config = Klira.get_config()
                            if config.evals_run is not None:
                                root_attrs["klira.evals.evals_run"] = config.evals_run
                    except Exception:
                        logger.debug("Failed to inject klira.evals.evals_run", exc_info=True)
                    root_span = tracer.start_span(
                        "klira.user.message",
                        attributes=root_attrs,
                    )
                    _capture_input(root_span, *args, **kwargs)
                    ctx_with_span = trace.set_span_in_context(
                        root_span, context.get_current()
                    )
                    root_span_token = context.attach(ctx_with_span)

                # Enrich span attrs with context-resolved values if not set
                runtime_attrs = dict(attributes)
                if "klira.user_id" not in runtime_attrs:
                    resolved = _resolve_user_id(**kwargs)
                    if resolved != "anonymous":
                        runtime_attrs["klira.user_id"] = resolved
                if "klira.conversation_id" not in runtime_attrs:
                    from klira.sdk.tracing.propagation import (
                        get_klira_conversation_id,
                    )

                    ctx_cid = get_klira_conversation_id()
                    if ctx_cid is not None:
                        runtime_attrs["klira.conversation_id"] = ctx_cid
                if "klira.framework" not in runtime_attrs:
                    try:
                        from klira.sdk import Klira

                        if Klira.is_initialized():
                            fw = Klira.get_config().framework
                            if fw is not None:
                                runtime_attrs["klira.framework"] = fw
                    except Exception:
                        logger.debug("Failed to inject klira.framework", exc_info=True)

                with tracer.start_as_current_span(
                    span_name, attributes=runtime_attrs
                ) as span:
                    if entity_type == "workflow":
                        _capture_input(span, *args, **kwargs)
                    result = func(*args, **kwargs)
                    _capture_output(span, result)
                    return result
            finally:
                if root_span is not None:
                    root_span.end()
                if root_span_token is not None:
                    context.detach(root_span_token)
                if token is not None:
                    context.detach(token)

        return sync_wrapper


@overload
def workflow(func: F) -> F: ...
@overload
def workflow(
    *,
    name: str | None = None,
    user_id: str | None = None,
    conversation_id: str | None = None,
) -> Callable[[F], F]: ...


def workflow(
    func: F | None = None,
    *,
    name: str | None = None,
    user_id: str | None = None,
    conversation_id: str | None = None,
) -> F | Callable[[F], F]:
    """Decorator for workflow functions.

    Creates a klira.workflow.{name} span. Auto-creates a
    klira.user.message root span if not already in a trace context.

    Args:
        name: Custom workflow name (defaults to function name).
        user_id: User identifier for tracing and reporting.
        conversation_id: Conversation identifier for tracing.

    Usage:
        @workflow
        def my_flow(query: str) -> str: ...

        @workflow(name="custom_name", user_id="user_123")
        def another_flow(query: str) -> str: ...
    """

    def decorator(fn: F) -> F:
        workflow_name = name or fn.__name__
        span_name = f"klira.workflow.{workflow_name}"
        wrapped = _make_span_wrapper(
            fn,
            span_name=span_name,
            entity_type="workflow",
            entity_name=workflow_name,
            auto_trace=True,
            user_id=user_id,
            conversation_id=conversation_id,
        )
        return wrapped  # type: ignore[return-value]

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


@overload
def agent(func: F) -> F: ...
@overload
def agent(
    *,
    name: str | None = None,
    user_id: str | None = None,
    conversation_id: str | None = None,
) -> Callable[[F], F]: ...


def agent(
    func: F | None = None,
    *,
    name: str | None = None,
    user_id: str | None = None,
    conversation_id: str | None = None,
) -> F | Callable[[F], F]:
    """Decorator for agent functions.

    Creates a klira.agent.{name} span.

    Args:
        name: Custom agent name (defaults to function name).
        user_id: User identifier for tracing and reporting.
        conversation_id: Conversation identifier for tracing.
    """

    def decorator(fn: F) -> F:
        agent_name = name or fn.__name__
        span_name = f"klira.agent.{agent_name}"
        wrapped = _make_span_wrapper(
            fn,
            span_name=span_name,
            entity_type="agent",
            entity_name=agent_name,
            user_id=user_id,
            conversation_id=conversation_id,
        )
        return wrapped  # type: ignore[return-value]

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


@overload
def task(func: F) -> F: ...
@overload
def task(
    *,
    name: str | None = None,
    user_id: str | None = None,
    conversation_id: str | None = None,
) -> Callable[[F], F]: ...


def task(
    func: F | None = None,
    *,
    name: str | None = None,
    user_id: str | None = None,
    conversation_id: str | None = None,
) -> F | Callable[[F], F]:
    """Decorator for task functions.

    Creates a klira.task.{name} span.

    Args:
        name: Custom task name (defaults to function name).
        user_id: User identifier for tracing and reporting.
        conversation_id: Conversation identifier for tracing.
    """

    def decorator(fn: F) -> F:
        task_name = name or fn.__name__
        span_name = f"klira.task.{task_name}"
        wrapped = _make_span_wrapper(
            fn,
            span_name=span_name,
            entity_type="task",
            entity_name=task_name,
            user_id=user_id,
            conversation_id=conversation_id,
        )
        return wrapped  # type: ignore[return-value]

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


@overload
def tool(func: F) -> F: ...
@overload
def tool(
    *,
    name: str | None = None,
    fhir: str | None = None,
    user_id: str | None = None,
    conversation_id: str | None = None,
) -> Callable[[F], F]: ...


def tool(
    func: F | None = None,
    *,
    name: str | None = None,
    fhir: str | None = None,
    user_id: str | None = None,
    conversation_id: str | None = None,
) -> F | Callable[[F], F]:
    """Decorator for tool functions.

    Creates a klira.tool.{name} span.

    Args:
        name: Custom tool name (defaults to function name).
        fhir: FHIR resource type (e.g., "Patient"). Sets
            klira.fhir.resource_type on the span.
        user_id: User identifier for tracing and reporting.
        conversation_id: Conversation identifier for tracing.
    """

    def decorator(fn: F) -> F:
        tool_name = name or fn.__name__
        span_name = f"klira.tool.{tool_name}"
        extra: dict[str, Any] = {}
        if fhir is not None:
            extra["klira.fhir.resource_type"] = fhir
        wrapped = _make_span_wrapper(
            fn,
            span_name=span_name,
            entity_type="tool",
            entity_name=tool_name,
            extra_attributes=extra,
            user_id=user_id,
            conversation_id=conversation_id,
        )
        return wrapped  # type: ignore[return-value]

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