"""
Visibe SDK - AI Agent Observability Platform

This package provides observability tools for AI agents built with
CrewAI, LangChain, AutoGen, and other frameworks.

Usage:
    from visibe import Visibe
    from visibe.integrations.crewai import CrewAIIntegration

    obs = Visibe(api_url="http://localhost:3000")
    tracker = CrewAIIntegration(obs)

    with tracker.track(crew, name="my-workflow"):
        result = crew.kickoff()

Or use auto-instrumentation:
    import visibe
    visibe.init()  # Auto-instruments all detected frameworks
    # ... create clients, they're automatically traced
    visibe.shutdown()  # Optional cleanup
"""

import logging
from typing import Optional

from .client import Visibe
from .__version__ import __version__

logger = logging.getLogger("visibe")

# Global state for auto-instrumentation
_global_client: Optional[Visibe] = None
_auto_patched_frameworks: list = []


def _detect_frameworks() -> dict:
    """Detect which frameworks are installed.
    
    Returns:
        dict: {framework_name: True/False} for each supported framework
    """
    detectors = {
        "openai": lambda: __import__("openai"),
        "langchain": lambda: __import__("langchain_core"),
        "langgraph": lambda: __import__("langgraph"),
        "crewai": lambda: __import__("crewai"),
        "autogen": lambda: __import__("autogen_agentchat"),
        "bedrock": lambda: __import__("boto3"),
    }
    
    available = {}
    for name, detector in detectors.items():
        try:
            detector()
            available[name] = True
        except ImportError:
            available[name] = False
    
    return available


def init(
    *,
    api_key: Optional[str] = None,
    api_url: Optional[str] = None,
    frameworks: Optional[list] = None,
    content_limit: Optional[int] = None,
    debug: bool = False,
) -> Visibe:
    """Initialize Visibe auto-instrumentation.
    
    Automatically patches all detected AI framework constructors so every newly
    created client is instrumented without manual calls to obs.instrument().
    
    Environment variables can override parameters:
        VISIBE_API_KEY: API key (if not provided via parameter)
        VISIBE_API_URL: API URL (if not provided via parameter)
        VISIBE_AUTO_INSTRUMENT: Comma-separated framework list (e.g., "openai,langgraph")
                               or "all" for all detected frameworks (default)
        VISIBE_CONTENT_LIMIT: Max chars for content (overrides parameter if set)
        VISIBE_DEBUG: Enable debug logging ("1" or "true")
    
    Args:
        api_key: Your API key (reads from VISIBE_API_KEY env if not provided)
        api_url: Override the API endpoint (reads from VISIBE_API_URL env if not provided)
        frameworks: List of framework names to patch (e.g., ["openai", "langgraph"]).
                   None = auto-detect and patch all installed frameworks.
                   Can be overridden by VISIBE_AUTO_INSTRUMENT env var.
        content_limit: Max chars for LLM/tool input/output in trace spans.
                      None = defaults. 0 = omit content entirely.
                      Can be overridden by VISIBE_CONTENT_LIMIT env var.
        debug: Enable debug logging.
               Can be enabled by VISIBE_DEBUG=1 env var.
    
    Returns:
        Visibe: The global client instance
    
    Example:
        >>> import visibe
        >>> visibe.init()  # Uses VISIBE_API_KEY from environment
        >>> from openai import OpenAI
        >>> client = OpenAI()  # automatically instrumented
        >>> response = client.chat.completions.create(
        ...     model="gpt-4o",
        ...     messages=[{"role": "user", "content": "Hello"}]
        ... )
    """
    import os
    
    global _global_client, _auto_patched_frameworks
    
    # Environment variable fallbacks (explicit parameters take precedence)
    if not debug:
        debug_env = os.getenv('VISIBE_DEBUG', '').lower()
        debug = debug_env in ('1', 'true', 'yes')
    
    if content_limit is None:
        env_content_limit = os.getenv('VISIBE_CONTENT_LIMIT')
        if env_content_limit is not None:
            try:
                content_limit = int(env_content_limit)
            except ValueError:
                logger.warning(f"Invalid VISIBE_CONTENT_LIMIT: {env_content_limit}")
    
    if frameworks is None:
        env_frameworks = os.getenv('VISIBE_AUTO_INSTRUMENT')
        if env_frameworks and env_frameworks.lower() != 'all':
            frameworks = [f.strip() for f in env_frameworks.split(',')]
    
    if _global_client is not None:
        logger.warning(
            "[Visibe] Already initialized — call shutdown() first to re-init"
        )
        print("[Visibe] Already initialized — call shutdown() first to re-init")
        return _global_client
    
    if debug:
        logging.basicConfig(level=logging.DEBUG)
        logging.getLogger("visibe").setLevel(logging.DEBUG)
    
    # Create global client (Visibe constructor reads VISIBE_API_KEY/VISIBE_API_URL if None)
    _global_client = Visibe(api_key=api_key, api_url=api_url)
    
    # Determine which frameworks to patch
    available = _detect_frameworks()
    
    if frameworks is None:
        # Auto-detect: patch all installed frameworks
        targets = [name for name, installed in available.items() if installed]
    else:
        # Use explicit allowlist, but warn if not installed
        targets = []
        for name in frameworks:
            if available.get(name):
                targets.append(name)
            else:
                logger.warning(f"[Visibe] Framework '{name}' not installed — skipping")
    
    # Patchers per framework (langchain excluded — requires explicit instrument())
    patchers = {
        "openai": _patch_openai,
        "langgraph": _patch_langgraph,
        "langchain": _patch_langchain,
        "crewai": _patch_crewai,
        "autogen": _patch_autogen,
        "bedrock": _patch_bedrock,
    }
    
    # Apply patches
    for fw in targets:
        patcher = patchers.get(fw)
        if patcher:
            try:
                patcher(_global_client, content_limit=content_limit)
                _auto_patched_frameworks.append(fw)
                logger.debug(f"[Visibe] Auto-patched {fw}")
            except Exception as e:
                logger.error(f"[Visibe] Failed to patch {fw}: {e}")
    
    if _auto_patched_frameworks:
        msg = f"[Visibe] Auto-instrumented: {', '.join(_auto_patched_frameworks)}"
        print(msg)
        logger.debug(msg)
    else:
        print("[Visibe] Initialized (no supported frameworks detected)")
    
    # Register atexit hook so users never need to call shutdown() manually
    import atexit
    atexit.register(shutdown)
    
    return _global_client


def shutdown() -> None:
    """Shut down Visibe auto-instrumentation.
    
    Restores all patched constructors to their original state and clears
    the global client. Safe to call multiple times.
    
    Example:
        >>> visibe.shutdown()
    """
    global _global_client, _auto_patched_frameworks
    
    if _global_client is None:
        logger.debug("[Visibe] Not initialized — nothing to shutdown")
        return
    
    # Unpatchers per framework
    unpatchers = {
        "openai": _unpatch_openai,
        "langgraph": _unpatch_langgraph,
        "langchain": _unpatch_langchain,
        "crewai": _unpatch_crewai,
        "autogen": _unpatch_autogen,
        "bedrock": _unpatch_bedrock,
    }
    
    for fw in _auto_patched_frameworks:
        unpatcher = unpatchers.get(fw)
        if unpatcher:
            try:
                unpatcher()
                logger.debug(f"[Visibe] Unpatched {fw}")
            except Exception as e:
                logger.error(f"[Visibe] Failed to unpatch {fw}: {e}")
    
    # Flush any remaining spans
    if _global_client:
        _global_client.flush_spans()
    
    _global_client = None
    _auto_patched_frameworks.clear()
    print("[Visibe] Shutdown complete")


_original_openai_init = None
_original_async_openai_init = None
_original_pregel_init = None
_original_crew_init = None
_original_autogen_client_init = None
_original_runnable_sequence_init = None
_original_runnable_parallel_init = None
_original_boto3_session_client = None


def _patch_openai(client: Visibe, *, content_limit=None):
    """Patch OpenAI constructors for auto-instrumentation."""
    global _original_openai_init, _original_async_openai_init
    
    try:
        import openai
    except ImportError:
        logger.warning("OpenAI not installed — skipping patch")
        return
    
    # Patch sync OpenAI.__init__
    _original_openai_init = openai.OpenAI.__init__
    
    def _patched_openai_init(self, *args, **kwargs):
        _original_openai_init(self, *args, **kwargs)
        client._get_openai_integration().instrument(self)
    
    openai.OpenAI.__init__ = _patched_openai_init
    
    # Patch async OpenAI.__init__
    _original_async_openai_init = openai.AsyncOpenAI.__init__
    
    def _patched_async_openai_init(self, *args, **kwargs):
        _original_async_openai_init(self, *args, **kwargs)
        client._get_openai_integration().instrument(self)
    
    openai.AsyncOpenAI.__init__ = _patched_async_openai_init
    logger.debug("[Visibe] OpenAI patch applied")


def _unpatch_openai():
    """Restore original OpenAI constructors."""
    global _original_openai_init, _original_async_openai_init
    
    if _original_openai_init is None:
        return
    
    try:
        import openai
        if _original_openai_init:
            openai.OpenAI.__init__ = _original_openai_init
        if _original_async_openai_init:
            openai.AsyncOpenAI.__init__ = _original_async_openai_init
        _original_openai_init = None
        _original_async_openai_init = None
        logger.debug("[Visibe] OpenAI patch removed")
    except ImportError:
        pass


def _patch_langgraph(client: Visibe, *, content_limit=None):
    """Patch LangGraph constructors for auto-instrumentation."""
    global _original_pregel_init
    
    try:
        from langgraph.pregel import Pregel
    except ImportError:
        logger.warning("LangGraph not installed — skipping patch")
        return
    
    # Patch Pregel.__init__
    _original_pregel_init = Pregel.__init__
    
    def _patched_pregel_init(self, *args, **kwargs):
        _original_pregel_init(self, *args, **kwargs)
        client._get_langgraph_integration().instrument(
            self, content_limit=content_limit
        )
    
    Pregel.__init__ = _patched_pregel_init
    logger.debug("[Visibe] LangGraph patch applied")


def _unpatch_langgraph():
    """Restore original LangGraph constructors."""
    global _original_pregel_init
    
    if _original_pregel_init is None:
        return
    
    try:
        from langgraph.pregel import Pregel
        Pregel.__init__ = _original_pregel_init
        _original_pregel_init = None
        logger.debug("[Visibe] LangGraph patch removed")
    except ImportError:
        pass


def _wrap_dynamic_runnable(runnable, client: Visibe, content_limit=None, runnable_type="Runnable"):
    """Wrap methods on a dynamically created LangChain runnable (e.g., from pipe operator).
    
    Uses object.__setattr__ to bypass Pydantic validation when assigning wrapped methods.
    Tracks the runnable in _instrumented_runnables and wraps invoke/ainvoke methods.
    
    Note: LangChain flattens pipe chains — `(a | b) | c` creates a NEW sequence
    with all steps flattened. The intermediate sequence from `a | b` becomes orphaned.
    We instrument every RunnableSequence created, but since users only invoke the
    final chain, intermediate ones never produce traces.
    """
    try:
        from visibe.integrations.langchain import _instrumented_runnables
    except ImportError:
        return
    
    runnable_id = id(runnable)
    if runnable_id in _instrumented_runnables:
        return  # Already instrumented
    
    # Get LangChain integration to use its wrapping logic
    integration = client._get_langchain_integration()
    trace_name = f"{runnable_type}-pipe-chain"
    
    # Save originals
    original_invoke = runnable.invoke
    original_ainvoke = getattr(runnable, 'ainvoke', None)
    
    _instrumented_runnables[runnable_id] = {
        'invoke': original_invoke,
        'ainvoke': original_ainvoke,
        'name': trace_name,
    }
    
    # Wrap invoke method
    def wrapped_invoke(input, config=None, **kwargs):
        from visibe.integrations.langchain import LangGraphCallback, _active_callback_ctx, _merge_callback_into_config
        import visibe.integrations.langchain as _lc_mod
        from datetime import datetime, timezone
        
        # If a parent trace is already active, propagate it instead of creating a new trace.
        parent_cb = _lc_mod._active_root_callback
        if parent_cb is not None:
            config = _merge_callback_into_config(config, parent_cb)
            return original_invoke(input, config=config, **kwargs)
        
        callback = LangGraphCallback(content_limit=content_limit)
        integration.client.create_trace({
            'trace_id': callback.trace_id,
            'name': trace_name,
            'framework': 'langchain',
            'started_at': datetime.now(timezone.utc).isoformat(),
        })
        
        callback._span_sender = lambda span: integration.client.queue_span(
            callback.trace_id, span
        )
        
        _lc_mod._active_root_callback = callback
        token = _active_callback_ctx.set(callback)
        config = _merge_callback_into_config(config, callback)
        
        try:
            result = original_invoke(input, config=config, **kwargs)
            integration._complete_trace(callback, 'completed', trace_name)
            return result
        except Exception as exc:
            callback.status = "failed"
            callback.errors.append({
                'timestamp': datetime.now(timezone.utc).isoformat(),
                'error_type': type(exc).__name__,
                'message': str(exc),
                'agent_name': callback.current_agent or 'unknown',
            })
            integration._complete_trace(callback, 'failed', trace_name)
            raise
        finally:
            _lc_mod._active_root_callback = None
            _active_callback_ctx.reset(token)
    
    # Use object.__setattr__ to bypass Pydantic validation
    object.__setattr__(runnable, 'invoke', wrapped_invoke)
    
    # Wrap ainvoke if available
    if original_ainvoke is not None:
        async def wrapped_ainvoke(input, config=None, **kwargs):
            from visibe.integrations.langchain import LangGraphCallback, _active_callback_ctx, _merge_callback_into_config
            import visibe.integrations.langchain as _lc_mod
            from datetime import datetime, timezone
            
            # If a parent trace is already active, propagate it.
            parent_cb = _lc_mod._active_root_callback
            if parent_cb is not None:
                config = _merge_callback_into_config(config, parent_cb)
                return await original_ainvoke(input, config=config, **kwargs)
            
            callback = LangGraphCallback(content_limit=content_limit)
            integration.client.create_trace({
                'trace_id': callback.trace_id,
                'name': trace_name,
                'framework': 'langchain',
                'started_at': datetime.now(timezone.utc).isoformat(),
            })
            
            callback._span_sender = lambda span: integration.client.queue_span(
                callback.trace_id, span
            )
            
            _lc_mod._active_root_callback = callback
            token = _active_callback_ctx.set(callback)
            config = _merge_callback_into_config(config, callback)
            
            try:
                result = await original_ainvoke(input, config=config, **kwargs)
                integration._complete_trace(callback, 'completed', trace_name)
                return result
            except Exception as exc:
                callback.status = "failed"
                callback.errors.append({
                    'timestamp': datetime.now(timezone.utc).isoformat(),
                    'error_type': type(exc).__name__,
                    'message': str(exc),
                    'agent_name': callback.current_agent or 'unknown',
                })
                integration._complete_trace(callback, 'failed', trace_name)
                raise
            finally:
                _lc_mod._active_root_callback = None
                _active_callback_ctx.reset(token)
        
        object.__setattr__(runnable, 'ainvoke', wrapped_ainvoke)


def _patch_langchain(client: Visibe, *, content_limit=None):
    """Patch LangChain dynamic runnables for auto-instrumentation.
    
    Patches RunnableSequence and RunnableParallel constructors so chains built
    with the pipe operator (|) are automatically instrumented. Uses object.__setattr__
    to bypass Pydantic validation when assigning wrapped methods.
    """
    global _original_runnable_sequence_init, _original_runnable_parallel_init
    
    try:
        from langchain_core.runnables import RunnableSequence, RunnableParallel
        from datetime import datetime, timezone
    except ImportError:
        logger.warning("LangChain not installed — skipping patch")
        return
    
    # Patch RunnableSequence.__init__ (created via pipe operator)
    _original_runnable_sequence_init = RunnableSequence.__init__
    
    def _patched_sequence_init(self, *args, **kwargs):
        _original_runnable_sequence_init(self, *args, **kwargs)
        _wrap_dynamic_runnable(self, client, content_limit, "RunnableSequence")
    
    RunnableSequence.__init__ = _patched_sequence_init
    
    # Patch RunnableParallel.__init__ (created via pipe operator for parallel branches)
    _original_runnable_parallel_init = RunnableParallel.__init__
    
    def _patched_parallel_init(self, *args, **kwargs):
        _original_runnable_parallel_init(self, *args, **kwargs)
        _wrap_dynamic_runnable(self, client, content_limit, "RunnableParallel")
    
    RunnableParallel.__init__ = _patched_parallel_init
    logger.debug("[Visibe] LangChain patch applied")


def _unpatch_langchain():
    """Restore original LangChain constructors."""
    global _original_runnable_sequence_init, _original_runnable_parallel_init
    
    if _original_runnable_sequence_init is None:
        return
    
    try:
        from langchain_core.runnables import RunnableSequence, RunnableParallel
        
        if _original_runnable_sequence_init:
            RunnableSequence.__init__ = _original_runnable_sequence_init
            _original_runnable_sequence_init = None
        
        if _original_runnable_parallel_init:
            RunnableParallel.__init__ = _original_runnable_parallel_init
            _original_runnable_parallel_init = None
        
        # Clear the instrumented runnables tracking dict to avoid stale entries
        try:
            from visibe.integrations.langchain import _instrumented_runnables
            _instrumented_runnables.clear()
        except ImportError:
            pass
        
        logger.debug("[Visibe] LangChain patch removed")
    except ImportError:
        pass


def _patch_crewai(client: Visibe, *, content_limit=None):
    """Patch CrewAI constructors for auto-instrumentation."""
    global _original_crew_init
    
    try:
        from crewai import Crew
    except ImportError:
        logger.warning("CrewAI not installed — skipping patch")
        return
    
    # Patch Crew.__init__
    _original_crew_init = Crew.__init__
    
    def _patched_crew_init(self, *args, **kwargs):
        _original_crew_init(self, *args, **kwargs)
        client._get_crewai_integration().instrument(self)
    
    Crew.__init__ = _patched_crew_init
    logger.debug("[Visibe] CrewAI patch applied")


def _unpatch_crewai():
    """Restore original CrewAI constructors."""
    global _original_crew_init
    
    if _original_crew_init is None:
        return
    
    try:
        from crewai import Crew
        Crew.__init__ = _original_crew_init
        _original_crew_init = None
        
        # Clear the instrumented crews tracking dict to avoid stale entries
        try:
            from visibe.integrations.crewai import _instrumented_crews
            _instrumented_crews.clear()
        except ImportError:
            pass
        
        logger.debug("[Visibe] CrewAI patch removed")
    except ImportError:
        pass


def _patch_autogen(client: Visibe, *, content_limit=None):
    """Patch AutoGen constructors for auto-instrumentation."""
    global _original_autogen_client_init
    
    try:
        from autogen_ext.models.openai import OpenAIChatCompletionClient
    except ImportError:
        logger.warning("AutoGen not installed — skipping patch")
        return
    
    # Patch OpenAIChatCompletionClient.__init__
    _original_autogen_client_init = OpenAIChatCompletionClient.__init__
    
    def _patched_autogen_client_init(self, *args, **kwargs):
        _original_autogen_client_init(self, *args, **kwargs)
        client._get_autogen_integration().instrument(self)
    
    OpenAIChatCompletionClient.__init__ = _patched_autogen_client_init
    logger.debug("[Visibe] AutoGen patch applied")


def _unpatch_autogen():
    """Restore original AutoGen constructors."""
    global _original_autogen_client_init

    if _original_autogen_client_init is None:
        return

    try:
        from autogen_ext.models.openai import OpenAIChatCompletionClient
        OpenAIChatCompletionClient.__init__ = _original_autogen_client_init
        _original_autogen_client_init = None
        logger.debug("[Visibe] AutoGen patch removed")
    except ImportError:
        pass


def _patch_bedrock(client: Visibe, *, content_limit=None):
    """Patch boto3.Session.client() to auto-instrument bedrock-runtime clients.

    Patches at the Session level (not boto3.client directly) so we catch
    both boto3.client('bedrock-runtime') and session.client('bedrock-runtime').
    Non-bedrock services pass through with zero overhead beyond a string check.
    """
    global _original_boto3_session_client

    try:
        import boto3
    except ImportError:
        logger.warning("boto3 not installed — skipping Bedrock patch")
        return

    _original_boto3_session_client = boto3.Session.client

    def _patched_session_client(self, service_name, *args, **kwargs):
        result = _original_boto3_session_client(self, service_name, *args, **kwargs)
        if service_name == 'bedrock-runtime':
            client._get_bedrock_integration().instrument(result)
        return result

    boto3.Session.client = _patched_session_client
    logger.debug("[Visibe] Bedrock patch applied")


def _unpatch_bedrock():
    """Restore original boto3.Session.client()."""
    global _original_boto3_session_client

    if _original_boto3_session_client is None:
        return

    try:
        import boto3
        boto3.Session.client = _original_boto3_session_client
        _original_boto3_session_client = None
        logger.debug("[Visibe] Bedrock patch removed")
    except ImportError:
        pass


__all__ = ['Visibe', '__version__', 'init', 'shutdown']
