# Copyright (c) Microsoft. All rights reserved.

"""Utility functions for DevUI."""

import inspect
import json
import logging
from dataclasses import fields, is_dataclass
from types import UnionType
from typing import Any, Union, cast, get_args, get_origin, get_type_hints

from agent_framework import Message

logger = logging.getLogger(__name__)


def _string_key_dict(value: object) -> dict[str, Any] | None:
    """Cast value to a dict."""
    if not isinstance(value, dict):
        return None
    return cast(dict[str, Any], value)


# ============================================================================
# Agent Metadata Extraction
# ============================================================================


def extract_agent_metadata(entity_object: Any) -> dict[str, Any]:
    """Extract agent-specific metadata from an entity object.

    Args:
        entity_object: Agent Framework agent object

    Returns:
        Dictionary with agent metadata: instructions, model, chat_client_type,
        context_providers, and middleware
    """
    metadata = {
        "instructions": None,
        "model": None,
        "chat_client_type": None,
        "context_provider": None,
        "middleware": None,
    }

    # Try to get instructions
    if hasattr(entity_object, "default_options"):
        chat_opts = entity_object.default_options
        chat_opts_dict = _string_key_dict(chat_opts)
        if chat_opts_dict is not None:
            if "instructions" in chat_opts_dict:
                metadata["instructions"] = chat_opts_dict.get("instructions")
        elif hasattr(chat_opts, "instructions"):
            metadata["instructions"] = chat_opts.instructions

    # Try to get model - check both default_options and client
    if hasattr(entity_object, "default_options"):
        chat_opts = entity_object.default_options
        chat_opts_dict = _string_key_dict(chat_opts)
        if chat_opts_dict is not None:
            model = chat_opts_dict.get("model")
            if model:
                metadata["model"] = model
        elif hasattr(chat_opts, "model") and chat_opts.model:
            metadata["model"] = chat_opts.model
    if metadata["model"] is None and hasattr(entity_object, "client") and hasattr(entity_object.client, "model"):
        metadata["model"] = entity_object.client.model

    # Try to get chat client type
    if hasattr(entity_object, "client"):
        metadata["chat_client_type"] = entity_object.client.__class__.__name__

    # Try to get context providers
    if (
        hasattr(entity_object, "context_provider")
        and entity_object.context_provider
        and hasattr(entity_object.context_provider, "__class__")
    ):
        metadata["context_provider"] = [entity_object.context_provider.__class__.__name__]  # type: ignore

    # Try to get middleware
    if hasattr(entity_object, "middleware") and entity_object.middleware:
        middlewares_list: list[str] = []
        for m in entity_object.middleware:
            # Try multiple ways to get a good name for middleware
            if hasattr(m, "__name__"):  # Function or callable
                middlewares_list.append(m.__name__)
            elif hasattr(m, "__class__"):  # Class instance
                middlewares_list.append(m.__class__.__name__)
            else:
                middlewares_list.append(str(m))
        metadata["middleware"] = middlewares_list  # type: ignore

    return metadata


# ============================================================================
# Workflow Input Type Utilities
# ============================================================================


def extract_executor_message_types(executor: Any) -> list[Any]:
    """Extract declared input types for the given executor.

    Args:
        executor: Workflow executor object

    Returns:
        List of message types that the executor accepts
    """
    message_types: list[Any] = []

    try:
        input_types = getattr(executor, "input_types", None)
    except Exception as exc:  # pragma: no cover - defensive logging path
        logger.debug(f"Failed to access executor input_types: {exc}")
    else:
        if input_types:
            message_types = list(input_types)

    if not message_types and hasattr(executor, "_handlers"):
        try:
            handlers = executor._handlers
            if isinstance(handlers, dict):
                message_types = list(handlers.keys())  # type: ignore[arg-type]
        except Exception as exc:  # pragma: no cover - defensive logging path
            logger.debug(f"Failed to read executor handlers: {exc}")

    return message_types


def _contains_chat_message(type_hint: Any) -> bool:
    """Check whether the provided type hint directly or indirectly references Message."""
    if type_hint is Message:
        return True

    origin = get_origin(type_hint)
    if origin in (list, tuple):
        return any(_contains_chat_message(arg) for arg in get_args(type_hint))

    if origin in (Union, UnionType):
        return any(_contains_chat_message(arg) for arg in get_args(type_hint))

    return False


def _is_list_message_type(type_hint: Any) -> bool:
    """Return True if type_hint is exactly list[Message]."""
    return get_origin(type_hint) is list and bool(get_args(type_hint)) and get_args(type_hint)[0] is Message


def _find_chat_message_type(type_hint: Any) -> Any | None:
    """Return ``list[Message]`` or ``Message`` if present in type_hint, else None.

    Recursively inspects union members so that a single-element ``message_types``
    list like ``[dict | str | list[Message] | ...]`` (the real ``JoinExecutor``
    form) is handled correctly.  ``list[Message]`` takes priority over bare
    ``Message``.
    """
    if _is_list_message_type(type_hint):
        return type_hint
    if type_hint is Message:
        return Message

    origin = get_origin(type_hint)
    if origin in (Union, UnionType):
        fallback = None
        for arg in get_args(type_hint):
            found = _find_chat_message_type(arg)
            if found is None:
                continue
            if _is_list_message_type(found):
                return found  # list[Message] wins immediately
            if fallback is None:
                fallback = found  # bare Message — keep searching
        return fallback

    return None


def select_primary_input_type(message_types: list[Any]) -> Any | None:
    """Choose the most user-friendly input type for workflow inputs.

    Prefers Message (or containers thereof) and then falls back to primitives.
    When the executor's union contains ``list[Message]`` (as the declarative
    entry ``JoinExecutor`` does), that type is returned so that
    ``parse_input_for_type`` can wrap the user's text in a list before
    dispatching — avoiding a "cannot handle message of type Message" error.

    ``message_types`` may contain full union types as single elements (the real
    ``JoinExecutor`` emits ``[dict | str | list[Message] | ...]``), so each
    element is searched recursively.

    Args:
        message_types: List of possible message types

    Returns:
        Selected primary input type, or None if list is empty
    """
    if not message_types:
        return None

    # First pass: search each type (including union members) for list[Message] or Message.
    for message_type in message_types:
        found = _find_chat_message_type(message_type)
        if found is not None:
            return found

    # Second pass: broader fallback for deeply nested or unusual containers.
    for message_type in message_types:
        if _contains_chat_message(message_type):
            return Message

    preferred = (str, dict)

    for candidate in preferred:
        for message_type in message_types:
            if message_type is candidate:
                return candidate
            origin = get_origin(message_type)
            if origin is candidate:
                return candidate

    return message_types[0]


# ============================================================================
# Type System Utilities
# ============================================================================


def is_serialization_mixin(cls: type) -> bool:
    """Check if class is a SerializationMixin subclass.

    Args:
        cls: Class to check

    Returns:
        True if class is a SerializationMixin subclass
    """
    try:
        from agent_framework._serialization import SerializationMixin

        return isinstance(cls, type) and issubclass(cls, SerializationMixin)
    except ImportError:
        return False


def _type_to_schema(type_hint: Any, field_name: str) -> dict[str, Any]:
    """Convert a type hint to JSON schema.

    Args:
        type_hint: Type hint to convert
        field_name: Name of the field (for documentation)

    Returns:
        JSON schema dict
    """
    type_str = str(type_hint)

    # Handle None/Optional
    if type_hint is type(None):
        return {"type": "null"}

    # Handle basic types
    if type_hint is str or "str" in type_str:
        return {"type": "string"}
    if type_hint is int or "int" in type_str:
        return {"type": "integer"}
    if type_hint is float or "float" in type_str:
        return {"type": "number"}
    if type_hint is bool or "bool" in type_str:
        return {"type": "boolean"}

    # Handle Literal types (for enum-like values)
    if "Literal" in type_str:
        origin = get_origin(type_hint)
        if origin is not None:
            args = get_args(type_hint)
            if args:
                return {"type": "string", "enum": list(args)}

    # Handle Union/Optional
    if "Union" in type_str or "Optional" in type_str:
        origin = get_origin(type_hint)
        if origin is not None:
            args = get_args(type_hint)
            # Filter out None type
            non_none_args = [arg for arg in args if arg is not type(None)]
            if len(non_none_args) == 1:
                return _type_to_schema(non_none_args[0], field_name)
            # Multiple types - pick first non-None
            if non_none_args:
                return _type_to_schema(non_none_args[0], field_name)

    # Handle collections
    if "list" in type_str or "List" in type_str or "Sequence" in type_str:
        origin = get_origin(type_hint)
        if origin is not None:
            args = get_args(type_hint)
            if args:
                items_schema = _type_to_schema(args[0], field_name)
                return {"type": "array", "items": items_schema}
        return {"type": "array"}

    if "dict" in type_str or "Dict" in type_str or "Mapping" in type_str:
        return {"type": "object"}

    # Default fallback
    return {"type": "string", "description": f"Type: {type_hint}"}


def generate_schema_from_serialization_mixin(cls: type[Any]) -> dict[str, Any]:
    """Generate JSON schema from SerializationMixin class.

    Introspects the __init__ signature to extract parameter types and defaults.

    Args:
        cls: SerializationMixin subclass

    Returns:
        JSON schema dict
    """
    sig = inspect.signature(cls)

    # Get type hints
    try:
        type_hints = get_type_hints(cls)
    except Exception:
        type_hints = {}

    properties: dict[str, Any] = {}
    required: list[str] = []

    for param_name, param in sig.parameters.items():
        if param_name in ("self", "kwargs"):
            continue

        # Get type annotation
        param_type = type_hints.get(param_name, str)

        # Generate schema for this parameter
        param_schema = _type_to_schema(param_type, param_name)
        properties[param_name] = param_schema

        # Check if required (no default value, not VAR_KEYWORD)
        if param.default == inspect.Parameter.empty and param.kind != inspect.Parameter.VAR_KEYWORD:
            required.append(param_name)

    schema: dict[str, Any] = {"type": "object", "properties": properties}

    if required:
        schema["required"] = required

    return schema


def generate_schema_from_dataclass(cls: type[Any]) -> dict[str, Any]:
    """Generate JSON schema from dataclass.

    Args:
        cls: Dataclass type

    Returns:
        JSON schema dict
    """
    if not is_dataclass(cls):
        return {"type": "object"}

    properties: dict[str, Any] = {}
    required: list[str] = []

    for field in fields(cls):
        # Generate schema for field type
        field_schema = _type_to_schema(field.type, field.name)
        properties[field.name] = field_schema

        # Check if required (no default value)
        if field.default == field.default_factory:  # No default
            required.append(field.name)

    schema: dict[str, Any] = {"type": "object", "properties": properties}

    if required:
        schema["required"] = required

    return schema


def extract_response_type_from_executor(executor: Any, request_type: type) -> type | None:
    """Extract the expected response type from an executor's response handler.

    Looks for methods decorated with @response_handler that have signature:
       async def handler(self, original_request: RequestType, response: ResponseType, ctx)

    Args:
        executor: Executor object that should have a handler for the request type
        request_type: The request message type

    Returns:
        The response type class, or None if not found
    """
    try:
        # Introspect handler methods for @response_handler pattern
        for attr_name in dir(executor):
            if attr_name.startswith("_"):
                continue
            attr = getattr(executor, attr_name, None)
            if not callable(attr):
                continue

            # Get type hints for this method
            try:
                type_hints = get_type_hints(attr)

                # Check for @response_handler pattern:
                # async def handler(self, original_request: RequestType, response: ResponseType, ctx)
                type_hint_params = {k: v for k, v in type_hints.items() if k not in ("self", "return")}

                # Look for at least 2 parameters: original_request, response (ctx is optional)
                if len(type_hint_params) >= 2:
                    param_items = list(type_hint_params.items())
                    # First param should be original_request matching request_type
                    _, first_param_type = param_items[0]
                    _, second_param_type = param_items[1] if len(param_items) > 1 else (None, None)

                    # Check if first param matches request_type
                    first_matches_request = first_param_type == request_type
                    if not first_matches_request and isinstance(first_param_type, type):
                        request_type_name = request_type.__name__
                        first_matches_request = first_param_type.__name__ == request_type_name

                    # Verify we have a matching request type and valid response type (must be a type class)
                    if first_matches_request and second_param_type is not None and isinstance(second_param_type, type):
                        response_type_class: type = second_param_type
                        logger.debug(
                            f"Found response type {response_type_class} for request {request_type} "
                            f"via @response_handler"
                        )
                        return response_type_class

            except Exception as e:
                logger.debug(f"Failed to get type hints for {attr_name}: {e}")
                continue

    except Exception as e:
        logger.debug(f"Failed to extract response type from executor: {e}")

    return None


def generate_input_schema(input_type: type) -> dict[str, Any]:
    """Generate JSON schema for workflow input type.

    Supports multiple input types in priority order:
    0. list[Message] — rendered as a plain string input (DevUI presents a text
       box; the text is later wrapped in a list by parse_input_for_type)
    1. Built-in types (str, dict, int, etc.)
    2. Pydantic models (via model_json_schema)
    3. SerializationMixin classes (via __init__ introspection)
    4. Dataclasses (via fields introspection)
    5. Fallback to string

    Args:
        input_type: Input type to generate schema for

    Returns:
        JSON schema dict
    """
    # 0. list[Message] — treat as simple text so DevUI shows a text box
    if _is_list_message_type(input_type):
        return {"type": "string"}

    # 1. Built-in types
    if input_type is str:
        return {"type": "string"}
    if input_type is dict:
        return {"type": "object"}
    if input_type is int:
        return {"type": "integer"}
    if input_type is float:
        return {"type": "number"}
    if input_type is bool:
        return {"type": "boolean"}

    # 2. Pydantic models (legacy support)
    if hasattr(input_type, "model_json_schema"):
        return input_type.model_json_schema()  # type: ignore

    # 3. SerializationMixin classes (Message, etc.)
    if is_serialization_mixin(input_type):
        return generate_schema_from_serialization_mixin(input_type)

    # 4. Dataclasses
    if is_dataclass(input_type):
        return generate_schema_from_dataclass(input_type)

    # 5. Fallback to string
    type_name = input_type.__name__ if isinstance(input_type, type) else str(cast(Any, input_type))
    return {"type": "string", "description": f"Input type: {type_name}"}


# ============================================================================
# Input Parsing Utilities
# ============================================================================


def parse_input_for_type(input_data: Any, target_type: type) -> Any:
    """Parse input data to match the target type.

    Handles conversion from raw input (string, dict) to the expected type:
    - list[Message]: build a Message from the raw input and wrap it in a list
    - Built-in types: direct conversion
    - Pydantic models: use model_validate or model_validate_json
    - SerializationMixin: use from_dict or construct from string
    - Dataclasses: construct from dict

    Args:
        input_data: Raw input data (string, dict, or already correct type)
        target_type: Expected type for the input

    Returns:
        Parsed input matching target_type, or original input if parsing fails
    """
    # list[Message]: generic aliases cannot be used with isinstance, handle first.
    if _is_list_message_type(target_type):
        raw: object = input_data
        if isinstance(raw, list):
            items = cast(list[object], raw)
            if all(isinstance(m, Message) for m in items):
                return items
            # Try to convert each item (serialized str/dict OpenAI message) to Message.
            converted: list[Message] = []
            ok = True
            for item in items:
                if isinstance(item, Message):
                    converted.append(item)
                elif isinstance(item, str):
                    converted.append(_build_message_from_legacy_payload(item))
                elif isinstance(item, dict):
                    converted.append(_build_message_from_legacy_payload(cast(dict[str, Any], item)))
                else:
                    ok = False
                    break
            if ok:
                return converted
            # A list never matches the Message/str/dict checks below; stringify directly.
            return [_build_message_from_legacy_payload(str(items))]
        if isinstance(raw, Message):
            return [raw]
        if isinstance(raw, str):
            return [_build_message_from_legacy_payload(raw)]
        parsed_dict = _string_key_dict(raw)
        if parsed_dict is not None:
            if parsed_dict and _looks_like_message_dict(parsed_dict):
                return [_build_message_from_legacy_payload(parsed_dict)]
            return raw
        return [_build_message_from_legacy_payload(str(raw))]

    # If already correct type, return as-is
    if isinstance(input_data, target_type):
        return input_data

    # Handle string input
    if isinstance(input_data, str):
        return _parse_string_input(input_data, target_type)

    # Handle dict input
    parsed_dict = _string_key_dict(input_data)
    if parsed_dict is not None:
        return _parse_dict_input(parsed_dict, target_type)

    # Fallback: return original
    return input_data


def _looks_like_message_dict(d: dict[str, Any]) -> bool:
    """Return True if *d* is a recognisable serialised Message payload.

    Three recognised signatures (in priority order):
    1. ``type`` discriminator is literally ``"message"`` — output of ``Message.to_dict()``.
    2. Dict has a ``"role"`` key — all framework Message objects carry a role.
    3. Dict has exactly the single key ``"input"`` — DevUI ``WorkflowInputForm``
       submits ``{"input": "<user text>"}`` for workflows whose schema is ``{"type":
       "string"}``.

    Anything else is treated as a structured workflow input and passed through
    unchanged, to be handled by the dict-union branch of the executor.
    """
    if d.get("type") == "message":
        return True
    if "role" in d:
        return True
    return set(d.keys()) == {"input"}


def _build_message_from_legacy_payload(input_data: str | dict[str, Any]) -> Message:
    """Convert raw DevUI input into a framework Message.

    This preserves DevUI compatibility for older payloads that still send
    ``{"role": "...", "text": "..."}`` instead of the framework-native
    ``{"role": "...", "contents": [...]}`` shape.
    """
    if isinstance(input_data, str):
        return Message(role="user", contents=[input_data])

    role = input_data.get("role", "user")
    role = role if isinstance(role, str) else str(role)

    if "contents" in input_data:
        contents = input_data["contents"]
    else:
        contents = None
        for field in ("text", "message", "content", "input", "data"):
            if field in input_data:
                contents = input_data[field]
                break

    if contents is None:
        contents_list: list[Any] = []
    elif isinstance(contents, list):
        contents_list = contents  # type: ignore[reportUnknownVariableType]
    else:
        contents_list = [contents]

    kwargs: dict[str, Any] = {}
    for field in (
        "author_name",
        "message_id",
        "additional_properties",
        "raw_representation",
    ):
        if field in input_data:
            kwargs[field] = input_data[field]

    return Message(role=role, contents=contents_list, **kwargs)


def _parse_string_input(input_str: str, target_type: type) -> Any:
    """Parse string input to target type.

    Args:
        input_str: Input string
        target_type: Target type

    Returns:
        Parsed input or original string
    """
    # Built-in types
    if target_type is str:
        return input_str
    if target_type is int:
        try:
            return int(input_str)
        except ValueError:
            return input_str
    elif target_type is float:
        try:
            return float(input_str)
        except ValueError:
            return input_str
    elif target_type is bool:
        return input_str.lower() in ("true", "1", "yes")

    # Pydantic models
    if hasattr(target_type, "model_validate_json"):
        try:
            # Try parsing as JSON first
            if input_str.strip().startswith("{"):
                return target_type.model_validate_json(input_str)  # type: ignore

            # Try common field names with the string value
            common_fields = ["text", "message", "content", "input", "data"]
            for field in common_fields:
                try:
                    return target_type(**{field: input_str})
                except Exception as e:
                    logger.debug(f"Failed to parse string input with field '{field}': {e}")
                    continue
        except Exception as e:
            logger.debug(f"Failed to parse string as Pydantic model: {e}")

    # SerializationMixin (like Message)
    if is_serialization_mixin(target_type):
        try:
            if target_type is Message:
                if input_str.strip().startswith("{"):
                    data = json.loads(input_str)
                    parsed_dict = _string_key_dict(data)
                    if parsed_dict is not None:
                        return _build_message_from_legacy_payload(parsed_dict)
                return _build_message_from_legacy_payload(input_str)

            # Try parsing as JSON dict first
            if input_str.strip().startswith("{"):
                data = json.loads(input_str)
                if hasattr(target_type, "from_dict"):
                    return target_type.from_dict(data)  # type: ignore
                return target_type(**data)

            # Try other common fields
            common_fields = ["text", "message", "content"]
            sig = inspect.signature(target_type)
            params = list(sig.parameters.keys())
            for field in common_fields:
                if field in params:
                    try:
                        return target_type(**{field: input_str})
                    except Exception as e:
                        logger.debug(f"Failed to create SerializationMixin with field '{field}': {e}")
                        continue
        except Exception as e:
            logger.debug(f"Failed to parse string as SerializationMixin: {e}")

    # Dataclasses
    if is_dataclass(target_type):
        try:
            # Try parsing as JSON
            if input_str.strip().startswith("{"):
                data = json.loads(input_str)
                return target_type(**data)

            # Try common field names
            common_fields = ["text", "message", "content", "input", "data"]
            for field in common_fields:
                try:
                    return target_type(**{field: input_str})
                except Exception as e:
                    logger.debug(f"Failed to create dataclass with field '{field}': {e}")
                    continue
        except Exception as e:
            logger.debug(f"Failed to parse string as dataclass: {e}")

    # Fallback: return original string
    return input_str


def _parse_dict_input(input_dict: dict[str, Any], target_type: type) -> Any:
    """Parse dict input to target type.

    Args:
        input_dict: Input dictionary
        target_type: Target type

    Returns:
        Parsed input or original dict
    """
    # Handle primitive types - extract from common field names
    if target_type in (str, int, float, bool):
        try:
            # If it's already the right type, return as-is
            if isinstance(input_dict, target_type):
                return input_dict

            # Try "input" field first (common for workflow inputs)
            if "input" in input_dict:
                return target_type(input_dict["input"])

            # If single-key dict, extract the value
            if len(input_dict) == 1:
                value = next(iter(input_dict.values()))
                return target_type(value)

            # Otherwise, return as-is
            return input_dict
        except (ValueError, TypeError) as e:
            logger.debug(f"Failed to convert dict to {target_type}: {e}")
            return input_dict

    # If target is dict, return as-is
    if target_type is dict:
        return input_dict

    # Pydantic models
    if hasattr(target_type, "model_validate"):
        try:
            return target_type.model_validate(input_dict)  # type: ignore
        except Exception as e:
            logger.debug(f"Failed to validate dict as Pydantic model: {e}")

    # SerializationMixin
    if is_serialization_mixin(target_type):
        try:
            if target_type is Message:
                return _build_message_from_legacy_payload(input_dict)
            if hasattr(target_type, "from_dict"):
                return target_type.from_dict(input_dict)  # type: ignore
            return target_type(**input_dict)
        except Exception as e:
            logger.debug(f"Failed to parse dict as SerializationMixin: {e}")

    # Dataclasses
    if is_dataclass(target_type):
        try:
            return target_type(**input_dict)
        except Exception as e:
            logger.debug(f"Failed to parse dict as dataclass: {e}")

    # Fallback: return original dict
    return input_dict
