"""
Core decorators for execution monitoring.
"""
import functools
import logging
from datetime import datetime
from typing import Callable, TypeVar, Optional, Any, Protocol, runtime_checkable, Type
from uuid_extensions import uuid7str

from .models import ExecutionRun, TaskRun, StepRun, ExecutionStatus, TaskStatus, StepStatus
from .context import (
    get_current_execution_run,
    set_current_execution_run,
    get_current_task_run,
    set_current_task_run,
    set_current_step_run,
)
from . import observers

logger = logging.getLogger('zsynctech_studio_sdk.core')


T = TypeVar("T")


@runtime_checkable
class ExecutionFunction(Protocol[T]):
    """Protocol for execution functions with deploy method."""

    def __call__(self, *args: Any, **kwargs: Any) -> T:
        """Execute the function."""
        ...

    def deploy(
        self,
        instance_id: str,
        secret_key: str,
        server: str,
        port: int,
        host: str = '0.0.0.0',
        **config: Any
    ) -> 'ExecutionFunction[T]':
        """Deploy this execution to the server for remote triggering.

        Args:
            instance_id (str): Instance ID for this robot deployment.
            secret_key (str): Secret key for platform authentication.
            server (str): Platform server URL for status updates.
            port (int): Port to run server on.
            host (str, optional): Host to bind server to. Defaults to '0.0.0.0'.
            **config (Any): Additional configuration available to all tasks/steps.
        """
        ...


def execution(
    func: Optional[Callable[..., T]] = None,
    *,
    name: Optional[str] = None,
    version: Optional[str] = None,
    description: Optional[str] = None,
    exception_handlers: Optional[dict[Type[Exception], ExecutionStatus]] = None,
) -> Callable[[Callable[..., T]], ExecutionFunction[T]]:
    """Decorator to mark a function as an execution.

    Args:
        func (Callable, optional): The function to decorate.
        name (str, optional): Custom name for the execution.
        version (str, optional): Version of the execution.
        description (str, optional): Description of the execution.
        exception_handlers (dict, optional): Map of exception types to ExecutionStatus.
            When an exception occurs, if it matches a key, the execution will be marked
            with the corresponding status instead of ERROR.
    """

    def decorator(fn: Callable[..., T]) -> ExecutionFunction[T]:
        execution_name = name or fn.__name__

        @functools.wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> T:
            # Check if custom execution_id was passed
            custom_execution_id = kwargs.pop('execution_id', None)

            # Create execution run
            if custom_execution_id:
                execution_run = ExecutionRun(id=custom_execution_id, name=execution_name)
            else:
                execution_run = ExecutionRun(name=execution_name)

            execution_run.status = ExecutionStatus.RUNNING
            execution_run.start_time = datetime.now()

            # Set in context
            set_current_execution_run(execution_run)

            # Notify observers
            observers.notify_execution_start(execution_run)

            try:
                # Execute function
                result = fn(*args, **kwargs)

                # Mark as completed
                execution_run.status = ExecutionStatus.FINISHED
                execution_run.end_time = datetime.now()
                observers.notify_execution_complete(execution_run)

                return result

            except Exception as e:
                # Check if exception has custom handler
                status = ExecutionStatus.ERROR
                if exception_handlers:
                    for exc_type, exc_status in exception_handlers.items():
                        if isinstance(e, exc_type):
                            status = exc_status
                            logger.debug(f"Exception {type(e).__name__} mapped to status {status}")
                            break

                execution_run.status = status
                execution_run.error = str(e)
                execution_run.end_time = datetime.now()
                observers.notify_execution_fail(execution_run)
                raise

            finally:
                # Clear context
                set_current_execution_run(None)

        wrapper.__execution_name__ = execution_name  # type: ignore
        wrapper.__execution_version__ = version  # type: ignore
        wrapper.__execution_description__ = description  # type: ignore

        def deploy(
            instance_id: str,
            secret_key: str,
            server: str,
            port: int,
            host: str = '0.0.0.0',
            encryption_key: Optional[str] = None,
            **config: Any
        ):
            """Deploy this execution to the server for remote triggering.

            Args:
                instance_id (str): Instance ID for this robot deployment.
                secret_key (str): Secret key for platform authentication.
                server (str): Platform server URL for status updates.
                port (int): Port to run server on.
                host (str, optional): Host to bind server to. Defaults to '0.0.0.0'.
                encryption_key (str, optional): Encryption key for decrypting credentials.
                **config (Any): Additional configuration available to all tasks/steps.
            """
            from .server import _registry, serve
            from .config import set_deployment_config
            from .platform import PlatformObserver

            _registry.set_instance_id(instance_id)
            if encryption_key:
                _registry.set_encryption_key(encryption_key)

            config['instance_id'] = instance_id
            config['secret_key'] = secret_key
            config['server'] = server
            config['encryption_key'] = encryption_key
            set_deployment_config(**config)
            logger.info(f"Registered execution '{execution_name}' with instance ID '{instance_id[:13]}...'")

            _registry.register_by_instance(instance_id, wrapper)

            platform_observer = PlatformObserver(
                server=server,
                secret_key=secret_key,
                instance_id=instance_id
            )
            observers.register_observer(platform_observer)
            logger.info(f"Platform integration enabled: {server}")

            serve(host=host, port=port, blocking=True)

            return wrapper

        wrapper.deploy = deploy  # type: ignore

        # Return wrapper as ExecutionFunction for type checking
        return wrapper  # type: ignore[return-value]

    # Support both @execution and @execution()
    if func is None:
        return decorator
    return decorator(func)  # type: ignore[return-value]


def task(
    func: Optional[Callable[..., T]] = None,
    *,
    name: Optional[str] = None,
    retries: int = 0,
    description: Optional[str] = None,
    code: Optional[str] = None,
    observation: Optional[str] = None,
) -> Callable[..., T]:
    """Decorator to mark a function as a task.

    Args:
        func (Callable, optional): The function to decorate.
        name (str, optional): Custom name for the task.
        retries (int, optional): Number of retries on failure. Defaults to 0.
        description (str, optional): Description of the task.
        code (str, optional): Task code for platform identification.
        observation (str, optional): Initial observation message.
    """

    def decorator(fn: Callable[..., T]) -> Callable[..., T]:
        task_name = name or fn.__name__

        @functools.wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> T:
            # Check if custom task_id was passed
            custom_task_id = kwargs.pop('task_id', None)

            # Get current execution run from context (automatic!)
            execution_run = get_current_execution_run()

            # Generate uuid7 for code if not provided
            task_code = code if code is not None else uuid7str()
            task_description = description if description is not None else "Não informado"
            task_observation = observation if observation is not None else "Não informado"

            # Create task run
            if custom_task_id:
                task_run = TaskRun(
                    id=custom_task_id,
                    name=task_name,
                    execution_run_id=execution_run.id if execution_run else None,
                    code=task_code,
                    description=task_description,
                    observation=task_observation
                )
            else:
                task_run = TaskRun(
                    name=task_name,
                    execution_run_id=execution_run.id if execution_run else None,
                    code=task_code,
                    description=task_description,
                    observation=task_observation
                )

            task_run.status = TaskStatus.RUNNING
            task_run.start_time = datetime.now()

            set_current_task_run(task_run)

            observers.notify_task_start(task_run)

            last_exception = None
            max_attempts = retries + 1  # retries = 3 means 4 total attempts (1 original + 3 retries)
            task_completed = False

            try:
                for attempt in range(max_attempts):
                    try:
                        # Execute function
                        result = fn(*args, **kwargs)

                        # Mark as completed
                        task_run.status = TaskStatus.SUCCESS
                        task_run.end_time = datetime.now()
                        task_completed = True

                        observers.notify_task_complete(task_run)

                        return result

                    except Exception as e:
                        last_exception = e

                        # Check if we should retry
                        if attempt < max_attempts - 1:
                            # Not the last attempt - will retry
                            logger.warning(f"Task run '{task_name}' - Encountered error: {str(e)}")
                            logger.info(f"Task run '{task_name}' - Retrying (attempt {attempt + 2}/{max_attempts})")
                            import time
                            time.sleep(1)
                            continue
                        else:
                            # Last attempt failed
                            task_run.status = TaskStatus.FAIL
                            task_run.error = str(e)
                            task_run.end_time = datetime.now()
                            task_completed = True

                            observers.notify_task_fail(task_run)

                            # Update execution failed tasks counter
                            if execution_run:
                                execution_run.failed_tasks += 1

                            raise

                # This should never be reached, but just in case
                if last_exception:
                    raise last_exception
                return None  # type: ignore

            finally:
                # Always update completed_tasks counter and notify platform
                if task_completed and execution_run:
                    execution_run.completed_tasks += 1
                    observers.notify_execution_update(execution_run)

                # Clear context
                set_current_task_run(None)

        wrapper.__task_name__ = task_name  # type: ignore
        wrapper.__task_retries__ = retries  # type: ignore
        wrapper.__task_description__ = description  # type: ignore

        return wrapper

    # Support both @task and @task()
    if func is None:
        return decorator
    return decorator(func)


def step(
    func: Optional[Callable[..., T]] = None,
    *,
    code: str,  # REQUIRED
    name: Optional[str] = None,
    description: Optional[str] = None,
    observation: Optional[str] = None,
) -> Callable[..., T]:
    """Decorator to mark a function as a step.

    Args:
        func (Callable, optional): The function to decorate.
        code (str): Step code for platform identification. Required.
        name (str, optional): Custom name for the step.
        description (str, optional): Description of the step.
        observation (str, optional): Initial observation message.
    """

    def decorator(fn: Callable[..., T]) -> Callable[..., T]:
        step_name = name or fn.__name__

        @functools.wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> T:
            # Check if custom step_id was passed
            custom_step_id = kwargs.pop('step_id', None)

            # Get current task run from context (automatic!)
            task_run = get_current_task_run()

            # Get instance_id from deployment config (if available)
            from .config import _get_config_value
            instance_id = _get_config_value('instance_id')

            # Observation starts with provided value or default
            step_observation = observation if observation is not None else "Step iniciado"

            # Create step run
            if custom_step_id:
                step_run = StepRun(
                    id=custom_step_id,
                    name=step_name,
                    task_run_id=task_run.id if task_run else None,
                    automation_on_client_id=instance_id,
                    step_code=code,
                    observation=step_observation
                )
            else:
                step_run = StepRun(
                    name=step_name,
                    task_run_id=task_run.id if task_run else None,
                    automation_on_client_id=instance_id,
                    step_code=code,
                    observation=step_observation
                )

            step_run.status = StepStatus.RUNNING
            step_run.start_time = datetime.now()

            # Update task metrics
            if task_run:
                task_run.total_steps += 1

            # Set in context
            set_current_step_run(step_run)

            # Notify observers
            observers.notify_step_start(step_run)

            try:
                # Execute function
                result = fn(*args, **kwargs)

                # Mark as completed
                step_run.status = StepStatus.SUCCESS
                step_run.end_time = datetime.now()
                step_run.observation = "Step concluído com sucesso"

                # Update task metrics
                if task_run:
                    task_run.completed_steps += 1

                observers.notify_step_complete(step_run)

                return result

            except Exception as e:
                step_run.status = StepStatus.FAIL
                step_run.error = str(e)
                step_run.observation = f"Erro: {str(e)}"
                step_run.end_time = datetime.now()

                # Update task metrics
                if task_run:
                    task_run.failed_steps += 1

                observers.notify_step_fail(step_run)
                raise

            finally:
                # Clear context
                set_current_step_run(None)

        # Store metadata
        wrapper.__step_name__ = step_name  # type: ignore
        wrapper.__step_description__ = description  # type: ignore

        return wrapper

    # Support both @step and @step()
    if func is None:
        return decorator
    return decorator(func)


# Convenience functions for accessing current context
def get_current_execution() -> Optional[ExecutionRun]:
    """Get the current execution run"""
    return get_current_execution_run()


def get_current_task() -> Optional[TaskRun]:
    """Get the current task run"""
    return get_current_task_run()


def set_total_tasks(count: int) -> None:
    """
    Set the total number of tasks expected for the current execution.

    This should be called at the beginning of an execution to inform
    the platform how many tasks are expected to be processed.

    Usage:
        @execution
        def my_execution():
            set_total_tasks(10)  # Inform that 10 tasks will be processed
            for item in items:
                process_task(item)

    Args:
        count: Total number of tasks expected
    """
    execution_run = get_current_execution_run()
    if execution_run:
        execution_run.total_tasks = count
        # Notify observers about the update
        observers.notify_execution_update(execution_run)
