"""
Task Planning Workflow Factory - Creates and configures planning workflows

This module provides factory functions to create Agno workflows for task planning:
- create_planning_workflow(): Standard 2-step workflow
- create_fast_workflow(): Single-step workflow for --local mode

The workflow orchestrates agents through defined steps with:
- Pre-fetched data caching
- Model configuration (Sonnet for reliability)
- Validation hooks
"""
import json
from typing import Optional, Dict, Any
import os
import structlog

from sqlalchemy.orm import Session
from agno.workflow import Workflow
from agno.models.litellm import LiteLLM

from .cache import get_cached_prefetch, set_cached_prefetch
from .agents import (
    create_analysis_and_selection_agent,
    create_plan_generation_agent,
    create_fast_selection_agent,
)
from .hooks import validate_prefetch_context

logger = structlog.get_logger()


# ============================================================================
# Model Configuration
# ============================================================================

def get_litellm_config() -> tuple[str, str]:
    """
    Get LiteLLM API configuration from environment.

    Returns:
        Tuple of (api_url, api_key)

    Raises:
        ValueError: If LITELLM_API_KEY is not set
    """
    api_url = (
        os.getenv("LITELLM_API_URL") or
        os.getenv("LITELLM_API_BASE") or
        "https://llm-proxy.kubiya.ai"
    ).strip()

    api_key = os.getenv("LITELLM_API_KEY", "").strip()
    if not api_key:
        raise ValueError("LITELLM_API_KEY environment variable not set")

    return api_url, api_key


def create_model(
    model_id: str,
    api_url: str,
    api_key: str,
    timeout: int = 60,
    task_execution_id: str = "",
) -> LiteLLM:
    """
    Create a LiteLLM model instance.

    Args:
        model_id: Model identifier (e.g., "kubiya/claude-sonnet-4")
        api_url: LiteLLM API base URL
        api_key: LiteLLM API key
        timeout: Request timeout in seconds
        task_execution_id: Task execution ID for Langfuse tracing (from DB executions table id column)

    Returns:
        Configured LiteLLM instance
    """
    # Build metadata for Langfuse tracking
    metadata = {}
    if task_execution_id:
        metadata["trace_metadata"] = json.dumps({"task_execution_id": task_execution_id})

    return LiteLLM(
        id=f"openai/{model_id}",
        api_base=api_url,
        api_key=api_key,
        request_params={"timeout": timeout},
        metadata=metadata,
    )


# ============================================================================
# Pre-fetch Resources
# ============================================================================

def prefetch_resources(
    db: Session,
    organization_id: str,
    api_token: str,
    outer_context: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """
    Pre-fetch organization resources with caching.

    Fetches top 20 of each resource type (agents, teams, environments, queues)
    and caches for 5 minutes to avoid repeated DB calls.

    Args:
        db: Database session
        organization_id: Organization ID
        api_token: API token for authentication
        outer_context: Optional pre-provided context (skips fetch)

    Returns:
        Dict with pre-fetched resources
    """
    # Check cache first
    cached = get_cached_prefetch(organization_id)
    if cached and not outer_context:
        logger.info("using_cached_prefetch", organization_id=organization_id[:8])
        return cached.copy()

    # Validate any provided outer_context
    if outer_context:
        outer_context = validate_prefetch_context(outer_context)
        # If outer_context has data, use it (don't override with DB fetch)
        if outer_context.get("agents") or outer_context.get("teams"):
            logger.info(
                "using_provided_context",
                agents=len(outer_context.get("agents", [])),
                teams=len(outer_context.get("teams", []))
            )
            return outer_context

    # Fetch fresh data from database
    logger.info("pre_fetching_top_resources", organization_id=organization_id[:8])

    from control_plane_api.app.lib.planning_tools.planning_service import PlanningService

    planning_service = PlanningService(db, organization_id, api_token)
    result = {}

    try:
        # Execute fetches SEQUENTIALLY to avoid SQLAlchemy session issues
        agents_data = planning_service.list_agents(limit=20, status=None)
        teams_data = planning_service.list_teams(limit=20, status=None)
        # Don't filter by status - environments may have status="ready" or "active"
        envs_data = planning_service.list_environments(status=None, limit=20)
        queues_data = planning_service.list_worker_queues(limit=20)

        result["agents"] = agents_data[:20]
        result["teams"] = teams_data[:20]
        result["environments"] = envs_data[:20]

        # Sort worker queues by active_workers DESC
        sorted_queues = sorted(
            queues_data[:20],
            key=lambda q: q.get("active_workers", 0),
            reverse=True
        )
        result["worker_queues"] = sorted_queues

        # Log worker queue availability for debugging
        queues_with_workers = [q for q in sorted_queues if q.get("active_workers", 0) > 0]
        if sorted_queues:
            logger.info(
                "worker_queues_fetched",
                total_queues=len(sorted_queues),
                queues_with_active_workers=len(queues_with_workers),
                top_queue_name=sorted_queues[0].get("name") if sorted_queues else None,
                top_queue_active_workers=sorted_queues[0].get("active_workers", 0) if sorted_queues else 0
            )
        else:
            logger.warning(
                "no_worker_queues_available",
                message="No worker queues with active workers found - plan execution will fail"
            )

        result["pre_fetch_note"] = "Top 20 resources. Use tools for more specific searches."

        logger.info(
            "pre_fetch_completed",
            agents_count=len(result["agents"]),
            teams_count=len(result["teams"]),
            envs_count=len(result["environments"]),
            queues_count=len(result["worker_queues"])
        )

        # Cache for subsequent requests
        set_cached_prefetch(organization_id, result)

    except Exception as e:
        import traceback
        logger.warning(
            "pre_fetch_failed",
            error=str(e),
            traceback=traceback.format_exc()
        )
        # Return empty context - agent will use tools instead
        result = {}

    return result


# ============================================================================
# Main Workflow Factory
# ============================================================================

def create_planning_workflow(
    db: Session,
    organization_id: str,
    api_token: str,
    quick_mode: bool = False,
    outer_context: Optional[Dict[str, Any]] = None
) -> Workflow:
    """
    Create the task planning workflow.

    Creates a 2-step workflow:
    - Step 1: Analysis & Resource Selection (tool calling + structured output)
    - Step 2: Plan Generation (structured output only)

    Args:
        db: Database session
        organization_id: Organization ID (required)
        api_token: API token (required)
        quick_mode: Deprecated, kept for compatibility
        outer_context: Optional pre-fetched context

    Returns:
        Configured Workflow instance

    Raises:
        ValueError: If required parameters are missing
    """
    if not organization_id:
        raise ValueError("organization_id is required")
    if not api_token:
        raise ValueError("api_token is required")

    # Get LiteLLM configuration
    api_url, api_key = get_litellm_config()

    # Get model IDs from environment
    step1_model_id = os.getenv("STEP1_MODEL", "kubiya/claude-sonnet-4").strip()
    step2_model_id = os.getenv("STEP2_MODEL", "kubiya/claude-sonnet-4").strip()

    logger.info(
        "model_configuration",
        step1_model=step1_model_id,
        step2_model=step2_model_id,
        message="Using Sonnet for both steps (reliable structured output)"
    )

    # Create model instances
    step1_model = create_model(step1_model_id, api_url, api_key, timeout=60)
    step2_model = create_model(step2_model_id, api_url, api_key, timeout=60)

    # Pre-fetch resources with caching
    outer_context = prefetch_resources(db, organization_id, api_token, outer_context)

    # Create planning toolkit for Step 1 tools
    from control_plane_api.app.lib.planning_tools.agno_toolkit import PlanningToolkit
    planning_toolkit = PlanningToolkit(db, organization_id, api_token)

    # Create agents
    step1_agent = create_analysis_and_selection_agent(
        model=step1_model,
        planning_toolkit=planning_toolkit,
        outer_context=outer_context
    )

    step2_agent = create_plan_generation_agent(step2_model)

    # Create 2-step workflow
    workflow = Workflow(
        name="Task Planning Workflow",
        steps=[step1_agent, step2_agent],
        description="2-step task planning: (1) Analysis & Selection, (2) Plan Generation"
    )

    # Store references for validation and runtime access
    workflow._planning_toolkit = planning_toolkit
    workflow._outer_context = outer_context
    workflow._organization_id = organization_id

    logger.info(
        "planning_workflow_created",
        steps=2,
        pre_fetched_data=bool(outer_context)
    )

    return workflow


# ============================================================================
# Fast Workflow Factory (--local mode)
# ============================================================================

def create_fast_planning_workflow(
    db: Session,
    organization_id: str,
    api_token: str,
    outer_context: Dict[str, Any]
) -> Workflow:
    """
    Create a fast single-step workflow for --local mode.

    Uses pre-fetched data only (no API calls) for maximum speed.

    Args:
        db: Database session
        organization_id: Organization ID
        api_token: API token
        outer_context: Pre-fetched context (required for fast mode)

    Returns:
        Configured single-step Workflow

    Raises:
        ValueError: If outer_context is empty
    """
    if not outer_context or (not outer_context.get("agents") and not outer_context.get("teams")):
        raise ValueError("outer_context with agents or teams is required for fast mode")

    # Get LiteLLM configuration
    api_url, api_key = get_litellm_config()

    # Use faster model for single-step
    model_id = os.getenv("FAST_MODEL", "kubiya/claude-sonnet-4").strip()
    model = create_model(model_id, api_url, api_key, timeout=30)

    # Validate context
    outer_context = validate_prefetch_context(outer_context)

    # Create fast selection agent
    fast_agent = create_fast_selection_agent(model, outer_context)

    # Create single-step workflow
    workflow = Workflow(
        name="Fast Selection Workflow",
        steps=[fast_agent],
        description="Single-step fast selection for --local mode"
    )

    workflow._outer_context = outer_context
    workflow._organization_id = organization_id

    logger.info(
        "fast_workflow_created",
        agents_count=len(outer_context.get("agents", [])),
        teams_count=len(outer_context.get("teams", []))
    )

    return workflow


# ============================================================================
# Workflow Configuration Helpers
# ============================================================================

class WorkflowConfig:
    """Configuration for task planning workflows."""

    def __init__(
        self,
        step1_model: str = "kubiya/claude-sonnet-4",
        step2_model: str = "kubiya/claude-sonnet-4",
        step1_timeout: int = 60,
        step2_timeout: int = 60,
        cache_ttl: int = 300,
        prefer_runtime: str = "claude_code"
    ):
        self.step1_model = step1_model
        self.step2_model = step2_model
        self.step1_timeout = step1_timeout
        self.step2_timeout = step2_timeout
        self.cache_ttl = cache_ttl
        self.prefer_runtime = prefer_runtime

    @classmethod
    def from_env(cls) -> 'WorkflowConfig':
        """Create config from environment variables."""
        return cls(
            step1_model=os.getenv("STEP1_MODEL", "kubiya/claude-sonnet-4"),
            step2_model=os.getenv("STEP2_MODEL", "kubiya/claude-sonnet-4"),
            step1_timeout=int(os.getenv("STEP1_TIMEOUT", "60")),
            step2_timeout=int(os.getenv("STEP2_TIMEOUT", "60")),
            cache_ttl=int(os.getenv("PREFETCH_CACHE_TTL", "300")),
            prefer_runtime=os.getenv("PREFER_RUNTIME", "claude_code")
        )


def create_workflow_with_config(
    db: Session,
    organization_id: str,
    api_token: str,
    config: WorkflowConfig,
    outer_context: Optional[Dict[str, Any]] = None
) -> Workflow:
    """
    Create workflow with explicit configuration.

    Args:
        db: Database session
        organization_id: Organization ID
        api_token: API token
        config: Workflow configuration
        outer_context: Optional pre-fetched context

    Returns:
        Configured Workflow
    """
    api_url, api_key = get_litellm_config()

    step1_model = create_model(config.step1_model, api_url, api_key, config.step1_timeout)
    step2_model = create_model(config.step2_model, api_url, api_key, config.step2_timeout)

    # Pre-fetch with config
    outer_context = prefetch_resources(db, organization_id, api_token, outer_context)

    # Add preferred runtime to context
    if outer_context and not outer_context.get("preferred_runtime"):
        outer_context["preferred_runtime"] = config.prefer_runtime

    from control_plane_api.app.lib.planning_tools.agno_toolkit import PlanningToolkit
    planning_toolkit = PlanningToolkit(db, organization_id, api_token)

    step1_agent = create_analysis_and_selection_agent(
        model=step1_model,
        planning_toolkit=planning_toolkit,
        outer_context=outer_context
    )
    step2_agent = create_plan_generation_agent(step2_model)

    workflow = Workflow(
        name="Task Planning Workflow",
        steps=[step1_agent, step2_agent],
        description="Configured 2-step task planning workflow"
    )

    workflow._planning_toolkit = planning_toolkit
    workflow._outer_context = outer_context
    workflow._config = config

    return workflow
