"""Phase 3: Fleet organization — pool groups, scaler, federation (AD-517).

Registers pool groups, starts the pool scaler, and sets up federation
transport/bridge if enabled.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, Callable

from probos.startup.results import FleetOrganizationResult
from probos.substrate.pool_group import PoolGroup

if TYPE_CHECKING:
    from probos.cognitive.llm_client import BaseLLMClient
    from probos.config import SystemConfig
    from probos.consensus.escalation import EscalationManager
    from probos.consensus.trust import TrustNetwork
    from probos.mesh.intent import IntentBus
    from probos.substrate.pool import ResourcePool
    from probos.substrate.pool_group import PoolGroupRegistry

logger = logging.getLogger(__name__)


async def organize_fleet(
    *,
    config: "SystemConfig",
    pools: dict[str, "ResourcePool"],
    pool_groups: "PoolGroupRegistry",
    escalation_manager: "EscalationManager",
    intent_bus: "IntentBus",
    trust_network: "TrustNetwork",
    llm_client: "BaseLLMClient",
    build_pool_intent_map_fn: Callable[[], dict[str, list[str]]],
    find_consensus_pools_fn: Callable[[], set[str]],
    build_self_model_fn: Callable[..., Any],
    validate_remote_result_fn: Callable[..., Any] | None,
    nats_bus: Any | None = None,
) -> FleetOrganizationResult:
    """Register pool groups, start scaler, set up federation."""
    logger.info("Startup [fleet_organization]: starting")

    # Register pool groups (crew teams) — AD-291
    pool_groups.register(PoolGroup(
        name="core",
        display_name="Core Systems",
        pool_names={"system", "filesystem", "filesystem_writers", "directory", "search", "shell", "http", "introspect", "medical_vitals", "red_team", "system_qa"},
        exclude_from_scaler=True,
        startup_phase=1,  # AD-447: infrastructure first
    ))

    if config.utility_agents.enabled:
        pool_groups.register(PoolGroup(
            name="utility",
            display_name="Utility Agents",
            pool_names={"web_search", "page_reader", "weather", "news", "translator", "summarizer", "calculator", "todo_manager", "note_taker", "scheduler"},
            startup_phase=4,
        ))

    if config.medical.enabled:
        pool_groups.register(PoolGroup(
            name="medical",
            display_name="Medical",
            pool_names={"medical_diagnostician", "medical_surgeon", "medical_pharmacist", "medical_pathologist"},
            exclude_from_scaler=True,
            startup_phase=2,
        ))

    if config.self_mod.enabled:
        sm_pools = {"skills"}
        pool_groups.register(PoolGroup(
            name="self_mod",
            display_name="Self-Modification",
            pool_names=sm_pools,
            exclude_from_scaler=True,
            startup_phase=3,
        ))

    # Security pool group (AD-398: cognitive security officer)
    pool_groups.register(PoolGroup(
        name="security",
        display_name="Security",
        pool_names={"security_officer"},
        exclude_from_scaler=True,
        startup_phase=2,
    ))

    # Engineering pool group (AD-302, AD-398: add engineering_officer)
    pool_groups.register(PoolGroup(
        name="engineering",
        display_name="Engineering",
        pool_names={"builder", "engineering_officer"},
        exclude_from_scaler=True,
        startup_phase=2,
    ))

    # Science pool group (AD-307)
    pool_groups.register(PoolGroup(
        name="science",
        display_name="Science",
        pool_names={"architect", "scout", "science_data_analyst", "science_systems_analyst", "science_research_specialist"},
        exclude_from_scaler=True,
        startup_phase=3,
    ))

    # Operations pool group (AD-398)
    pool_groups.register(PoolGroup(
        name="operations",
        display_name="Operations",
        pool_names={"operations_officer", "training_officer"},
        exclude_from_scaler=True,
        startup_phase=2,
    ))

    # Bridge pool group (BF-015: Counselor was ungrouped)
    pool_groups.register(PoolGroup(
        name="bridge",
        display_name="Bridge",
        pool_names={"counselor"},
        exclude_from_scaler=True,
        startup_phase=1,  # AD-447: bridge is infrastructure
    ))

    # Start pool scaler if scaling is enabled
    pool_scaler = None
    if config.scaling.enabled:
        from probos.substrate.scaler import PoolScaler

        pool_intent_map = build_pool_intent_map_fn()
        consensus_pools = find_consensus_pools_fn()
        pool_scaler = PoolScaler(
            pools=pools,
            intent_bus=intent_bus,
            pool_config=config.pools,
            scaling_config=config.scaling,
            pool_intent_map=pool_intent_map,
            excluded_pools=pool_groups.excluded_pools(),
            trust_network=trust_network,
            consensus_pools=consensus_pools,
            consensus_min_agents=config.consensus.min_votes,
        )
        await pool_scaler.start()

        # PATCH(AD-517): Wire surge function into escalation manager
        escalation_manager._surge_fn = pool_scaler.request_surge

    # Start federation if enabled (AD-637e: NATS-first, ZeroMQ fallback)
    federation_bridge = None
    federation_transport = None
    if config.federation.enabled:
        from probos.federation import FederationRouter, FederationBridge

        peer_node_ids = [p.node_id for p in config.federation.peers]
        transport = None

        # Try NATS transport first (AD-637e)
        if nats_bus is not None and nats_bus.connected:
            try:
                from probos.federation.nats_transport import NATSFederationTransport

                transport = NATSFederationTransport(
                    node_id=config.federation.node_id,
                    nats_bus=nats_bus,
                    peer_node_ids=peer_node_ids,
                )
                await transport.start()
                # AD-479f: TLS pass-through surface — NATSBus consumes config.tls
                # at start. v1 logs whether TLS is requested for observability;
                # actual TLS context is configured on NATSBus during AD-637a startup.
                if config.federation.tls.enabled:
                    logger.info(
                        "AD-479f: Federation TLS requested (NATS path); cert_file=%s ca_file=%s verify_peer=%s",
                        config.federation.tls.cert_file,
                        config.federation.tls.ca_file,
                        config.federation.tls.verify_peer,
                    )
                logger.info("AD-637e: Federation using NATS transport")
            except Exception as e:
                logger.warning("AD-637e: NATS federation transport failed, falling back to ZeroMQ: %s", e)
                transport = None

        # ZeroMQ fallback
        if transport is None:
            try:
                from probos.federation.transport import FederationTransport

                transport = FederationTransport(
                    node_id=config.federation.node_id,
                    bind_address=config.federation.bind_address,
                    peers=config.federation.peers,
                )
                await transport.start()
            except ImportError:
                logger.warning("pyzmq not available; federation transport disabled")
            except Exception as e:
                logger.warning("Federation transport failed to start: %s", e)

        if transport is not None:
            router = FederationRouter()
            validate_fn = (
                validate_remote_result_fn
                if config.federation.validate_remote_results
                else None
            )
            bridge = FederationBridge(
                node_id=config.federation.node_id,
                transport=transport,
                router=router,
                intent_bus=intent_bus,
                config=config.federation,
                self_model_fn=build_self_model_fn,
                validate_fn=validate_fn,
            )
            await bridge.start()
            # PATCH(AD-517): Wire federation function into intent bus
            intent_bus.set_federation_handler(bridge.forward_intent)
            federation_bridge = bridge
            federation_transport = transport
            logger.info("Federation started: node=%s", config.federation.node_id)

    logger.info("Startup [fleet_organization]: complete")
    return FleetOrganizationResult(
        pool_scaler=pool_scaler,
        federation_bridge=federation_bridge,
        federation_transport=federation_transport,
    )
