import time
import traceback

from concurrent.futures import ThreadPoolExecutor
from typing import Any, Dict, Optional, Tuple, Type

from clipped.utils.tz import now

from polyaxon import settings
from polyaxon.auxiliaries import V1PolyaxonInitContainer, V1PolyaxonSidecarContainer
from polyaxon.connections import V1Connection
from polyaxon.exceptions import PolyaxonConverterError
from polyaxon.lifecycle import LiveState, V1Statuses
from polyaxon.logger import logger
from polyaxon.runner.agent.client import AgentClient
from polyaxon.runner.executor import BaseExecutor
from polyaxon.schemas.cli.checks_config import ChecksConfig
from polyaxon.schemas.responses.v1_agent_state_response import V1AgentStateResponse


class BaseAgent:
    EXECUTOR: Type[BaseExecutor]
    HEALTH_FILE = "/tmp/.healthz"
    SLEEP_STOP_TIME = 60 * 5
    SLEEP_ARCHIVED_TIME = 60 * 60
    IS_ASYNC = False

    def __init__(
        self,
        owner: Optional[str] = None,
        agent_uuid: Optional[str] = None,
        sleep_interval: Optional[int] = None,
    ):
        self.sleep_interval = sleep_interval
        if self.sleep_interval:
            self.sleep_interval = max(self.sleep_interval, 1)
        self.executor = None
        self._executor_refreshed_at = now()
        self._graceful_shutdown = False
        self.client = AgentClient(
            owner=owner, agent_uuid=agent_uuid, is_async=self.IS_ASYNC
        )
        self.executor = self.EXECUTOR()
        self.content = settings.AGENT_CONFIG.to_json()

    def sync(self):
        raise NotImplementedError

    def sync_compatible_updates(self, compatible_updates: Dict):
        if compatible_updates and settings.AGENT_CONFIG:
            init = compatible_updates.get("init")
            if init and settings.AGENT_CONFIG.init:
                init = V1PolyaxonInitContainer.from_dict(init)
                settings.AGENT_CONFIG.init = settings.AGENT_CONFIG.init.patch(init)

            sidecar = compatible_updates.get("sidecar")
            if sidecar and settings.AGENT_CONFIG.sidecar:
                sidecar = V1PolyaxonSidecarContainer.from_dict(sidecar)
                settings.AGENT_CONFIG.sidecar = settings.AGENT_CONFIG.sidecar.patch(
                    sidecar
                )
            connections = compatible_updates.get("connections")
            if connections:
                settings.AGENT_CONFIG.connections = [
                    V1Connection.from_dict(c) for c in connections
                ]

            self.content = settings.AGENT_CONFIG.to_json()
            return self.sync()

    @classmethod
    def get_healthz_config(cls) -> Optional[ChecksConfig]:
        try:
            return ChecksConfig.read(cls.HEALTH_FILE, config_type=".json")
        except Exception:  # noqa
            return None

    @classmethod
    def ping(cls):
        if not settings.AGENT_CONFIG.enable_health_checks:
            return
        ChecksConfig.init_file(cls.HEALTH_FILE)
        config = cls.get_healthz_config()
        if config:
            config.last_check = now()
            config.write(cls.HEALTH_FILE, mode=config._WRITE_MODE)

    @classmethod
    def pong(cls, interval: int = 15) -> bool:
        config = cls.get_healthz_config()
        if not config:
            return False
        return not config.should_check(interval=interval)

    def refresh_executor(self):
        raise NotImplementedError

    def start(self):
        raise NotImplementedError

    def end(self, sleep: Optional[int] = None):
        self._graceful_shutdown = True
        if sleep:
            time.sleep(sleep)
        else:
            logger.info("Agent is shutting down.")

    def _check_status(self, agent_state):
        if agent_state.status == V1Statuses.STOPPED:
            print(
                "Agent has been stopped from the platform,"
                "but the deployment is still running."
                "Please either set the agent to starting or teardown the agent deployment."
            )
            return self.end(sleep=self.SLEEP_STOP_TIME)
        elif agent_state.live_state < LiveState.LIVE:
            print(
                "Agent has been archived from the platform,"
                "but the deployment is still running."
                "Please either restore the agent or teardown the agent deployment."
            )
            return self.end(sleep=self.SLEEP_ARCHIVED_TIME)

    def process(self, pool: "ThreadPoolExecutor") -> V1AgentStateResponse:
        raise NotImplementedError

    def make_run_resource(
        self,
        owner_name: str,
        project_name: str,
        run_name: str,
        run_uuid: str,
        content: str,
        default_auth: bool = False,
    ) -> Optional[Any]:
        try:
            return self.executor.make_and_convert(
                owner_name=owner_name,
                project_name=project_name,
                run_uuid=run_uuid,
                run_name=run_name,
                content=content,
                default_auth=default_auth,
            )
        except PolyaxonConverterError as e:
            logger.info(
                "Run could not be cleaned. Agent failed converting run manifest: {}\n{}".format(
                    repr(e), traceback.format_exc()
                )
            )
        except Exception as e:
            logger.info(
                "Agent failed during compilation with unknown exception: {}\n{}".format(
                    repr(e), traceback.format_exc()
                )
            )
        return None

    def prepare_run_resource(
        self,
        owner_name: str,
        project_name: str,
        run_name: str,
        run_uuid: str,
        content: str,
    ) -> Optional[Any]:
        raise NotImplementedError

    def submit_run(self, run_data: Tuple[str, str, str, str]):
        raise NotImplementedError

    def make_and_create_run(
        self, run_data: Tuple[str, str, str, str], default_auth: bool = False
    ):
        raise NotImplementedError

    def apply_run(self, run_data: Tuple[str, str, str, str]):
        raise NotImplementedError

    def check_run(self, run_data: Tuple[str, str]):
        raise NotImplementedError

    def stop_run(self, run_data: Tuple[str, str]):
        raise NotImplementedError

    def delete_run(self, run_data: Tuple[str, str, str, str]):
        raise NotImplementedError

    def clean_run(self, run_uuid: str, run_kind: str):
        raise NotImplementedError
