"""
Core Visibe SDK client.
Used by all framework integrations.
"""
import logging
import os
import sys
import time
import warnings
from datetime import datetime, timezone
from typing import Optional
from uuid import uuid4

from .api import APIClient, SpanBatcher
from .context import _SessionContext, _active_session

# Production API URL (hardcoded - users never need to change this)
_PRODUCTION_API_URL = "https://api.visibe.ai"


class VisibleSDKWarning(UserWarning):
    """Warnings issued by the Visibe SDK. Users can suppress with:
    warnings.filterwarnings("ignore", category=visibe.VisibleSDKWarning)
    """


_orig_formatwarning = warnings.formatwarning

def _formatwarning(message, category, filename, lineno, line=None):
    if issubclass(category, VisibleSDKWarning):
        return f"{message}\n"
    return _orig_formatwarning(message, category, filename, lineno, line)

warnings.formatwarning = _formatwarning


def _yellow(text: str) -> str:
    """Apply yellow ANSI color when stderr is a real terminal."""
    return f"\033[33m{text}\033[0m" if sys.stderr.isatty() else text


def _is_openai_client(obj) -> bool:
    """Duck-type check: does this look like an OpenAI client?"""
    return hasattr(obj, 'chat') and hasattr(obj.chat, 'completions')


def _is_langchain_runnable(obj) -> bool:
    """Duck-type check: does this look like a LangChain/LangGraph runnable?"""
    return hasattr(obj, 'invoke') and hasattr(obj, 'get_graph')


def _is_langgraph_runnable(obj) -> bool:
    """Duck-type check: does this look like a LangGraph compiled graph?"""
    try:
        from langgraph.pregel import Pregel
        return isinstance(obj, Pregel)
    except ImportError:
        return False


def _is_crewai_crew(obj) -> bool:
    """Duck-type check: does this look like a CrewAI Crew?"""
    return hasattr(obj, 'kickoff') and hasattr(obj, 'agents') and hasattr(obj, 'tasks')


def _is_autogen_model_client(obj) -> bool:
    """Duck-type check: does this look like an AutoGen model client?"""
    return (
        hasattr(obj, 'create') and
        hasattr(obj, 'create_stream') and
        hasattr(obj, 'model_info')
    )


def _is_bedrock_client(obj) -> bool:
    """Check if this is a Bedrock runtime client.

    Uses boto3's service metadata for reliable detection rather than
    pure duck-typing (which could false-positive on any object with
    converse/converse_stream/invoke_model methods).
    """
    # Primary: use boto3 service metadata (most reliable)
    try:
        if hasattr(obj, 'meta') and hasattr(obj.meta, 'service_model'):
            return obj.meta.service_model.service_name == 'bedrock-runtime'
    except Exception:
        pass

    # Fallback: duck-type check for non-standard boto3 wrappers
    return (
        hasattr(obj, 'converse') and
        hasattr(obj, 'converse_stream') and
        hasattr(obj, 'invoke_model') and
        hasattr(obj, 'meta')  # all boto3 clients have .meta
    )


# SDK logger
logger = logging.getLogger("visibe")
logger.addHandler(logging.NullHandler())  # No output by default


class _InstrumentContext:
    """Optional context manager for auto-cleanup after instrument().

    Instrumentation happens eagerly in instrument(), not deferred to __enter__.
    If the return value is ignored, behavior is identical to calling instrument()
    without a context manager — the caller must call uninstrument() manually.
    """

    def __init__(self, visibe: 'Visibe', client):
        self._visibe = visibe
        self._client = client

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._visibe.uninstrument(self._client)
        return False


class _SessionContextManager:
    """Returned by Visibe.session().

    Implements both __enter__/__exit__ (sync) and __aenter__/__aexit__ (async)
    so the same call works with both 'with' and 'async with'.
    """

    def __init__(self, visibe: 'Visibe', name: str):
        self._visibe      = visibe
        self._name        = name
        self._token       = None
        self._session_ctx = None
        self._start       = None

    def _begin(self) -> _SessionContext:
        self._start    = time.monotonic()
        trace_id       = str(uuid4())
        self._session_ctx = _SessionContext(trace_id=trace_id)

        # Create the session trace directly on the API client so we bypass the
        # create_trace() interception (which would be a no-op inside a session).
        self._visibe.api_client.create_trace({
            'trace_id':   trace_id,
            'name':       self._name,
            'framework':  'session',
            'started_at': datetime.now(timezone.utc).isoformat(),
            **(({'session_id': self._visibe.session_id})
               if self._visibe.session_id else {}),
        })

        self._token = _active_session.set(self._session_ctx)
        return self._session_ctx

    def _end(self, failed: bool) -> None:
        if self._token is not None:
            _active_session.reset(self._token)

        sc = self._session_ctx
        if sc is None:
            return

        duration_ms = int((time.monotonic() - self._start) * 1000)
        status      = 'failed' if failed else 'completed'

        self._visibe._batcher.flush()
        sent = self._visibe.api_client.complete_trace(sc.trace_id, {
            'status':              status,
            'ended_at':            datetime.now(timezone.utc).isoformat(),
            'duration_ms':         duration_ms,
            'llm_call_count':      sc.llm_call_count,
            'total_cost':          round(sc.total_cost, 6),
            'total_tokens':        sc.total_input + sc.total_output,
            'total_input_tokens':  sc.total_input,
            'total_output_tokens': sc.total_output,
        })

        total = sc.total_input + sc.total_output
        print(
            f"[Visibe] Session: {self._name} | "
            f"{sc.llm_call_count} LLM calls | "
            f"{total:,} tokens | "
            f"${sc.total_cost:.6f} | "
            f"{duration_ms / 1000:.1f}s | "
            f"status: {status} | "
            f"sent: {'OK' if sent else 'FAILED'}"
        )

    # Sync context manager
    def __enter__(self) -> _SessionContext:
        return self._begin()

    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
        self._end(failed=exc_type is not None)
        return False  # Never suppress exceptions

    # Async context manager — same logic, no extra await needed
    async def __aenter__(self) -> _SessionContext:
        return self._begin()

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:
        self._end(failed=exc_type is not None)
        return False


class Visibe:
    """
    Main SDK client for Visibe observability platform.

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

        # Production
        obs = Visibe(api_key="sk_live_abc123")

        # Or use environment variable
        export VISIBE_API_KEY=sk_live_abc123
        obs = Visibe()

    Environment Variables:
        VISIBE_API_KEY: Your API key (required for production)

    Example:
        >>> import visibe
        >>> obs = visibe.Visibe(api_key="sk_live_abc123")
        >>> # Traces automatically sent to https://api.visibe.ai
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        session_id: Optional[str] = None,
        timeout: int = 10,
        api_url: Optional[str] = None  # Hidden parameter for testing
    ):
        """
        Initialize Visibe client.

        Args:
            api_key: Your API key (reads from VISIBE_API_KEY env if not provided)
            session_id: Optional session ID for grouping traces
            timeout: Request timeout in seconds (default: 10)
            api_url: Override API URL (for internal testing only - not documented)

        Example:
            >>> obs = Visibe(api_key="sk_live_abc123")
            >>> obs = Visibe()  # Uses VISIBE_API_KEY environment variable
        """
        # API key from parameter or environment variable
        self.api_key = api_key or os.getenv('VISIBE_API_KEY')
        
        # Warn if no API key provided
        if not self.api_key:
            warnings.warn(
                _yellow("[Visibe] No API key provided — tracing is disabled. "
                        "Set VISIBE_API_KEY env var or pass api_key= to enable."),
                VisibleSDKWarning,
                stacklevel=2,
            )

        # Validate API key format (if provided)
        if self.api_key and not self.api_key.startswith(('sk_live_', 'sk_test_')):
            warnings.warn(
                _yellow("[Visibe] API key format unrecognized "
                        "(expected sk_live_* or sk_test_*)."),
                VisibleSDKWarning,
                stacklevel=2,
            )

        # Session ID (parameter only, not env var)
        self.session_id = session_id
        
        # Timeout validation
        if timeout <= 0:
            raise ValueError(f"timeout must be positive, got: {timeout}")
        self.timeout = timeout
        
        # API URL - use production by default, allow override via env or parameter
        self.api_url = api_url or os.getenv('VISIBE_API_URL') or _PRODUCTION_API_URL
        
        # Validate API URL format
        if not self.api_url.startswith(('http://', 'https://')):
            raise ValueError(
                f"api_url must start with http:// or https://, got: {self.api_url}"
            )
        
        # Normalize API URL
        self.api_url = self.api_url.rstrip('/')
        
        # Initialize API client and span batcher
        self.api_client = APIClient(self.api_url, self.api_key, self.timeout)
        self._batcher = SpanBatcher(self.api_client)

        # Fire-and-forget background key validation.
        # Runs ~80ms after startup — well before the 2s batcher flush window.
        # On invalid key: disables tracing and prints a warning.
        # On network error: stays enabled (lenient).
        if self.api_client._enabled:
            import threading
            def _validate():
                if not self.api_client.validate_key():
                    self.api_client._enabled = False
                    warnings.warn(
                        _yellow("[Visibe] Invalid API key — tracing is disabled. "
                                "Check your VISIBE_API_KEY."),
                        VisibleSDKWarning,
                        stacklevel=1,  # stacklevel=2 is meaningless in a background thread
                    )
            threading.Thread(target=_validate, daemon=True).start()

        # Log initialization at debug level
        logger.debug(f"Visibe initialized: url={self.api_url}, has_key={bool(self.api_key)}")

    def create_trace(self, trace_data: dict) -> bool:
        """Create a trace header on the backend.

        When a session is active the integration's sub-trace is suppressed —
        the session owns the single trace for this request/context.
        """
        session = _active_session.get()
        if session is not None:
            # Session manages the trace; don't create a new one.
            return True
        if self.session_id and 'session_id' not in trace_data:
            trace_data['session_id'] = self.session_id
        return self.api_client.create_trace(trace_data)

    def send_span(self, trace_id: str, span_data: dict) -> bool:
        """Send a single span directly (bypasses batcher).

        Redirects to the session trace when a session is active.
        """
        session = _active_session.get()
        if session is not None:
            return self.api_client.send_span(session.trace_id, span_data)
        return self.api_client.send_span(trace_id, span_data)

    def queue_span(self, trace_id: str, span_data: dict):
        """Queue a span for batched sending. Non-blocking.

        Redirects to the session trace when a session is active.
        """
        session = _active_session.get()
        if session is not None:
            self._batcher.add(session.trace_id, span_data)
            return
        self._batcher.add(trace_id, span_data)

    def flush_spans(self):
        """Flush all queued spans immediately."""
        self._batcher.flush()

    def complete_trace(self, trace_id: str, summary: dict) -> bool:
        """Flush pending spans, then mark trace as complete.

        When a session is active, accumulates stats into the session context
        instead of completing the trace — the session completes it on exit.
        """
        session = _active_session.get()
        if session is not None:
            session.llm_call_count += summary.get('llm_call_count', 0)
            session.total_input    += summary.get('total_input_tokens', 0)
            session.total_output   += summary.get('total_output_tokens', 0)
            session.total_cost     += summary.get('total_cost', 0.0)
            return True
        self._batcher.flush()
        return self.api_client.complete_trace(trace_id, summary)

    def _get_openai_integration(self):
        """Lazily create and cache the OpenAI integration."""
        if not hasattr(self, '_openai_integration'):
            from .integrations.openai import OpenAIIntegration
            self._openai_integration = OpenAIIntegration(self)
        return self._openai_integration

    def _get_langchain_integration(self):
        """Lazily create and cache the LangChain integration."""
        if not hasattr(self, '_langchain_integration'):
            from .integrations.langchain import LangChainIntegration
            self._langchain_integration = LangChainIntegration(self)
        return self._langchain_integration

    def _get_langgraph_integration(self):
        """Lazily create and cache the LangGraph integration."""
        if not hasattr(self, '_langgraph_integration'):
            from .integrations.langgraph import LangGraphIntegration
            self._langgraph_integration = LangGraphIntegration(self)
        return self._langgraph_integration

    def _get_crewai_integration(self):
        """Lazily create and cache the CrewAI integration."""
        if not hasattr(self, '_crewai_integration'):
            from .integrations.crewai import CrewAIIntegration
            self._crewai_integration = CrewAIIntegration(self)
        return self._crewai_integration

    def _get_autogen_integration(self):
        """Lazily create and cache the AutoGen integration."""
        if not hasattr(self, '_autogen_integration'):
            from .integrations.autogen import AutoGenIntegration
            self._autogen_integration = AutoGenIntegration(self)
        return self._autogen_integration

    def _get_bedrock_integration(self):
        """Lazily create and cache the Bedrock integration."""
        if not hasattr(self, '_bedrock_integration'):
            from .integrations.bedrock import BedrockIntegration
            self._bedrock_integration = BedrockIntegration(self)
        return self._bedrock_integration

    def instrument(self, client, *, name: Optional[str] = None,
                   content_limit: Optional[int] = None):
        """Auto-instrument a client, runnable, or crew.

        Supports OpenAI clients, LangChain/LangGraph runnables, CrewAI crews,
        and AutoGen model clients.
        Returns a context manager for optional automatic cleanup.

        Args:
            client: An OpenAI client, LangChain/LangGraph runnable, or CrewAI Crew
            name: Optional trace name. Shown in the Visibe dashboard.
            content_limit: Max chars for LLM/tool input/output in trace spans.
                          None = defaults. 0 = omit content entirely.

        Returns:
            _InstrumentContext: Use with 'with' for auto-cleanup, or ignore.

        Example:
            >>> obs = Visibe()
            >>> with obs.instrument(crew, name="my-workflow"):
            ...     crew.kickoff()
            >>> # Or classic style:
            >>> obs.instrument(graph, name="trip-planner")
            >>> graph.invoke(...)
            >>> obs.uninstrument(graph)
        """
        if _is_openai_client(client):
            self._get_openai_integration().instrument(client, name=name)
        elif _is_crewai_crew(client):
            self._get_crewai_integration().instrument(client, name=name)
        elif _is_autogen_model_client(client):
            self._get_autogen_integration().instrument(client, name=name)
        elif _is_bedrock_client(client):
            self._get_bedrock_integration().instrument(client, name=name)
        elif _is_langchain_runnable(client):
            if _is_langgraph_runnable(client):
                self._get_langgraph_integration().instrument(
                    client, name=name, content_limit=content_limit
                )
            else:
                self._get_langchain_integration().instrument(
                    client, name=name, content_limit=content_limit
                )
        else:
            raise TypeError(
                f"Unsupported type for instrument(): {type(client).__name__}. "
                f"Expected an OpenAI client, Bedrock client, "
                f"LangChain/LangGraph runnable, CrewAI Crew, or AutoGen model client."
            )

        if name:
            print(f'[Visibe] Tracking enabled for "{name}"')
        else:
            print('[Visibe] Tracking enabled')

        return _InstrumentContext(self, client)

    def uninstrument(self, client):
        """Remove instrumentation from a previously instrumented client, runnable, or crew.

        Args:
            client: A previously instrumented OpenAI client, LangChain/LangGraph runnable, or CrewAI Crew
        """
        if _is_openai_client(client):
            self._get_openai_integration().uninstrument(client)
        elif _is_crewai_crew(client):
            self._get_crewai_integration().uninstrument(client)
        elif _is_autogen_model_client(client):
            self._get_autogen_integration().uninstrument(client)
        elif _is_bedrock_client(client):
            self._get_bedrock_integration().uninstrument(client)
        elif _is_langchain_runnable(client):
            if _is_langgraph_runnable(client):
                self._get_langgraph_integration().uninstrument(client)
            else:
                self._get_langchain_integration().uninstrument(client)
        else:
            raise TypeError(
                f"Unsupported type for uninstrument(): {type(client).__name__}."
            )

    def session(self, name: str) -> '_SessionContextManager':
        """Group all LLM calls within a block into one shared trace.

        Works as both a sync and async context manager — no separate
        method needed for async code.

        Uses ContextVar internally so concurrent requests/tasks are
        fully isolated even in multi-user servers.

        Args:
            name: Trace name shown in the Visibe dashboard.

        Usage (sync — scripts, background jobs, Flask routes):
            with visibe.session("my-task"):
                client.chat.completions.create(...)

        Usage (async — FastAPI handlers, async agents):
            async with visibe.session("my-task"):
                await async_client.chat.completions.create(...)
        """
        return _SessionContextManager(self, name)

    def flask_middleware(self, app) -> None:
        """Register request-scoped tracing on a Flask application.

        Groups all LLM calls within each HTTP request into one trace.
        Uses ContextVar so concurrent requests are fully isolated —
        each request gets its own trace regardless of how many are
        in flight at once.

        Args:
            app: A Flask application instance.

        Usage:
            app = Flask(__name__)
            visibe.flask_middleware(app)
        """
        try:
            from flask import g, request as flask_request
        except ImportError:
            raise ImportError(
                "flask_middleware() requires Flask. "
                "Install it with: pip install flask"
            )

        visibe = self

        @app.before_request
        def _visibe_before():
            if not visibe.api_client._enabled:
                return
            trace_id = str(uuid4())
            g._visibe_session = _SessionContext(trace_id=trace_id)
            g._visibe_token   = _active_session.set(g._visibe_session)
            g._visibe_start   = time.monotonic()
            g._visibe_name    = f"{flask_request.method} {flask_request.path}"
            g._visibe_done    = False
            # Create the trace directly — bypasses the interception in create_trace()
            visibe.api_client.create_trace({
                'trace_id':   trace_id,
                'name':       g._visibe_name,
                'framework':  'http',
                'started_at': datetime.now(timezone.utc).isoformat(),
                **(({'session_id': visibe.session_id}) if visibe.session_id else {}),
            })

        @app.after_request
        def _visibe_after(response):
            _finish(exc=None, status_code=response.status_code)
            return response

        @app.teardown_request
        def _visibe_teardown(exc):
            # Fallback: fires when after_request didn't run (unhandled exception).
            _finish(exc=exc, status_code=500)

        def _finish(exc, status_code: int):
            sc = getattr(g, '_visibe_session', None)
            if sc is None or getattr(g, '_visibe_done', True):
                return
            g._visibe_done = True
            _active_session.reset(g._visibe_token)

            duration_ms = int((time.monotonic() - g._visibe_start) * 1000)
            failed      = exc is not None or status_code >= 400
            status      = 'failed' if failed else 'completed'

            visibe._batcher.flush()
            visibe.api_client.complete_trace(sc.trace_id, {
                'status':              status,
                'ended_at':            datetime.now(timezone.utc).isoformat(),
                'duration_ms':         duration_ms,
                'llm_call_count':      sc.llm_call_count,
                'total_cost':          round(sc.total_cost, 6),
                'total_tokens':        sc.total_input + sc.total_output,
                'total_input_tokens':  sc.total_input,
                'total_output_tokens': sc.total_output,
            })

    def starlette_middleware(self):
        """Return a Starlette/FastAPI-compatible middleware class.

        Groups all LLM calls within each HTTP request into one trace.
        Uses ContextVar so concurrent async requests are fully isolated.

        The sync HTTP calls (create_trace, complete_trace) are dispatched
        to a thread-pool executor so they don't block the async event loop.

        Usage:
            app.add_middleware(visibe.starlette_middleware())
        """
        try:
            from starlette.middleware.base import BaseHTTPMiddleware
        except ImportError:
            raise ImportError(
                "starlette_middleware() requires Starlette. "
                "Install it with: pip install starlette (or fastapi)"
            )

        visibe = self

        class _VisibMiddleware(BaseHTTPMiddleware):
            async def dispatch(self, request, call_next):
                if not visibe.api_client._enabled:
                    return await call_next(request)

                import asyncio

                trace_id    = str(uuid4())
                session_ctx = _SessionContext(trace_id=trace_id)
                name        = f"{request.method} {request.url.path}"
                start       = time.monotonic()

                # Fire create_trace in thread pool — don't block the event loop
                # and don't add a network round-trip to every request's hot path.
                loop = asyncio.get_running_loop()
                loop.run_in_executor(None, lambda: visibe.api_client.create_trace({
                    'trace_id':   trace_id,
                    'name':       name,
                    'framework':  'http',
                    'started_at': datetime.now(timezone.utc).isoformat(),
                    **(({'session_id': visibe.session_id}) if visibe.session_id else {}),
                }))

                token       = _active_session.set(session_ctx)
                status      = 'completed'
                status_code = 200

                try:
                    response    = await call_next(request)
                    status_code = response.status_code
                    return response
                except Exception:
                    status = 'failed'
                    raise
                finally:
                    _active_session.reset(token)

                    if status_code >= 400:
                        status = 'failed'

                    duration_ms = int((time.monotonic() - start) * 1000)
                    sc          = session_ctx

                    # Fire-and-forget in thread pool — don't block the response.
                    def _finalize():
                        visibe._batcher.flush()
                        visibe.api_client.complete_trace(sc.trace_id, {
                            'status':              status,
                            'ended_at':            datetime.now(timezone.utc).isoformat(),
                            'duration_ms':         duration_ms,
                            'llm_call_count':      sc.llm_call_count,
                            'total_cost':          round(sc.total_cost, 6),
                            'total_tokens':        sc.total_input + sc.total_output,
                            'total_input_tokens':  sc.total_input,
                            'total_output_tokens': sc.total_output,
                        })

                    loop.run_in_executor(None, _finalize)

        return _VisibMiddleware

    def track(self, client=None, name: str = "trace"):
        """Context manager to group calls into one trace.

        Supports OpenAI clients, CrewAI crews, LangChain/LangGraph runnables,
        and AutoGen model clients.

        Args:
            client: An OpenAI client, CrewAI Crew, or LangChain/LangGraph runnable
            name: Name for the grouped trace

        Example:
            >>> with obs.track(openai_client, name="my-chat"):
            ...     r1 = client.chat.completions.create(...)
            >>> with obs.track(crew, name="my-workflow"):
            ...     crew.kickoff()
        """
        if client is None:
            raise TypeError("track() requires a client argument")

        if _is_openai_client(client):
            return self._get_openai_integration().track(client=client, name=name)
        elif _is_crewai_crew(client):
            return self._get_crewai_integration().track(client, name=name)
        elif _is_autogen_model_client(client):
            return self._get_autogen_integration().track(
                model_client=client, name=name
            )
        elif _is_bedrock_client(client):
            return self._get_bedrock_integration().track(client=client, name=name)
        elif _is_langchain_runnable(client):
            if _is_langgraph_runnable(client):
                return self._get_langgraph_integration().track(name=name)
            return self._get_langchain_integration().track(name=name)
        else:
            raise TypeError(
                f"Unsupported type for track(): {type(client).__name__}. "
                f"Expected an OpenAI client, Bedrock client, "
                f"LangChain/LangGraph runnable, CrewAI Crew, or AutoGen model client."
            )
