"""Sentinel value and function introspection for AI tasks.

This module provides the _ai sentinel used to mark LLM-generated outputs
in AI function definitions, and introspection tools to extract metadata
from those functions.

Example:
    >>> from elemai import ai, _ai
    >>>
    >>> @ai
    ... def analyze(text: str) -> str:
    ...     '''Analyze the text'''
    ...     thinking: str = _ai["Think carefully"]
    ...     return _ai
"""

import inspect
from typing import Any, Optional


class AISentinel:
    """Sentinel value representing an LLM-generated output.

    When used in a function body decorated with @ai, indicates that
    this value should be generated by the LLM. Supports subscript
    notation for adding descriptions to guide the LLM.

    Attributes:
        description: Optional description for the LLM output

    Example:
        >>> from elemai.sentinel import _ai
        >>> # Basic usage
        >>> sentinel = _ai
        >>> sentinel.description is None
        True
        >>>
        >>> # With description
        >>> sentinel_with_desc = _ai["Think step by step"]
        >>> sentinel_with_desc.description
        'Think step by step'
        >>>
        >>> # String representation
        >>> str(_ai)
        '_ai'
        >>> str(_ai["reason carefully"])
        "_ai['reason carefully']"
    """

    def __init__(self, description: Optional[str] = None):
        """Initialize the sentinel.

        Args:
            description: Optional description for the output

        Example:
            >>> from elemai.sentinel import AISentinel
            >>> s1 = AISentinel()
            >>> s1.description is None
            True
            >>> s2 = AISentinel("analyze this")
            >>> s2.description
            'analyze this'
        """
        self.description = description

    def __getitem__(self, description: str) -> 'AISentinel':
        """Support _ai["description"] syntax.

        Args:
            description: Description for the output

        Returns:
            AISentinel: New sentinel with the description

        Example:
            >>> from elemai.sentinel import _ai, AISentinel
            >>> result = _ai["Think carefully about this"]
            >>> result.description
            'Think carefully about this'
            >>> isinstance(result, AISentinel)
            True
        """
        return AISentinel(description=description)

    def __repr__(self) -> str:
        """Return string representation.

        Returns:
            str: String representation of the sentinel

        Example:
            >>> from elemai.sentinel import _ai
            >>> repr(_ai)
            '_ai'
            >>> repr(_ai["some description"])
            "_ai['some description']"
        """
        if self.description:
            return f"_ai[{self.description!r}]"
        return "_ai"


# The global _ai sentinel
_ai = AISentinel()


class FunctionIntrospector:
    """Introspects a function to extract AI task metadata.

    Analyzes a Python function to extract input parameters, output fields,
    docstrings, and type annotations. Used by the @ai decorator to build
    prompts automatically.

    Attributes:
        func: The function being introspected
        signature: The function's signature
        source: The function's source code (if available)

    Example:
        >>> from elemai.sentinel import FunctionIntrospector, _ai
        >>>
        >>> def example_func(text: str, count: int = 5) -> str:
        ...     '''Process the text'''
        ...     thinking: str = _ai["Analyze"]
        ...     return _ai
        >>>
        >>> inspector = FunctionIntrospector(example_func)
        >>> inspector.func.__name__
        'example_func'
    """

    def __init__(self, func):
        """Initialize the introspector.

        Args:
            func: Function to introspect

        Example:
            >>> def my_func(x: int) -> int:
            ...     return x * 2
            >>> from elemai.sentinel import FunctionIntrospector
            >>> inspector = FunctionIntrospector(my_func)
            >>> inspector.func.__name__
            'my_func'
        """
        self.func = func
        self.signature = inspect.signature(func)
        self.source = None
        try:
            self.source = inspect.getsource(func)
        except (OSError, TypeError):
            pass

    def get_input_fields(self):
        """Extract input fields from function signature.

        Returns:
            list: List of dicts with 'name', 'type', and 'default' keys

        Example:
            >>> from elemai.sentinel import FunctionIntrospector
            >>> def func(text: str, count: int = 5):
            ...     pass
            >>> inspector = FunctionIntrospector(func)
            >>> fields = inspector.get_input_fields()
            >>> len(fields)
            2
            >>> fields[0]['name']
            'text'
            >>> fields[0]['type']
            <class 'str'>
            >>> fields[1]['default']
            5
        """
        fields = []
        for name, param in self.signature.parameters.items():
            if name == 'self':
                continue

            field_type = param.annotation if param.annotation != inspect.Parameter.empty else Any
            fields.append({
                'name': name,
                'type': field_type,
                'default': param.default if param.default != inspect.Parameter.empty else None
            })
        return fields

    def get_output_fields(self):
        """Extract output fields by analyzing function body for _ai assignments.

        Parses the function source code to find variables assigned to _ai,
        extracting their names, types, and descriptions.

        Returns:
            list: List of dicts with 'name', 'type', 'description' keys

        Example:
            >>> from elemai.sentinel import FunctionIntrospector, _ai
            >>> def func(text: str) -> str:
            ...     thinking: str = _ai["Think carefully"]
            ...     return _ai
            >>> inspector = FunctionIntrospector(func)
            >>> fields = inspector.get_output_fields()
            >>> len(fields) >= 1
            True
            >>> any(f['name'] == 'thinking' for f in fields)
            True
        """
        fields = []

        if not self.source:
            # If we can't get source, infer from return annotation
            return_type = self.signature.return_annotation
            if return_type != inspect.Signature.empty:
                fields.append({
                    'name': 'result',
                    'type': return_type,
                    'description': None
                })
            return fields

        # Parse the source to find _ai assignments
        lines = self.source.split('\n')

        for line in lines:
            line = line.strip()

            # Match patterns like: thinking: str = _ai["description"]
            # or: thinking: str = _ai
            if '= _ai' in line and not line.startswith('#'):
                # Extract variable name and type
                parts = line.split('=')[0].strip()

                if ':' in parts:
                    var_name = parts.split(':')[0].strip()
                    type_hint = parts.split(':')[1].strip()

                    # Try to resolve type hint
                    var_type = self._resolve_type(type_hint)
                else:
                    var_name = parts.strip()
                    var_type = Any

                # Extract description if present
                description = None
                if '[' in line and ']' in line:
                    # Extract description from _ai["description"]
                    desc_match = line.split('[')[1].split(']')[0]
                    description = desc_match.strip('"\'')

                fields.append({
                    'name': var_name,
                    'type': var_type,
                    'description': description
                })

        # Also check return statement
        return_var = None
        for line in lines:
            line = line.strip()
            if line.startswith('return '):
                # Extract what's being returned
                return_part = line[7:].strip()
                if return_part == '_ai':
                    return_var = 'result'  # Direct return of _ai
                else:
                    # Check if returning a variable that was assigned to _ai
                    var_name = return_part.split()[0].strip('(),')
                    if any(f['name'] == var_name for f in fields):
                        return_var = var_name
                break

        # If we found a direct 'return _ai', ensure we have a 'result' field
        if return_var == 'result':
            return_type = self.signature.return_annotation
            if return_type != inspect.Signature.empty:
                if not any(f['name'] == 'result' for f in fields):
                    fields.append({
                        'name': 'result',
                        'type': return_type,
                        'description': None
                    })

        # Mark which field is the final return value
        if return_var and fields:
            for field in fields:
                if field['name'] == return_var:
                    field['is_final_output'] = True
                else:
                    field['is_final_output'] = False

        return fields

    def _resolve_type(self, type_str: str):
        """Attempt to resolve a type string to an actual type.

        Args:
            type_str: String representation of a type

        Returns:
            type: Resolved type object, or Any if resolution fails

        Example:
            >>> from elemai.sentinel import FunctionIntrospector
            >>> def dummy(): pass
            >>> inspector = FunctionIntrospector(dummy)
            >>> inspector._resolve_type('str')
            <class 'str'>
            >>> inspector._resolve_type('int')
            <class 'int'>
        """
        # Get the function's globals for type resolution
        func_globals = self.func.__globals__

        # Common built-in types
        builtin_types = {
            'str': str,
            'int': int,
            'float': float,
            'bool': bool,
            'list': list,
            'dict': dict,
            'Any': Any,
        }

        # Try built-ins first
        if type_str in builtin_types:
            return builtin_types[type_str]

        # Try to evaluate in function's namespace
        try:
            return eval(type_str, func_globals)
        except:
            return Any

    def get_instruction(self):
        """Extract instruction from docstring.

        Uses the first line of the function's docstring as the instruction,
        or generates a default if no docstring exists.

        Returns:
            str: Instruction string for the AI

        Example:
            >>> from elemai.sentinel import FunctionIntrospector
            >>> def func1():
            ...     '''Analyze the input text'''
            ...     pass
            >>> FunctionIntrospector(func1).get_instruction()
            'Analyze the input text'
            >>>
            >>> def func2():
            ...     pass
            >>> FunctionIntrospector(func2).get_instruction()
            'Perform func2'
        """
        doc = inspect.getdoc(self.func)
        if doc:
            # Use first line as instruction
            return doc.split('\n')[0].strip()
        return f"Perform {self.func.__name__}"

    def get_full_docstring(self):
        """Get the complete docstring.

        Returns:
            str: Full docstring or empty string if none exists

        Example:
            >>> from elemai.sentinel import FunctionIntrospector
            >>> def func():
            ...     '''First line.
            ...
            ...     Second paragraph.'''
            ...     pass
            >>> doc = FunctionIntrospector(func).get_full_docstring()
            >>> 'First line.' in doc
            True
            >>> 'Second paragraph.' in doc
            True
        """
        return inspect.getdoc(self.func) or ""

    def introspect(self):
        """Perform full introspection of the function.

        Extracts all metadata needed to build an AI task from the function.

        Returns:
            dict: Dictionary with keys: 'fn_name', 'instruction', 'doc',
                  'input_fields', 'output_fields', 'return_type'

        Example:
            >>> from elemai.sentinel import FunctionIntrospector, _ai
            >>> def analyze(text: str, context: str = "") -> str:
            ...     '''Analyze the text carefully'''
            ...     thinking: str = _ai
            ...     return _ai
            >>>
            >>> metadata = FunctionIntrospector(analyze).introspect()
            >>> metadata['fn_name']
            'analyze'
            >>> metadata['instruction']
            'Analyze the text carefully'
            >>> len(metadata['input_fields'])
            2
            >>> metadata['return_type']
            <class 'str'>
        """
        return {
            'fn_name': self.func.__name__,
            'instruction': self.get_instruction(),
            'doc': self.get_full_docstring(),
            'input_fields': self.get_input_fields(),
            'output_fields': self.get_output_fields(),
            'return_type': self.signature.return_annotation,
        }
