"""Provider aggregate root - the main domain entity."""

import threading
import time
from typing import Any, Dict, List, Optional

from ...logging_config import get_logger
from ..contracts.metrics_publisher import IMetricsPublisher, NullMetricsPublisher
from ..events import (
    HealthCheckFailed,
    HealthCheckPassed,
    ProviderDegraded,
    ProviderIdleDetected,
    ProviderStarted,
    ProviderStateChanged,
    ProviderStopped,
    ToolInvocationCompleted,
    ToolInvocationFailed,
    ToolInvocationRequested,
)
from ..exceptions import (
    CannotStartProviderError,
    InvalidStateTransitionError,
    ProviderStartError,
    ToolInvocationError,
    ToolNotFoundError,
)
from ..value_objects import (
    CorrelationId,
    HealthCheckInterval,
    IdleTTL,
    ProviderId,
    ProviderMode,
    ProviderState,
)
from .aggregate import AggregateRoot
from .health_tracker import HealthTracker
from .tool_catalog import ToolCatalog, ToolSchema

logger = get_logger(__name__)


# Valid state transitions
VALID_TRANSITIONS = {
    ProviderState.COLD: {ProviderState.INITIALIZING},
    ProviderState.INITIALIZING: {
        ProviderState.READY,
        ProviderState.DEAD,
        ProviderState.DEGRADED,
    },
    ProviderState.READY: {
        ProviderState.COLD,
        ProviderState.DEAD,
        ProviderState.DEGRADED,
    },
    ProviderState.DEGRADED: {ProviderState.INITIALIZING, ProviderState.COLD},
    ProviderState.DEAD: {ProviderState.INITIALIZING, ProviderState.DEGRADED},
}


class Provider(AggregateRoot):
    """
    Provider aggregate root.

    Manages the complete lifecycle of an MCP provider including:
    - State machine with valid transitions
    - Health tracking and circuit breaker logic
    - Tool catalog management
    - Process/client management

    All public operations are thread-safe using internal locking.
    """

    def __init__(
        self,
        provider_id: str,
        mode: str | ProviderMode,  # Accept both string and enum
        command: Optional[List[str]] = None,
        image: Optional[str] = None,
        endpoint: Optional[str] = None,
        env: Optional[Dict[str, str]] = None,
        idle_ttl_s: int | IdleTTL = 300,  # Accept both int and value object
        health_check_interval_s: int | HealthCheckInterval = 60,  # Accept both int and value object
        max_consecutive_failures: int = 3,
        # Container-specific options
        volumes: Optional[List[str]] = None,
        build: Optional[Dict[str, str]] = None,
        resources: Optional[Dict[str, str]] = None,
        network: str = "none",
        read_only: bool = True,
        user: Optional[str] = None,  # UID:GID or username
        description: Optional[str] = None,  # Description/preprompt for AI models
        # Pre-defined tools (allows visibility before provider starts)
        tools: Optional[List[Dict[str, Any]]] = None,
        # Dependencies
        metrics_publisher: Optional[IMetricsPublisher] = None,
    ):
        super().__init__()

        # Identity
        self._id = ProviderId(provider_id)

        # Mode - normalize to ProviderMode enum (container -> docker)
        self._mode = ProviderMode.normalize(mode)

        self._description = description

        # Configuration - normalize to value objects
        self._command = command
        self._image = image
        self._endpoint = endpoint
        self._env = env or {}

        # Idle TTL - normalize to value object
        if isinstance(idle_ttl_s, IdleTTL):
            self._idle_ttl = idle_ttl_s
        else:
            self._idle_ttl = IdleTTL(idle_ttl_s)

        # Health check interval - normalize to value object
        if isinstance(health_check_interval_s, HealthCheckInterval):
            self._health_check_interval = health_check_interval_s
        else:
            self._health_check_interval = HealthCheckInterval(health_check_interval_s)

        # Container-specific configuration
        self._volumes = volumes or []
        self._build = build  # {"dockerfile": "...", "context": "..."}
        self._resources = resources or {"memory": "512m", "cpu": "1.0"}
        self._network = network
        self._read_only = read_only
        self._user = user

        # Dependencies (Dependency Inversion Principle)
        self._metrics_publisher = metrics_publisher or NullMetricsPublisher()

        # State
        self._state = ProviderState.COLD
        self._health = HealthTracker(max_consecutive_failures=max_consecutive_failures)
        self._tools = ToolCatalog()
        self._client: Optional[Any] = None  # StdioClient
        self._meta: Dict[str, Any] = {}
        self._last_used: float = 0.0

        # Pre-load tools from configuration (allows visibility before start)
        self._tools_predefined = False
        if tools:
            self._tools.update_from_list(tools)
            self._tools_predefined = True

        # Thread safety
        self._lock = threading.RLock()

    # --- Properties ---

    @property
    def id(self) -> ProviderId:
        """Provider identifier."""
        return self._id

    @property
    def provider_id(self) -> str:
        """Provider identifier as string (for backward compatibility)."""
        return str(self._id)

    @property
    def mode(self) -> ProviderMode:
        """Provider mode enum."""
        return self._mode

    @property
    def mode_str(self) -> str:
        """Provider mode as string (for backward compatibility)."""
        return self._mode.value

    @property
    def description(self) -> Optional[str]:
        """Provider description for AI models."""
        return self._description

    @property
    def state(self) -> ProviderState:
        """Current provider state."""
        with self._lock:
            return self._state

    @property
    def health(self) -> HealthTracker:
        """Health tracker."""
        return self._health

    @property
    def tools(self) -> ToolCatalog:
        """Tool catalog."""
        return self._tools

    @property
    def has_tools(self) -> bool:
        """Check if provider has any tools registered (predefined or discovered)."""
        return self._tools.count() > 0

    @property
    def tools_predefined(self) -> bool:
        """Check if tools were predefined in configuration (no startup needed for visibility)."""
        return self._tools_predefined

    @property
    def is_alive(self) -> bool:
        """Check if provider client is alive."""
        with self._lock:
            return self._client is not None and self._client.is_alive()

    @property
    def last_used(self) -> float:
        """Timestamp of last tool invocation."""
        with self._lock:
            return self._last_used

    @property
    def idle_time(self) -> float:
        """Time since last use in seconds."""
        with self._lock:
            if self._last_used == 0:
                return 0.0
            return time.time() - self._last_used

    @property
    def is_idle(self) -> bool:
        """Check if provider has been idle longer than TTL."""
        with self._lock:
            if self._state != ProviderState.READY:
                return False
            if self._last_used == 0:
                return False
            return self.idle_time > self._idle_ttl.seconds

    @property
    def meta(self) -> Dict[str, Any]:
        """Provider metadata."""
        with self._lock:
            return dict(self._meta)

    @property
    def lock(self) -> threading.RLock:
        """Get the internal lock (for backward compatibility)."""
        return self._lock

    # --- State Management ---

    def _transition_to(self, new_state: ProviderState) -> None:
        """
        Transition to a new state (must hold lock).

        Validates the transition is valid according to state machine rules.
        Records a ProviderStateChanged event.
        """
        if new_state == self._state:
            return

        if new_state not in VALID_TRANSITIONS.get(self._state, set()):
            raise InvalidStateTransitionError(self.provider_id, str(self._state.value), str(new_state.value))

        old_state = self._state
        self._state = new_state
        self._increment_version()

        self._record_event(
            ProviderStateChanged(
                provider_id=self.provider_id,
                old_state=str(old_state.value),
                new_state=str(new_state.value),
            )
        )

    def _can_start(self) -> tuple:
        """
        Check if provider can be started (must hold lock).

        Returns: (can_start, reason, time_until_retry)
        """
        if self._state == ProviderState.READY:
            if self._client and self._client.is_alive():
                return True, "already_ready", 0

        if self._state == ProviderState.DEGRADED:
            if not self._health.can_retry():
                time_left = self._health.time_until_retry()
                return False, "backoff_not_elapsed", time_left

        return True, "", 0

    # --- Business Operations ---

    def ensure_ready(self) -> None:
        """
        Ensure provider is in READY state, starting if necessary.

        Thread-safe. Blocks until ready or raises exception.

        Raises:
            CannotStartProviderError: If backoff hasn't elapsed
            ProviderStartError: If provider fails to start
        """
        with self._lock:
            # Fast path - already ready
            if self._state == ProviderState.READY:
                if self._client and self._client.is_alive():
                    return
                # Client died
                logger.warning(f"provider_dead: {self.provider_id}")
                self._state = ProviderState.DEAD

            # Check if we can start
            can_start, reason, time_left = self._can_start()
            if not can_start:
                raise CannotStartProviderError(
                    self.provider_id,
                    f"backoff not elapsed, retry in {time_left:.1f}s",
                    time_left,
                )

            # Start if needed
            if self._state in (
                ProviderState.COLD,
                ProviderState.DEAD,
                ProviderState.DEGRADED,
            ):
                self._start()

    def _start(self) -> None:
        """
        Start provider process (must hold lock).

        Handles subprocess, docker, container modes.
        """
        start_time = time.time()
        self._transition_to(ProviderState.INITIALIZING)

        cold_start_time = self._begin_cold_start_tracking()

        try:
            client = self._create_client()
            self._perform_mcp_handshake(client)
            self._finalize_start(client, start_time)
            self._end_cold_start_tracking(cold_start_time, success=True)

        except ProviderStartError:
            self._end_cold_start_tracking(cold_start_time, success=False)
            self._handle_start_failure(None)
            raise
        except Exception as e:
            self._end_cold_start_tracking(cold_start_time, success=False)
            self._handle_start_failure(e)
            raise ProviderStartError(self.provider_id, str(e)) from e

    def _begin_cold_start_tracking(self) -> Optional[float]:
        """Begin tracking cold start metrics. Returns start timestamp."""
        try:
            self._metrics_publisher.begin_cold_start(self.provider_id)
            return time.time()
        except Exception:
            return None

    def _end_cold_start_tracking(self, start_time: Optional[float], success: bool) -> None:
        """End cold start tracking and record metrics."""
        if start_time is None:
            return
        try:
            if success:
                duration = time.time() - start_time
                self._metrics_publisher.record_cold_start(self.provider_id, duration, self._mode.value)
            self._metrics_publisher.end_cold_start(self.provider_id)
        except Exception:
            pass

    def _create_client(self) -> Any:
        """Create and return the appropriate client based on mode."""
        from ..services.provider_launcher import get_launcher

        launcher = get_launcher(self._mode.value)
        config = self._get_launch_config()
        return launcher.launch(**config)

    def _get_launch_config(self) -> Dict[str, Any]:
        """Get launch configuration for the current mode."""
        if self._mode == ProviderMode.SUBPROCESS:
            return {"command": self._command, "env": self._env}

        if self._mode == ProviderMode.DOCKER:
            return {
                "image": self._image,
                "volumes": self._volumes,
                "env": self._env,
                "memory_limit": self._resources.get("memory", "512m"),
                "cpu_limit": self._resources.get("cpu", "1.0"),
                "network": self._network,
                "read_only": self._read_only,
                "user": self._user,
            }

        if self._mode.value in ("container", "podman"):
            return {
                "image": self._get_container_image(),
                "volumes": self._volumes,
                "env": self._env,
                "memory_limit": self._resources.get("memory", "512m"),
                "cpu_limit": self._resources.get("cpu", "1.0"),
                "network": self._network,
                "read_only": self._read_only,
                "user": self._user,
            }

        raise ValueError(f"unsupported_mode: {self._mode.value}")

    def _get_container_image(self) -> str:
        """Get or build container image."""
        from ..services.image_builder import BuildConfig, get_image_builder

        if self._build and self._build.get("dockerfile"):
            runtime = "podman" if self._mode.value == "podman" else "auto"
            builder = get_image_builder(runtime=runtime)
            build_config = BuildConfig(
                dockerfile=self._build["dockerfile"],
                context=self._build.get("context", "."),
                tag=self._build.get("tag"),
            )
            image = builder.build_if_needed(build_config)
            logger.info(f"Built image for {self.provider_id}: {image}")
            return image

        if not self._image:
            raise ProviderStartError(
                self.provider_id,
                "Container mode requires 'image' or 'build.dockerfile'",
            )
        return self._image

    def _perform_mcp_handshake(self, client: Any) -> None:
        """Perform MCP initialize and tools/list handshake."""
        # Initialize
        init_resp = client.call(
            "initialize",
            {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {"name": "mcp-registry", "version": "1.0.0"},
            },
            timeout=10.0,
        )

        if "error" in init_resp:
            error_msg = init_resp["error"].get("message", "unknown")
            self._log_client_error(client, error_msg)
            raise ProviderStartError(self.provider_id, f"init_failed: {error_msg}")

        # Discover tools
        tools_resp = client.call("tools/list", {}, timeout=10.0)
        if "error" in tools_resp:
            error_msg = tools_resp["error"].get("message", "unknown")
            raise ProviderStartError(self.provider_id, f"tools_list_failed: {error_msg}")

        tool_list = tools_resp.get("result", {}).get("tools", [])
        self._tools.update_from_list(tool_list)

    def _log_client_error(self, client: Any, error_msg: str) -> None:
        """Log detailed error info for debugging (especially in CI)."""
        if error_msg != "reader_died":
            return

        proc = getattr(client, "process", None)
        if not proc:
            return

        # Try to capture stderr
        stderr = getattr(proc, "stderr", None)
        if stderr:
            try:
                err_bytes = stderr.read()
                if err_bytes:
                    err_text = (err_bytes if isinstance(err_bytes, str) else err_bytes.decode(errors="replace")).strip()
                    if err_text:
                        logger.error(f"provider_container_stderr: {err_text}")
            except Exception:
                pass

        # Log exit code
        try:
            rc = proc.poll()
            if rc is not None:
                logger.error(f"provider_process_exit_code: {rc}")
        except Exception:
            pass

    def _finalize_start(self, client: Any, start_time: float) -> None:
        """Finalize successful provider start."""
        self._client = client
        self._meta = {
            "init_result": {},
            "tools_count": self._tools.count(),
            "started_at": time.time(),
        }
        self._transition_to(ProviderState.READY)
        self._health.record_success()
        self._last_used = time.time()

        startup_duration_ms = (time.time() - start_time) * 1000
        self._record_event(
            ProviderStarted(
                provider_id=self.provider_id,
                mode=self._mode.value,
                tools_count=self._tools.count(),
                startup_duration_ms=startup_duration_ms,
            )
        )

        logger.info(f"provider_started: {self.provider_id}, mode={self._mode.value}, tools={self._tools.count()}")

    def _handle_start_failure(self, error: Optional[Exception]) -> None:
        """Handle start failure (must hold lock)."""
        # Clean up client if partially started
        if self._client:
            try:
                self._client.close()
            except Exception:
                pass
            self._client = None

        self._health.record_failure()

        error_str = str(error) if error else "unknown error"

        # Determine new state
        if self._health.should_degrade():
            # Use direct assignment to avoid transition validation issues
            self._state = ProviderState.DEGRADED
            self._increment_version()

            logger.warning(f"provider_degraded: {self.provider_id}, " f"failures={self._health.consecutive_failures}")

            self._record_event(
                ProviderDegraded(
                    provider_id=self.provider_id,
                    consecutive_failures=self._health.consecutive_failures,
                    total_failures=self._health.total_failures,
                    reason=error_str,
                )
            )
        else:
            self._state = ProviderState.DEAD
            self._increment_version()

        logger.error(f"provider_start_failed: {self.provider_id}, error={error_str}")

    def invoke_tool(self, tool_name: str, arguments: Dict[str, Any], timeout: float = 30.0) -> Dict[str, Any]:
        """
        Invoke a tool on this provider.

        Thread-safe. Ensures provider is ready before invocation.

        Args:
            tool_name: Name of the tool to invoke
            arguments: Tool arguments
            timeout: Timeout in seconds

        Returns:
            Tool result dictionary

        Raises:
            CannotStartProviderError: If provider cannot be started
            ToolNotFoundError: If tool doesn't exist
            ToolInvocationError: If invocation fails
        """
        correlation_id = str(CorrelationId())

        with self._lock:
            # Ensure ready
            self.ensure_ready()

            # Check tool exists
            if not self._tools.has(tool_name):
                # Try refreshing tools once
                self._refresh_tools()

            if not self._tools.has(tool_name):
                raise ToolNotFoundError(self.provider_id, tool_name)

            self._health._total_invocations += 1

            # Record start event
            self._record_event(
                ToolInvocationRequested(
                    provider_id=self.provider_id,
                    tool_name=tool_name,
                    correlation_id=correlation_id,
                    arguments=arguments,
                )
            )

            start_time = time.time()

            try:
                response = self._client.call(
                    "tools/call",
                    {"name": tool_name, "arguments": arguments},
                    timeout=timeout,
                )

                if "error" in response:
                    error_msg = response["error"].get("message", "unknown")
                    self._health.record_invocation_failure()

                    self._record_event(
                        ToolInvocationFailed(
                            provider_id=self.provider_id,
                            tool_name=tool_name,
                            correlation_id=correlation_id,
                            error_message=error_msg,
                            error_type=str(response["error"].get("code", "unknown")),
                        )
                    )

                    raise ToolInvocationError(
                        self.provider_id,
                        f"tool_error: {error_msg}",
                        {"tool_name": tool_name, "correlation_id": correlation_id},
                    )

                # Success
                duration_ms = (time.time() - start_time) * 1000
                self._health.record_success()
                self._last_used = time.time()

                result = response.get("result", {})
                self._record_event(
                    ToolInvocationCompleted(
                        provider_id=self.provider_id,
                        tool_name=tool_name,
                        correlation_id=correlation_id,
                        duration_ms=duration_ms,
                        result_size_bytes=len(str(result)),
                    )
                )

                logger.debug(f"tool_invoked: {correlation_id}, " f"provider={self.provider_id}, tool={tool_name}")

                return result

            except ToolInvocationError:
                raise
            except Exception as e:
                self._health.record_failure()

                self._record_event(
                    ToolInvocationFailed(
                        provider_id=self.provider_id,
                        tool_name=tool_name,
                        correlation_id=correlation_id,
                        error_message=str(e),
                        error_type=type(e).__name__,
                    )
                )

                logger.error(
                    f"tool_invocation_failed: {correlation_id}, "
                    f"provider={self.provider_id}, tool={tool_name}, error={e}"
                )

                raise ToolInvocationError(
                    self.provider_id,
                    str(e),
                    {"tool_name": tool_name, "correlation_id": correlation_id},
                ) from e

    def _refresh_tools(self) -> None:
        """Refresh tool catalog from provider (must hold lock)."""
        if not self._client or not self._client.is_alive():
            return

        try:
            tools_resp = self._client.call("tools/list", {}, timeout=5.0)
            if "result" in tools_resp:
                tool_list = tools_resp.get("result", {}).get("tools", [])
                self._tools.update_from_list(tool_list)
        except Exception as e:
            logger.warning(f"tool_refresh_failed: {self.provider_id}, error={e}")

    def health_check(self) -> bool:
        """
        Perform active health check.

        Thread-safe. Returns True if healthy.
        """
        with self._lock:
            if self._state != ProviderState.READY:
                return False

            if not self._client or not self._client.is_alive():
                self._state = ProviderState.DEAD
                self._increment_version()
                return False

            try:
                start_time = time.time()
                response = self._client.call("tools/list", {}, timeout=5.0)

                if "error" in response:
                    raise Exception(response["error"].get("message", "unknown"))

                duration_ms = (time.time() - start_time) * 1000
                self._health.record_success()

                self._record_event(HealthCheckPassed(provider_id=self.provider_id, duration_ms=duration_ms))

                return True

            except Exception as e:
                self._health.record_failure()

                self._record_event(
                    HealthCheckFailed(
                        provider_id=self.provider_id,
                        consecutive_failures=self._health.consecutive_failures,
                        error_message=str(e),
                    )
                )

                logger.warning(f"health_check_failed: {self.provider_id}, error={e}")

                if self._health.should_degrade():
                    self._state = ProviderState.DEGRADED
                    self._increment_version()

                    logger.warning(f"provider_degraded_by_health_check: {self.provider_id}")

                    self._record_event(
                        ProviderDegraded(
                            provider_id=self.provider_id,
                            consecutive_failures=self._health.consecutive_failures,
                            total_failures=self._health.total_failures,
                            reason="health_check_failures",
                        )
                    )

                return False

    def maybe_shutdown_idle(self) -> bool:
        """
        Shutdown if idle past TTL.

        Thread-safe. Returns True if shutdown was performed.
        """
        with self._lock:
            if self._state != ProviderState.READY:
                return False

            idle_time = time.time() - self._last_used
            if idle_time > self._idle_ttl.seconds:
                self._record_event(
                    ProviderIdleDetected(
                        provider_id=self.provider_id,
                        idle_duration_s=idle_time,
                        last_used_at=self._last_used,
                    )
                )

                logger.info(f"provider_idle_shutdown: {self.provider_id}, idle={idle_time:.1f}s")
                self._shutdown_internal(reason="idle")
                return True

            return False

    def shutdown(self) -> None:
        """Explicit shutdown (public API). Thread-safe."""
        with self._lock:
            self._shutdown_internal(reason="shutdown")

    def _shutdown_internal(self, reason: str = "shutdown") -> None:
        """Shutdown implementation (must hold lock)."""
        if self._client:
            try:
                self._client.close()
            except Exception as e:
                logger.warning(f"shutdown_error: {self.provider_id}, error={e}")
            self._client = None

        self._state = ProviderState.COLD
        self._increment_version()
        self._tools.clear()
        self._meta.clear()

        self._record_event(ProviderStopped(provider_id=self.provider_id, reason=reason))

    # --- Compatibility Methods ---

    def get_tool_names(self) -> List[str]:
        """Get list of available tool names."""
        with self._lock:
            return self._tools.list_names()

    def get_tools_dict(self) -> Dict[str, ToolSchema]:
        """Get tools as dictionary (for backward compatibility)."""
        with self._lock:
            return self._tools.to_dict()

    def to_status_dict(self) -> Dict[str, Any]:
        """Get status as dictionary (for registry.list)."""
        with self._lock:
            return {
                "provider": self.provider_id,
                "state": self._state.value,
                "alive": self._client is not None and self._client.is_alive(),
                "mode": self._mode.value,
                "image_or_command": self._image or self._command,
                "tools_cached": self._tools.list_names(),
                "health": self._health.to_dict(),
                "meta": dict(self._meta),
            }
