"""Action module for phone copilot automation.

This module provides action classes for controlling Android and HarmonyOS devices
through ADB and HDC interfaces. It includes various touch actions, navigation,
text input, and application launch capabilities.
"""

import ast
import time

from loguru import logger
from pydantic import BaseModel, ConfigDict
from pydantic_core import core_schema
from xpdc.adb.device import ADBDevice
from xpdc.hdc.device import HDCDevice

from phone_copilot import apps
from phone_copilot.model_client import ModelResponse

__all__ = [
    'Action',
    'ActionResult',
    'Finish',
    'Tap',
    'DoubleTap',
    'LongPress',
    'Type',
    'Swipe',
    'Back',
    'Home',
    'Wait',
    'Launch',
]


class ActionParam(BaseModel):
    model_config = ConfigDict(extra='allow')

    # Finish
    message: str | None = None

    # Touch
    element: list[int] = None
    start: list[int] = None
    end: list[int] = None

    # Wait
    duration: str = None

    # Launch
    app: str = None

    # Type
    text: str = None


class ActionBase:
    """Base class for all device action types.

    Provides common functionality for executing actions on mobile devices,
    including coordinate calibration and parameter management.

    Attributes:
        _device: The ADB or HDC device instance.
        _width: Screen width in pixels.
        _height: Screen height in pixels.
        _ccf: Coordinate calibration factor for scaling coordinates.
        _param: Validated action parameters.
    """

    def __init__(self, *, device: ADBDevice | HDCDevice, width: int, height: int, ccf: int = None, **kwargs) -> None:
        """Initialize the action with device and screen parameters.

        Args:
            device: ADB or HDC device instance to control.
            width: Screen width in pixels.
            height: Screen height in pixels.
            ccf: Optional coordinate calibration factor.
            **kwargs: Additional parameters to be validated as ActionParam.
        """
        self._device, self._width, self._height, self._ccf = device, width, height, ccf
        self._param = ActionParam.model_validate(kwargs)

    @property
    def device(self) -> ADBDevice | HDCDevice:
        """Get the device instance.

        Returns:
            The ADB or HDC device instance.
        """
        return self._device

    @property
    def width(self) -> int:
        """Get the screen width.

        Returns:
            Screen width in pixels.
        """
        return self._width

    @property
    def height(self) -> int:
        """Get the screen height.

        Returns:
            Screen height in pixels.
        """
        return self._height

    @property
    def ccf(self) -> int | None:
        """Get the coordinate calibration factor.

        Returns:
            Coordinate calibration factor, or None if not set.
        """
        return self._ccf

    def coord_calibrate(self, *, element: list[int]) -> tuple[int, int]:
        """Calibrate coordinates based on calibration factor.

        Transforms coordinates from a potentially scaled coordinate system
        to the actual device screen coordinates using the calibration factor.

        Args:
            element: Original coordinates as [x, y].

        Returns:
            Tuple of calibrated (x, y) coordinates.
        """
        if self._ccf:
            x, y = int(element[0] / self._ccf * self._width), int(element[1] / self._ccf * self._height)

            if x > self._width or y > self._height:
                logger.warning(
                    f'Calibrated coordinates ({x}, {y}) exceed screen dimensions ({self._width}, {self._height}), '
                    f'using original coordinates {element} instead. Fix ccr or check the model output.'
                )
                return element[0], element[1]

            logger.debug(f'Calibrating coordinates {element} with ccf {self._ccf} to ({x}, {y})')
            return x, y

        return element[0], element[1]

    @property
    def param(self) -> ActionParam:
        """Get the action parameters.

        Returns:
            Validated ActionParam instance.
        """
        return self._param

    def _execute(self) -> 'ActionResult':
        """Execute the action implementation.

        This method should be overridden by subclasses to implement
        specific action behavior.

        Returns:
            ActionResult with execution status and details.
        """
        pass

    def __str__(self) -> str:
        """Get string representation of the action.

        Returns:
            The class name as string.
        """
        return self.__class__.__name__

    def __call__(self) -> 'ActionResult':
        """Execute the action with logging and error handling.

        Returns:
            ActionResult containing success status and any messages.
        """
        try:
            logger.info(
                f'Executing action: [{self}] with param: '
                f'{self._param.model_dump_json(exclude_none=True, ensure_ascii=False)}'
            )

            result = self._execute()

            logger.info(
                f'Action {self} executed with result: '
                f'{result.model_dump_json(exclude_none=True, ensure_ascii=False, exclude={"action"})}'
            )
            return result
        except Exception as e:  # noqa: BLE001
            msg = f'Exception during action: {e}'
            return ActionResult(success=False, should_finish=False, action=self, message=msg)

    @classmethod
    def __get_pydantic_core_schema__(cls, *args, **kwargs) -> core_schema.CoreSchema:
        """Custom serializer for Pydantic models.

        Returns:
            CoreSchema defining how to serialize ActionBase instances.
        """

        def serialize(v: 'ActionBase') -> dict:
            data = {'name': str(v)}
            data.update(v.param.model_dump(mode='json', exclude_none=True))
            return data

        return core_schema.no_info_plain_validator_function(
            lambda v: v, serialization=core_schema.plain_serializer_function_ser_schema(serialize)
        )


class ActionResult(BaseModel):
    """Result of an action execution.

    Contains information about the success status, whether the task should
    finish, and any associated messages.

    Attributes:
        action: The action instance that was executed.
        success: Whether the action executed successfully.
        should_finish: Whether the task should terminate after this action.
        message: Optional message with additional information.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    action: 'ActionBase'

    success: bool
    should_finish: bool

    message: str = None


class Action(ActionBase):
    """Factory class for creating action instances from model responses.

    Parses model output and instantiates the appropriate action subclass.
    Supports various action types including finish, tap, swipe, type, etc.
    """

    def __new__(cls, *, model_response: ModelResponse, **kwargs) -> 'ActionBase':
        """Create an action instance from model response.

        Args:
            model_response: Response from the model containing action string.
            **kwargs: Additional parameters passed to action constructor.

        Returns:
            An instance of the appropriate ActionBase subclass.
        """
        try:
            logger.info('Parsing action from model response...')

            action_str = model_response.action.strip()

            if action_str.startswith('finish'):
                return Finish(message=action_str.replace('finish(message=', '')[1:-2], **kwargs)

            elif action_str.startswith('do'):
                action_str = action_str.translate(str.maketrans({'\n': '\\n', '\r': '\\r', '\t': '\\t'})).strip(  # noqa: B005
                    '\\n\\r\\t'
                )
                tree = ast.parse(action_str, mode='eval')

                if not isinstance(tree.body, ast.Call):
                    return UnknownActionType(message=f'Expected a function call: {action_str}', **kwargs)

                call = tree.body
                keywords = {keyword.arg: ast.literal_eval(keyword.value) for keyword in call.keywords}

                try:
                    action = keywords.pop('action')
                except KeyError:
                    return UnknownActionType(message=f'Missing action: {action_str}', **kwargs)

                action_cls = {
                    'launch': Launch,
                    'tap': Tap,
                    'click': Tap,
                    'double tap': DoubleTap,
                    'double click': DoubleTap,
                    'long press': LongPress,
                    'type': Type,
                    'swipe': Swipe,
                    'back': Back,
                    'home': Home,
                    'wait': Wait,
                }.get(action.lower())

                return (
                    action_cls(**keywords, **kwargs)
                    if action_cls
                    else UnknownAction(message=f'Unknown action: {action_str}', **kwargs)
                )

            else:
                return UnknownActionType(message=f'Unknown action type: {action_str}', **kwargs)
        except Exception as e:  # noqa: BLE001
            msg = f'Failed to parse action: {model_response.action}, error: {e}'
            logger.error(msg)

            return UnknownActionType(message=msg, **kwargs)


class Launch(ActionBase):
    """Launch an application on the device.

    Opens the specified application using its package name and activity.
    """

    def _execute(self) -> ActionResult:
        """Execute the app launch action.

        Returns:
            ActionResult indicating success or failure with message.
        """
        if not self._param.app:
            return ActionResult(
                success=False, should_finish=False, message='No app specified for Launch action.', action=self
            )

        if app := apps.get(device=self._device, app_name=self._param.app):
            self._device.launch_app(**app.model_dump(exclude_none=True))
            time.sleep(1.0)
            return ActionResult(success=True, should_finish=False, action=self)

        return ActionResult(
            success=False,
            should_finish=False,
            message=f'Unknown app: {self._param.app}, please launch through the UI.',
            action=self,
        )


class Tap(ActionBase):
    """Perform a single tap at specified coordinates.

    Simulates a touch tap event on the device screen.
    """

    def _execute(self) -> ActionResult:
        """Execute the tap action.

        Returns:
            ActionResult indicating success.
        """
        x, y = self.coord_calibrate(element=self._param.element)
        self._param.element = [x, y]
        self._device.tap(x, y)
        time.sleep(1.0)
        return ActionResult(success=True, should_finish=False, action=self)


class DoubleTap(ActionBase):
    """Perform a double tap at specified coordinates.

    Simulates two quick consecutive taps on the device screen.
    """

    def _execute(self) -> ActionResult:
        """Execute the double tap action.

        Returns:
            ActionResult indicating success.
        """
        x, y = self.coord_calibrate(element=self._param.element)
        self._param.element = [x, y]
        self._device.double_tap(x=x, y=y)
        time.sleep(1.0)
        return ActionResult(success=True, should_finish=False, action=self)


class LongPress(ActionBase):
    """Perform a long press at specified coordinates.

    Simulates a prolonged touch on the device screen.
    """

    def _execute(self) -> ActionResult:
        """Execute the long press action.

        Returns:
            ActionResult indicating success.
        """
        x, y = self.coord_calibrate(element=self._param.element)
        self._param.element = [x, y]
        self._device.long_press(x=x, y=y)
        time.sleep(1.0)
        return ActionResult(success=True, should_finish=False, action=self)


class Type(ActionBase):
    """Type text input on the device.

    Clears existing text and inputs new text content.
    """

    def _execute(self) -> ActionResult:
        """Execute the text input action.

        Returns:
            ActionResult indicating success.
        """
        self._device.clear_text()
        self._device.type_text(self._param.text)
        time.sleep(1.0)
        return ActionResult(success=True, should_finish=False, action=self)


class Swipe(ActionBase):
    """Perform a swipe gesture from start to end coordinates.

    Simulates a drag or swipe motion across the screen.
    """

    def _execute(self) -> ActionResult:
        """Execute the swipe action.

        Returns:
            ActionResult indicating success.
        """
        start_x, start_y = self.coord_calibrate(element=self._param.start)
        end_x, end_y = self.coord_calibrate(element=self._param.end)
        self._param.start = [start_x, start_y]
        self._param.end = [end_x, end_y]
        self._device.swipe(start_x=start_x, start_y=start_y, end_x=end_x, end_y=end_y)
        time.sleep(1.0)
        return ActionResult(success=True, should_finish=False, action=self)


class Back(ActionBase):
    """Press the back button on the device.

    Simulates pressing the system back navigation button.
    """

    def _execute(self) -> ActionResult:
        """Execute the back button action.

        Returns:
            ActionResult indicating success.
        """
        self._device.back()
        time.sleep(1.0)
        return ActionResult(success=True, should_finish=False, action=self)


class Home(ActionBase):
    """Press the home button on the device.

    Simulates pressing the system home button to return to launcher.
    """

    def _execute(self) -> ActionResult:
        """Execute the home button action.

        Returns:
            ActionResult indicating success.
        """
        self._device.home()
        time.sleep(1.0)
        return ActionResult(success=True, should_finish=False, action=self)


class Wait(ActionBase):
    """Wait for a specified duration.

    Pauses execution for the given number of seconds.
    """

    def _execute(self) -> ActionResult:
        """Execute the wait action.

        Returns:
            ActionResult indicating success after waiting.
        """
        duration_str = self._param.duration or '1 seconds'

        try:
            duration = int(duration_str.replace('seconds', '').strip())
        except ValueError:
            duration = 1

        time.sleep(duration)

        return ActionResult(success=True, should_finish=False, action=self)


class Finish(ActionBase):
    """Signal task completion.

    Indicates that the automation task has finished successfully.
    """

    def _execute(self) -> ActionResult:
        """Execute the finish action.

        Returns:
            ActionResult with should_finish=True and optional message.
        """
        return ActionResult(success=True, should_finish=True, action=self, message=self._param.message)


class UnknownActionType(ActionBase):
    """Represents an action with unrecognized type format.

    Used when the action string cannot be parsed into a known format.
    """

    def _execute(self) -> ActionResult:
        """Return failure result for unknown action type.

        Returns:
            ActionResult with success=False and error message.
        """
        return ActionResult(success=False, should_finish=False, action=self, message=self._param.message)


class UnknownAction(ActionBase):
    """Represents an action with unrecognized action name.

    Used when the action name is not in the list of supported actions.
    """

    def _execute(self) -> ActionResult:
        """Return failure result for unknown action.

        Returns:
            ActionResult with success=False and error message.
        """
        return ActionResult(success=False, should_finish=False, action=self, message=self._param.message)
