"""
Runtime utilities for Outpost9 agents.

Primary export: `capture()` - wraps any async generator to capture raw SDK messages.

Usage:
    from outpost9 import runtime

    async for msg in runtime.capture(query(prompt="...", options=options)):
        # Messages flow through unchanged
        # Raw SDK messages are captured automatically
"""

import os
import json
import dataclasses
import urllib.request
import urllib.error
from datetime import datetime, timezone
from typing import Any, AsyncIterator, Callable, Dict, Optional, TypeVar
from functools import wraps

T = TypeVar('T')


# Environment variables set by Outpost9
_AGENT_ID = os.environ.get("OUTPOST9_AGENT_ID")
_LOG_URL = os.environ.get("OUTPOST9_LOG_URL")
_INTERNAL_TOKEN = os.environ.get("OUTPOST9_INTERNAL_TOKEN")


def _serialize(obj: Any) -> Any:
    """Recursively serialize dataclasses and objects to JSON-compatible dicts."""
    if obj is None or isinstance(obj, (str, int, float, bool)):
        return obj
    if isinstance(obj, (list, tuple)):
        return [_serialize(x) for x in obj]
    if isinstance(obj, dict):
        return {k: _serialize(v) for k, v in obj.items()}
    if dataclasses.is_dataclass(obj):
        return {k: _serialize(v) for k, v in dataclasses.asdict(obj).items()}
    if hasattr(obj, '__dict__'):
        return {k: _serialize(v) for k, v in obj.__dict__.items() if not k.startswith('_')}
    return str(obj)


def _send_sdk_message(msg_type: str, msg_data: dict) -> None:
    """POST raw SDK message to the log collector."""
    if not _AGENT_ID or not _LOG_URL:
        return

    from ._internal.auto_logging import get_run_id

    run_id = get_run_id()
    if not run_id:
        return

    payload = {
        "agentId": _AGENT_ID,
        "runId": run_id,
        "logs": [{
            "type": "sdk",
            "event": msg_type,
            "metadata": msg_data,
            "ts": datetime.now(timezone.utc).isoformat().replace('+00:00', 'Z'),
        }],
    }

    try:
        data_bytes = json.dumps(payload).encode("utf-8")
        headers = {
            "Content-Type": "application/json",
            "User-Agent": "Outpost9-SDK/0.0.1",
        }
        if _INTERNAL_TOKEN:
            headers["X-Internal-Token"] = _INTERNAL_TOKEN

        req = urllib.request.Request(
            _LOG_URL,
            data=data_bytes,
            headers=headers,
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as response:
            response.read()
    except Exception:
        pass  # Don't block on log failures


async def capture(messages: AsyncIterator[T]) -> AsyncIterator[T]:
    """
    Wrap any async generator to capture raw SDK messages.

    Works with any SDK method that returns an async iterator:
    - claude_code_sdk.query()
    - ClaudeSDKClient.receive_response()
    - Any future SDK methods

    Usage:
        from outpost9 import runtime

        async for msg in runtime.capture(query(prompt="...", options=options)):
            # msg is unchanged, but also captured to Outpost9

    Args:
        messages: Any async iterator yielding SDK messages

    Yields:
        Messages unchanged from the original iterator
    """
    async for msg in messages:
        # Capture raw message
        msg_type = type(msg).__name__
        msg_data = _serialize(msg)
        _send_sdk_message(msg_type, msg_data)

        # Yield unchanged
        yield msg


def _send_log(log_type: str, data: Dict[str, Any], event: Optional[str] = None) -> None:
    """Send a structured log entry to the collector."""
    if not _AGENT_ID or not _LOG_URL:
        return

    # Import here to avoid circular imports
    from ._internal.auto_logging import get_run_id

    log_entry: Dict[str, Any] = {
        "type": log_type,
        "msg": json.dumps(data) if isinstance(data, dict) else str(data),
        "metadata": {"service": "agent"},
    }
    if event:
        log_entry["event"] = event

    payload: Dict[str, Any] = {
        "agentId": _AGENT_ID,
        "logs": [log_entry],
    }

    # Include run ID if available (per-request context)
    run_id = get_run_id()
    if run_id:
        payload["runId"] = run_id

    try:
        data_bytes = json.dumps(payload).encode("utf-8")
        headers = {
            "Content-Type": "application/json",
            "User-Agent": "Outpost9-SDK/0.0.1",
        }
        if _INTERNAL_TOKEN:
            headers["X-Internal-Token"] = _INTERNAL_TOKEN

        req = urllib.request.Request(
            _LOG_URL,
            data=data_bytes,
            headers=headers,
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as response:
            response.read()
    except (urllib.error.URLError, Exception):
        pass  # Silently ignore


def _truncate(value: Any, max_len: int = 1000) -> Any:
    """Truncate large values for logging."""
    if value is None:
        return None
    if isinstance(value, str) and len(value) > max_len:
        return value[:max_len] + f"... (truncated, {len(value)} total chars)"
    elif isinstance(value, dict):
        return {k: _truncate(v, max_len) for k, v in value.items()}
    elif isinstance(value, list):
        if len(value) > 10:
            return value[:10] + [f"... ({len(value)} total items)"]
        return [_truncate(v, max_len) for v in value]
    return value


def log_event(event: str, data: Optional[Dict[str, Any]] = None) -> None:
    """
    Manually log a custom event.

    Args:
        event: Event name (e.g., "step_complete", "checkpoint")
        data: Optional data to include with the event
    """
    _send_log("app", data or {}, event=event)


def log_error(message: str, error: Optional[Exception] = None) -> None:
    """
    Log an error event.

    Args:
        message: Error message
        error: Optional exception to include
    """
    data: Dict[str, Any] = {"message": message}
    if error:
        data["error_type"] = type(error).__name__
        data["error_details"] = str(error)

    _send_log("error", data, event="error")


def instrument(func: Callable) -> Callable:
    """
    Decorator to instrument a function with logging.

    Logs function entry, exit, and any exceptions.
    """
    @wraps(func)
    async def async_wrapper(*args, **kwargs):
        func_name = func.__name__
        _send_log("app", {"func": func_name}, event="func_start")
        try:
            result = await func(*args, **kwargs)
            _send_log("app", {"func": func_name, "success": True}, event="func_end")
            return result
        except Exception as e:
            _send_log(
                "error",
                {"func": func_name, "error": str(e), "error_type": type(e).__name__},
                event="func_error",
            )
            raise

    @wraps(func)
    def sync_wrapper(*args, **kwargs):
        func_name = func.__name__
        _send_log("app", {"func": func_name}, event="func_start")
        try:
            result = func(*args, **kwargs)
            _send_log("app", {"func": func_name, "success": True}, event="func_end")
            return result
        except Exception as e:
            _send_log(
                "error",
                {"func": func_name, "error": str(e), "error_type": type(e).__name__},
                event="func_error",
            )
            raise

    import asyncio
    if asyncio.iscoroutinefunction(func):
        return async_wrapper
    return sync_wrapper
