"""
Omium Core Integration - Initialization and Configuration

This module provides the central initialization and configuration for Omium SDK.
It handles API key validation, auto-detection of frameworks, and global state management.
"""

import os
import logging
import threading
import atexit
from typing import Optional, Dict, Any, List
from dataclasses import dataclass, field
from enum import Enum

from omium.config import DEFAULT_API_BASE_URL

logger = logging.getLogger("omium")


class CheckpointStrategy(str, Enum):
    """Checkpoint creation strategy."""
    NODE = "node"       # Create checkpoint at each graph node (LangGraph)
    TASK = "task"       # Create checkpoint at each task (CrewAI)
    AGENT = "agent"     # Create checkpoint at each agent execution
    MANUAL = "manual"   # Only create checkpoints when explicitly called


@dataclass
class OmiumConfig:
    """
    Omium SDK configuration.
    
    Attributes:
        api_key: Omium API key (required, starts with 'om_')
        project: Project name for organizing traces
        auto_trace: Enable automatic tracing of framework calls
        auto_checkpoint: Enable automatic checkpoint creation
        checkpoint_strategy: When to create checkpoints
        api_base_url: Omium API base URL
        checkpoint_manager_url: Checkpoint Manager gRPC endpoint
        tracing_url: Trace ingestion endpoint
        batch_size: Number of spans to batch before sending
        flush_interval: Seconds between automatic flushes
        max_retries: Max retry attempts for failed requests
        timeout: Request timeout in seconds
        debug: Enable debug logging
    """
    
    # Required
    api_key: Optional[str] = None
    
    # Project settings
    project: str = "default"
    
    # Auto-instrumentation settings
    auto_trace: bool = True
    auto_checkpoint: bool = True
    checkpoint_strategy: CheckpointStrategy = CheckpointStrategy.NODE
    
    # Backend URLs (base without /api/v1; paths appended where needed)
    api_base_url: str = f"{DEFAULT_API_BASE_URL}/api/v1"
    checkpoint_manager_url: str = "checkpoint-manager.omium.ai:7001"
    tracing_url: str = f"{DEFAULT_API_BASE_URL}/api/v1/traces"
    
    # Performance settings
    batch_size: int = 100
    flush_interval: float = 5.0
    max_retries: int = 3
    timeout: float = 30.0
    
    # Debug
    debug: bool = False
    
    # Runtime state (set after validation)
    tenant_id: Optional[str] = None
    user_id: Optional[str] = None
    execution_id: Optional[str] = None
    
    # Detected frameworks
    detected_frameworks: List[str] = field(default_factory=list)
    
    def validate(self) -> bool:
        """
        Validate configuration.
        
        Raises:
            ValueError: If configuration is invalid
        """
        if not self.api_key:
            raise ValueError(
                "API key is required. Set OMIUM_API_KEY environment variable "
                "or pass api_key to omium.init()"
            )
        
        # Accept both 'om_' and 'omium_' prefixes
        if not (self.api_key.startswith("om_") or self.api_key.startswith("omium_")):
            raise ValueError(
                f"Invalid API key format: '{self.api_key[:10]}...'. "
                "API key should start with 'om_' or 'omium_'"
            )
        
        if len(self.api_key) < 20:
            raise ValueError("API key is too short. Please check your key.")
        
        return True
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert config to dictionary (excluding sensitive data)."""
        return {
            "project": self.project,
            "auto_trace": self.auto_trace,
            "auto_checkpoint": self.auto_checkpoint,
            "checkpoint_strategy": self.checkpoint_strategy.value,
            "api_base_url": self.api_base_url,
            "detected_frameworks": self.detected_frameworks,
        }


# Global state
_config: Optional[OmiumConfig] = None
_initialized: bool = False
_lock = threading.Lock()
_shutdown_registered: bool = False


def init(
    api_key: Optional[str] = None,
    project: Optional[str] = None,
    auto_trace: bool = True,
    auto_checkpoint: bool = True,
    checkpoint_strategy: str = "node",
    api_base_url: Optional[str] = None,
    debug: bool = False,
    **kwargs
) -> OmiumConfig:
    """
    Initialize Omium SDK.
    
    This is the primary entry point for Omium. Call this at the start of your
    application to enable automatic tracing and checkpointing.
    
    Args:
        api_key: Omium API key. If not provided, reads from OMIUM_API_KEY env var.
        project: Project name for organizing traces. Defaults to "default".
        auto_trace: Enable automatic tracing of LangGraph/CrewAI calls.
        auto_checkpoint: Enable automatic checkpoint creation.
        checkpoint_strategy: When to create checkpoints - "node", "task", "agent", or "manual".
        api_base_url: Override the default API URL (for self-hosted deployments).
        debug: Enable debug logging.
        **kwargs: Additional configuration options.
    
    Returns:
        OmiumConfig: The active configuration object.
    
    Raises:
        ValueError: If API key is missing or invalid.
    
    Example:
        ```python
        import omium
        
        # Simple initialization
        omium.init(api_key="om_xxx")
        
        # With project name
        omium.init(api_key="om_xxx", project="my-research-crew")
        
        # Disable auto-checkpointing
        omium.init(api_key="om_xxx", auto_checkpoint=False)
        ```
    """
    global _config, _initialized
    
    with _lock:
        if _initialized:
            logger.warning(
                "Omium already initialized. Use omium.configure() to update settings."
            )
            return _config
        
        # Load from environment variables first, then override with parameters
        resolved_api_key = api_key or os.environ.get("OMIUM_API_KEY")
        resolved_project = project or os.environ.get("OMIUM_PROJECT", "default")
        
        # Handle auto_trace from env
        env_tracing = os.environ.get("OMIUM_TRACING", "").lower()
        if api_key is None and env_tracing:
            auto_trace = env_tracing in ("true", "1", "yes")
        
        # Handle auto_checkpoint from env
        env_checkpoints = os.environ.get("OMIUM_CHECKPOINTS", "").lower()
        if api_key is None and env_checkpoints:
            auto_checkpoint = env_checkpoints not in ("false", "0", "no")
        
        # Parse checkpoint strategy
        try:
            strategy = CheckpointStrategy(checkpoint_strategy)
        except ValueError:
            logger.warning(f"Invalid checkpoint_strategy '{checkpoint_strategy}', using 'node'")
            strategy = CheckpointStrategy.NODE
        
        # Normalize the API URL to ensure it includes /api/v1
        raw_url = api_base_url or os.environ.get("OMIUM_API_URL", DEFAULT_API_BASE_URL)
        # Strip trailing slash for consistent handling
        raw_url = raw_url.rstrip("/")
        # Append /api/v1 if missing (common user mistake)
        if not raw_url.endswith("/api/v1"):
            raw_url = raw_url + "/api/v1"
        resolved_api_base_url = raw_url
        
        # Create config
        _config = OmiumConfig(
            api_key=resolved_api_key,
            project=resolved_project,
            auto_trace=auto_trace,
            auto_checkpoint=auto_checkpoint,
            checkpoint_strategy=strategy,
            api_base_url=resolved_api_base_url,
            debug=debug or os.environ.get("OMIUM_DEBUG", "").lower() in ("true", "1"),
            **kwargs
        )
        
        # Set up debug logging if requested
        if _config.debug:
            logging.getLogger("omium").setLevel(logging.DEBUG)
            handler = logging.StreamHandler()
            handler.setFormatter(logging.Formatter(
                "%(asctime)s [OMIUM] %(levelname)s: %(message)s"
            ))
            logging.getLogger("omium").addHandler(handler)
        
        # Validate
        _config.validate()
        
        # Detect and instrument frameworks
        if _config.auto_trace:
            _detect_and_instrument_frameworks()
        
        # Register shutdown handler
        _register_shutdown()
        
        _initialized = True
        
        logger.info(
            f"Omium initialized: project={_config.project}, "
            f"auto_trace={_config.auto_trace}, "
            f"frameworks={_config.detected_frameworks}"
        )
        
        # Auto-register workflow for the project (non-blocking, background task)
        # Can be disabled via OMIUM_SKIP_WORKFLOW_REGISTER=true (e.g. when API returns 502)
        skip_register = os.environ.get("OMIUM_SKIP_WORKFLOW_REGISTER", "").lower() in ("true", "1", "yes")
        if _config.project and _config.project != "default" and not skip_register:
            try:
                import threading
                thread = threading.Thread(
                    target=_ensure_workflow_registered_sync,
                    args=(_config.project, _config.api_key, _config.api_base_url),
                    daemon=True
                )
                thread.start()
            except Exception as e:
                logger.debug(f"Failed to start workflow registration thread: {e}")
        
        return _config


def configure(**kwargs) -> OmiumConfig:
    """
    Update configuration after initialization.
    
    If Omium hasn't been initialized yet, this will call init() with the
    provided parameters.
    
    Args:
        **kwargs: Configuration options to update.
    
    Returns:
        OmiumConfig: The updated configuration.
    """
    global _config
    
    if not _initialized:
        return init(**kwargs)
    
    with _lock:
        for key, value in kwargs.items():
            if hasattr(_config, key):
                setattr(_config, key, value)
            else:
                logger.warning(f"Unknown configuration option: {key}")
        
        return _config


def get_current_config() -> Optional[OmiumConfig]:
    """
    Get the current configuration.
    
    Returns:
        OmiumConfig or None if not initialized.
    """
    return _config


def set_execution_id(execution_id: str) -> None:
    """
    Set the current execution ID for correlation with Execution Engine.
    
    When a workflow is started via Execution Engine, call this function
    with the execution_id returned by the Engine so that all traces,
    checkpoints, and spans are correlated under the same ID.
    
    Args:
        execution_id: The execution ID from Execution Engine
        
    Example:
        ```python
        import omium
        import httpx
        
        omium.init(api_key="om_xxx")
        
        # Start execution via Execution Engine
        response = httpx.post(f"{omium.get_current_config().api_base_url}/executions", ...)
        exec_data = response.json()
        execution_id = exec_data["id"]
        
        # Set execution ID so LangGraph traces use it
        omium.set_execution_id(execution_id)
        
        # Now run your LangGraph workflow - traces will use this execution_id
        graph.invoke({"input": "hello"})
        ```
    """
    global _config
    if _config:
        _config.execution_id = execution_id
        logger.debug(f"Set execution_id: {execution_id}")
    else:
        logger.warning("Omium not initialized. Call omium.init() first.")


def get_execution_id() -> Optional[str]:
    """Get the current execution ID if set."""
    return _config.execution_id if _config else None


def is_initialized() -> bool:
    """
    Check if Omium has been initialized.
    
    Returns:
        bool: True if init() has been called successfully.
    """
    return _initialized


def shutdown():
    """
    Shutdown Omium SDK gracefully.
    
    This flushes any pending traces and closes connections.
    Called automatically on program exit.
    """
    global _initialized, _config
    
    if not _initialized:
        return
    
    logger.debug("Shutting down Omium SDK...")
    
    # Flush pending traces
    try:
        from omium.integrations.tracer import flush_all_tracers
        flush_all_tracers()
    except Exception as e:
        logger.warning(f"Error flushing traces on shutdown: {e}")
    
    _initialized = False
    logger.debug("Omium SDK shutdown complete")


def _detect_and_instrument_frameworks():
    """Detect installed frameworks and apply instrumentation."""
    global _config
    
    detected = []
    
    # Try LangGraph
    try:
        from omium.integrations.langgraph import instrument_langgraph
        instrument_langgraph()
        detected.append("langgraph")
        logger.info("✓ LangGraph auto-instrumentation enabled")
    except ImportError:
        logger.debug("LangGraph not installed, skipping instrumentation")
    except Exception as e:
        logger.warning(f"Failed to instrument LangGraph: {e}")
    
    # Try CrewAI
    try:
        from omium.integrations.crewai import instrument_crewai
        instrument_crewai()
        detected.append("crewai")
        logger.info("✓ CrewAI auto-instrumentation enabled")
    except ImportError:
        logger.debug("CrewAI not installed, skipping instrumentation")
    except Exception as e:
        logger.warning(f"Failed to instrument CrewAI: {e}")
    
    # Try LangChain (for general callback support)
    try:
        import langchain_core
        detected.append("langchain")
        logger.debug("LangChain detected (use OmiumCallbackHandler for tracing)")
    except ImportError:
        pass
    
    if _config:
        _config.detected_frameworks = detected
    
    if not detected:
        logger.info(
            "No supported frameworks detected. "
            "Use @omium.trace decorator for manual tracing."
        )


def _register_shutdown():
    """Register shutdown handler."""
    global _shutdown_registered
    
    if _shutdown_registered:
        return
    
    atexit.register(shutdown)
    _shutdown_registered = True


def _ensure_workflow_registered_sync(project: str, api_key: str, api_base_url: str):
    """
    Synchronously ensure workflow is registered for the project.
    
    This is called in a background thread to avoid blocking initialization.
    """
    import asyncio
    loop = None
    try:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        loop.run_until_complete(_ensure_workflow_registered(project, api_key, api_base_url))
    except Exception as e:
        logger.info(
            f"Workflow auto-registration failed for '{project}': {e}. "
            "Your app will still run. Set OMIUM_SKIP_WORKFLOW_REGISTER=true to skip."
        )
        logger.debug("Workflow registration traceback:", exc_info=True)
    finally:
        if loop is not None:
            try:
                loop.close()
            except Exception:
                pass


async def _ensure_workflow_registered(project: str, api_key: str, api_base_url: str):
    """
    Ensure workflow is registered for the project.
    
    If a workflow with the project name doesn't exist, create it via project sync endpoint.
    This makes the workflow appear in the Automations page automatically.
    """
    if not api_key or not project:
        return
    
    try:
        import httpx
        
        # Use the base URL without /api/v1 since project sync endpoint is at /api/v1/projects
        # Fix: Use replace instead of rstrip to avoid stripping partial matches
        if api_base_url.endswith('/api/v1'):
            base_url = api_base_url[:-7]  # Remove '/api/v1'
        else:
            base_url = api_base_url
        base_url = base_url.rstrip('/')
        if not base_url.startswith('http'):
            base_url = f"https://{base_url}"
        
        project_sync_url = f"{base_url}/api/v1/projects"
        
        # Check if workflow already exists by trying to get projects
        async with httpx.AsyncClient(timeout=10.0) as client:
            # First, try to get existing projects to see if this one exists
            try:
                response = await client.get(
                    f"{base_url}/api/v1/projects",
                    headers={"X-API-Key": api_key},
                    timeout=5.0
                )
                if response.status_code == 200:
                    projects = response.json()
                    # Check if project exists (projects might be a list or dict)
                    if isinstance(projects, list):
                        existing = any(p.get("name") == project for p in projects)
                    elif isinstance(projects, dict):
                        items = projects.get("items", [])
                        existing = any(p.get("name") == project for p in items)
                    else:
                        existing = False
                    
                    if existing:
                        logger.debug(f"Workflow '{project}' already registered")
                        return
            except Exception as e:
                logger.debug(f"Could not check existing projects: {e}")
            
            # Project doesn't exist, create it via project sync
            # Create a workflow entry with the project name
            project_data = {
                "name": project,
                "version": "1.0.0",
                "config": {},
                "workflows": [
                    {
                        "name": project,
                        "type": "langgraph",  # Default to langgraph, can be overridden
                        "config": {}
                    }
                ]
            }
            
            response = await client.post(
                project_sync_url,
                headers={
                    "X-API-Key": api_key,
                    "Content-Type": "application/json"
                },
                json=project_data,
                timeout=10.0
            )
            
            if response.status_code in (200, 201):
                logger.info(f"Auto-registered workflow '{project}' in Omium dashboard")
            elif response.status_code == 409:
                # Conflict means workflow already exists - that's fine
                logger.debug(f"Workflow '{project}' already exists in Omium")
            else:
                # 502 often means gateway can't reach API (e.g. use Cloudflare tunnel URL)
                if response.status_code == 502:
                    logger.info(
                        f"Workflow auto-registration skipped (502 Bad Gateway). "
                        f"The API at {base_url} may be unreachable. "
                        "If you use a Cloudflare tunnel, set OMIUM_API_URL to your tunnel base URL in .env. "
                        "To hide this, set OMIUM_SKIP_WORKFLOW_REGISTER=true. Your app will still run."
                    )
                else:
                    logger.info(
                        f"Could not auto-register workflow '{project}' in Omium dashboard: HTTP {response.status_code}. "
                        "Your app will still run; set OMIUM_SKIP_WORKFLOW_REGISTER=true to skip registration."
                    )
                logger.debug(f"Response: {response.text[:500]}")
                
    except Exception as e:
        # Log at info so 502/gateway issues don't look like hard failures
        logger.info(
            f"Workflow auto-registration failed for '{project}': {e}. "
            "Your app will still run. Set OMIUM_SKIP_WORKFLOW_REGISTER=true to skip."
        )
        logger.debug(f"Workflow registration traceback:", exc_info=True)


# Auto-initialization support
def _try_auto_init():
    """
    Attempt auto-initialization from environment variables.
    
    This is called lazily when frameworks are first used if
    OMIUM_TRACING=true is set.
    """
    if _initialized:
        return
    
    if os.environ.get("OMIUM_TRACING", "").lower() in ("true", "1", "yes"):
        if os.environ.get("OMIUM_API_KEY"):
            try:
                init()
            except Exception as e:
                logger.warning(f"Auto-initialization failed: {e}")


# ---------------------------------------------------------------------------
# Scoring API
# ---------------------------------------------------------------------------

async def score(
    trace_id: str,
    name: str,
    value: float,
    comment: Optional[str] = None,
    span_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Score a trace or span for evaluation.
    
    Use this to track quality metrics for your agent outputs.
    
    Args:
        trace_id: The trace ID to score
        name: Score name (e.g., 'accuracy', 'relevance', 'quality')
        value: Score value between 0.0 and 1.0
        comment: Optional comment explaining the score
        span_id: Optional span ID for granular scoring
    
    Returns:
        Dict with score details
    
    Example:
        await omium.score(
            trace_id="tr_xxx",
            name="accuracy",
            value=0.85,
            comment="Good but missed one key detail"
        )
    """
    config = get_current_config()
    if not config:
        raise RuntimeError("Omium not initialized. Call omium.init() first.")
    
    if not 0.0 <= value <= 1.0:
        raise ValueError("Score value must be between 0.0 and 1.0")
    
    import httpx
    
    try:
        async with httpx.AsyncClient(timeout=config.timeout) as client:
            response = await client.post(
                f"{config.api_base_url}/scores",
                headers={"X-API-Key": config.api_key},
                json={
                    "trace_id": trace_id,
                    "span_id": span_id,
                    "name": name,
                    "value": value,
                    "comment": comment,
                }
            )
            
            if response.status_code >= 400:
                logger.warning(f"Failed to create score: {response.status_code} - {response.text}")
                return {"error": response.text}
            
            return response.json()
            
    except Exception as e:
        logger.error(f"Error creating score: {e}")
        return {"error": str(e)}


def score_sync(
    trace_id: str,
    name: str,
    value: float,
    comment: Optional[str] = None,
    span_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Synchronous version of score().
    
    See score() for full documentation.
    """
    import asyncio
    return asyncio.run(score(trace_id, name, value, comment, span_id))

