#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   waiter.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Datature Vi SDK resource waiter module.
"""

import time
from collections.abc import Callable
from typing import Any

from rich import print as rprint
from rich.progress import SpinnerColumn, TimeElapsedColumn
from vi.api.resources.consts import RESOURCE_WAITER_POLL_INTERVAL_SECONDS
from vi.api.responses import Condition, ConditionStatus, ResourceCondition
from vi.client.errors import ViOperationError
from vi.utils.graceful_exit import graceful_exit
from vi.utils.progress import ViProgress


class ResourceWaiter:
    """Resource waiter with timeline tracking."""

    def __init__(self):
        """Initialize the resource waiter with status messages and icons.

        Sets up message templates and icons for different resource conditions.
        """
        self.status_messages = {
            (Condition.STARTED, ConditionStatus.WAITING): "🚀 Preparing to start...",
            (
                Condition.STARTED,
                ConditionStatus.REACHED,
            ): "✅ Process started successfully",
            (
                Condition.STARTED,
                ConditionStatus.FAILED_REACH,
            ): "❌ Failed to start process",
            (
                Condition.FINISHED,
                ConditionStatus.WAITING,
            ): "⏳ Processing in progress...",
            (
                Condition.FINISHED,
                ConditionStatus.REACHED,
            ): "🎉 Process completed successfully",
            (
                Condition.FINISHED,
                ConditionStatus.FAILED_REACH,
            ): "💥 Process failed with errors",
        }

    def wait_until_done(
        self,
        callback: Callable[..., Any],
        condition: Condition | str = Condition.ALL,
        status: ConditionStatus = ConditionStatus.REACHED,
    ) -> Any:
        """Wait for an async operation to complete with detailed progress tracking.

        Args:
            callback: Function to call repeatedly to check operation status.
            condition: The condition to wait for (default: Condition.ALL for all conditions).
            status: The status to wait for (default: ConditionStatus.REACHED).

        Returns:
            The final response from the callback when the condition is met.
            Returns None if the operation is cancelled by the user.

        Raises:
            ViOperationError: If the operation fails.

        """
        with graceful_exit("Waiting operation cancelled by user") as handler:
            try:
                response = None
                with ViProgress(
                    SpinnerColumn(),
                    "[progress.description]{task.description}",
                    TimeElapsedColumn(),
                    transient=True,
                ) as progress:
                    task = progress.add_task(description="🔄 Starting process...")

                    while True:
                        # Check for keyboard interrupt
                        if handler.exit_now:
                            rprint(f"\n[yellow]{handler.message}[/yellow]")
                            return None

                        response = callback()
                        resource_conditions: list[ResourceCondition] = (
                            response.status.conditions
                        )

                        progress_chain = self._build_progress_chain(resource_conditions)
                        progress.update(task, description=progress_chain)

                        # Check conditions in a single loop
                        conditions_met = 0
                        target_condition_found = False

                        for resource_condition in resource_conditions:
                            # Check for any failed conditions first
                            if (
                                resource_condition.status
                                == ConditionStatus.FAILED_REACH
                            ):
                                error_chain = self._build_progress_chain(
                                    resource_conditions
                                )
                                # Extract and display detailed error messages
                                error_details = self._extract_error_details(response)
                                if error_details:
                                    rprint(
                                        f"\n[red]Process failed:[/red] {error_chain}"
                                    )
                                    rprint(error_details)
                                raise ViOperationError(
                                    f"Process failed at: {error_chain}"
                                )

                            # Count conditions that have reached the target status
                            if resource_condition.status == status:
                                conditions_met += 1

                            # Check if specific condition is met (for non-ALL cases)
                            if (
                                condition != Condition.ALL
                                and resource_condition.condition == condition
                                and resource_condition.status == status
                            ):
                                target_condition_found = True

                        # Determine if completion criteria are met
                        completion_met = False
                        if condition == Condition.ALL:
                            # All conditions must reach the target status
                            completion_met = bool(
                                resource_conditions
                            ) and conditions_met == len(resource_conditions)
                        else:
                            # Specific condition must reach the target status
                            completion_met = target_condition_found

                        if completion_met:
                            # Exit the progress context first to clear the spinner cleanly
                            break

                        time.sleep(RESOURCE_WAITER_POLL_INTERVAL_SECONDS)

                # Progress context has exited and spinner is now cleared
                # Print completion message cleanly
                if response is not None:
                    rprint("[green]✅ Process completed successfully![/green]")
                    return response

            except KeyboardInterrupt:
                # Handle keyboard interrupt gracefully
                rprint(f"\n[yellow]{handler.message}[/yellow]")
                return None

    def _build_progress_chain(
        self, resource_conditions: list[ResourceCondition]
    ) -> str:
        """Build a progress chain showing the journey through conditions.

        Args:
            resource_conditions: List of resource conditions to display.

        Returns:
            A formatted string showing the progress chain with icons and messages.

        """
        if not resource_conditions:
            return "🔄 Initializing process..."

        # Reverse to chronological order
        chronological = list(reversed(resource_conditions))

        progress_parts = []
        for condition in chronological:
            condition_str = (
                condition.condition.value
                if isinstance(condition.condition, Condition)
                else condition.condition
            )
            # Create proper key for status_messages lookup
            condition_key = (
                condition.condition
                if isinstance(condition.condition, Condition)
                else condition_str
            )
            message = self.status_messages.get(
                (condition_key, condition.status),
                f"🔄 {condition_str} - {condition.status.value.lower()}",
            )

            progress_parts.append(message)

        return " → ".join(progress_parts)

    def _extract_error_details(self, response: Any) -> str | None:
        """Extract detailed error messages from a failed session response.

        Uses a generic approach to find error-related information in the response
        by traversing the object structure and looking for common error patterns:
        - Fields containing 'error' or 'failed' in their names
        - Numeric counts indicating failures
        - Error dictionaries and lists
        - Reason/message fields

        Args:
            response: The session response object.

        Returns:
            Formatted error message string if errors are found, None otherwise.

        """
        error_parts = []

        def extract_from_object(obj: Any, path: str = "") -> None:
            """Recursively extract error information from an object."""
            if obj is None:
                return

            # Handle lists (e.g., events, errors)
            if isinstance(obj, list):
                for item in obj:
                    extract_from_object(item, path)
                return

            # Handle dictionaries (e.g., file errors)
            if isinstance(obj, dict):
                # Check if it looks like an error mapping (filename -> error_code)
                for key, value in obj.items():
                    if isinstance(value, (str, int)):
                        error_parts.append(f"  • {key}: {value}")
                return

            # Handle objects with attributes
            if hasattr(obj, "__dict__") or hasattr(obj, "__dataclass_fields__"):
                # Look for error-related numeric fields
                for attr_name in dir(obj):
                    if attr_name.startswith("_"):
                        continue

                    try:
                        attr_value = getattr(obj, attr_name, None)
                    except (AttributeError, TypeError, ValueError):
                        # Ignore attribute access errors during introspection
                        continue

                    # Check for error count fields (errored, failed, error_count, etc.)
                    if (
                        "error" in attr_name.lower() or "failed" in attr_name.lower()
                    ) and isinstance(attr_value, int):
                        if attr_value > 0:
                            # Make the field name more readable
                            field_display = attr_name.replace("_", " ").title()
                            error_parts.append(f"  • {field_display}: {attr_value}")

                    # Check for reason/message fields
                    elif attr_name.lower() in ["reason", "message", "error", "detail"]:
                        if attr_value and isinstance(attr_value, str):
                            error_parts.append(f"  • {attr_name.title()}: {attr_value}")

                    # Check for nested error objects (files, events, etc.)
                    elif (
                        attr_name.lower()
                        in [
                            "events",
                            "files",
                            "errors",
                        ]
                        and attr_value
                    ):
                        extract_from_object(attr_value, attr_name)

        # Start extraction from response.status if it exists
        if hasattr(response, "status"):
            extract_from_object(response.status)

        if error_parts:
            # Remove duplicates while preserving order
            unique_parts = []
            seen = set()
            for part in error_parts:
                if part not in seen:
                    unique_parts.append(part)
                    seen.add(part)

            return "\n".join(["[red]Error details:[/red]"] + unique_parts)

        return None
