# Copyright (c) 2025 OpenStack Foundation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import datetime
import os
import signal
import socket
import sys
import threading
import time

from oslo_config import cfg
from oslo_log import log as logging
import oslo_messaging
from oslo_messaging import server as rpc_server
from oslo_messaging import notify as rpc_notify
import tooz
from tooz import coordination

from vigie import config
from vigie.db.sqlalchemy.database import get_session as get_db_session
from vigie.db.sqlalchemy.database import init_database_config
from vigie.db.sqlalchemy.models import VigieTasks, TaskState

# Initialize logging
LOG = logging.getLogger(__name__)

# Define RPC topics
RPC_TOPIC = 'vigie.tasks'


def sd_notify(state: str) -> bool:
    """Send a notification to systemd."""
    notify_socket = os.getenv("NOTIFY_SOCKET")
    if not notify_socket:
        return False

    # Abstract namespace socket if starts with '@'
    if notify_socket[0] == '@':
        notify_socket = '\0' + notify_socket[1:]

    sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
    try:
        sock.connect(notify_socket)
        sock.sendall(state.encode())
    except Exception as e:
        LOG.debug(f"🐛 Failed to notify systemd: {e}")
        return False
    finally:
        sock.close()
    return True


class TaskHandler:
    """RPC Message Handler for incoming task notifications."""

    def __init__(self, worker_service):
        self.worker_service = worker_service

    def create_task(self, ctxt, task_id, task_type, payload=None):
        """Handle incoming task creation request from RabbitMQ.

        :param ctxt: RPC context
        :param task_id: Unique task identifier
        :param task_type: Type of task (e.g., 'EVACUATE', 'PROJECT_PURGE')
        :param payload: Task-specific data dictionary
        """
        LOG.info(f"⚙ Received task notification: task_id={task_id}, type={task_type}")

        try:
            # Store task in database
            session = get_db_session()
            try:
                new_task = VigieTasks(
                    task_id=task_id,
                    task_type=task_type,
                    payload=payload or {},
                    state=TaskState.SCHEDULED,
                    created_at=datetime.datetime.utcnow()
                )
                session.add(new_task)
                session.commit()
                LOG.info(f"⚙ Task {task_id} stored in database")

                # Trigger immediate processing if we have the lock
                if self.worker_service.attempt_worker_lock():
                    try:
                        self.worker_service.run_tasks()
                    finally:
                        self.worker_service.release_worker_lock()

                return {'status': 'accepted', 'task_id': task_id}
            finally:
                session.close()

        except Exception as e:
            LOG.error(f"⧱ Error processing task {task_id}: {e}", exc_info=True)
            return {'status': 'failed', 'error': str(e)}


class WorkerService(object):
    """Main worker service."""

    def __init__(self, CONF):
        self.CONF = CONF
        self.node_id = CONF.worker.my_hostname.encode('utf-8')
        self.group_name = b'vigie_workers'
        self.running = False
        self.coordinator = None
        self.worker_lock = None
        self._stop_event = threading.Event()
        self.rpc_server = None
        self.task_handler = None

        # Setup signal handlers
        signal.signal(signal.SIGTERM, self._signal_handler)
        signal.signal(signal.SIGINT, self._signal_handler)

    def _signal_handler(self, signum, frame):
        """Handle shutdown signals."""
        LOG.info("⚙ Vigie worker is shutting down...")
        self.stop()

    def initialize_coordination(self) -> bool:
        """Initialize tooz coordination."""
        if not self.CONF.worker.coordination_url:
            LOG.error("⧱ Coordination URL not configured")
            return False

        try:
            self.coordinator = coordination.get_coordinator(
                self.CONF.worker.coordination_url,
                self.node_id
            )
            self.coordinator.start()

            # Join the worker group
            try:
                self.coordinator.join_group(self.group_name).get()
            except coordination.GroupNotCreated:
                # Try to create the group (first worker)
                try:
                    self.coordinator.create_group(self.group_name).get()
                    self.coordinator.join_group(self.group_name).get()
                except coordination.GroupAlreadyExist:
                    # Another worker created it, join normally
                    self.coordinator.join_group(self.group_name).get()
            except coordination.ToozError as e:
                LOG.error(f"⧱ Failed to join coordination group: {e}")
                return False

            # Create scheduler lock
            self.worker_lock = self.coordinator.get_lock(b'vigie_worker_lock')

            LOG.info(f"⚙ Successfully initialized coordination with node_id: {self.CONF.worker.my_hostname}")
            return True

        except coordination.ToozError as e:
            LOG.error(f"⧱ Tooz coordination error: {e}")
            return False
        except Exception as e:
            LOG.error(f"⧱ Failed to initialize coordination: {e}")
            return False

    def initialize_messaging(self) -> bool:
        """Initialize oslo.messaging for RabbitMQ."""
        try:
            # Create RPC transport
            self.rpc_transport = oslo_messaging.get_transport(self.CONF)

            # Create endpoint for handling messages
            self.task_handler = TaskHandler(self)
            endpoints = [self.task_handler]

            # Create server listening on the task topic
            self.rpc_server = oslo_messaging.get_rpc_server(
                self.rpc_transport,
                oslo_messaging.Target(topic=RPC_TOPIC),
                endpoints,
                executor='threading',
            )

            LOG.info(f"⚙ Messaging initialized on topic: {RPC_TOPIC}")
            return True

        except Exception as e:
            LOG.error(f"⧱ Failed to initialize messaging: {e}")
            return False

    def start(self) -> bool:
        """Start the worker service."""
        LOG.info("⚙ Vigie worker is starting up...")
        LOG.info(f"⚙ Vigie worker Service (node_id: {self.CONF.worker.my_hostname})")

        # Initialize database configuration
        init_database_config(self.CONF)

        if not self.initialize_coordination():
            LOG.error("⧱ Failed to initialize coordination, exiting")
            return False

        if not self.initialize_messaging():
            LOG.error("⧱ Failed to initialize messaging, exiting")
            return False

        self.running = True

        # Notify systemd that we're ready
        sd_notify("READY=1")
        LOG.info("⚙ Vigie worker is ready.")

        # Start RPC server in a separate thread
        self.rpc_thread = threading.Thread(target=self.rpc_server.wait)
        self.rpc_thread.daemon = True
        self.rpc_thread.start()

        # Main worker loop (for periodic checks)
        while self.running and not self._stop_event.is_set():
            try:
                self.run_periodic_check()
                self._stop_event.wait(self.CONF.worker.check_interval)
            except Exception as e:
                LOG.error(f"⧱ Error in periodic check: {e}", exc_info=True)
                self._stop_event.wait(10)  # Wait before retry

        self.cleanup()
        return True

    def stop(self):
        """Stop the worker service."""
        LOG.info("⚙ Vigie worker is shutting down...")
        self.running = False
        self._stop_event.set()
        sd_notify("STOPPING=1")

        # Stop RPC server
        if self.rpc_server:
            try:
                self.rpc_server.stop()
            except Exception as e:
                LOG.warning(f"⚠ Error stopping RPC server: {e}")

    def cleanup(self):
        """Cleanup resources."""
        if self.coordinator:
            try:
                self.coordinator.leave_group(self.group_name).get()
                self.coordinator.stop()
            except Exception as e:
                LOG.warning(f"⚠ Error during coordinator cleanup: {e}")

        if self.rpc_server:
            try:
                self.rpc_server.wait()
            except Exception as e:
                LOG.warning(f"⚠ Error during RPC cleanup: {e}")

    def run_tasks(self):
        """Process pending tasks from database."""
        session = get_db_session()
        try:
            # Fetch pending tasks
            pending_tasks = session.query(VigieTasks).filter(
                VigieTasks.state == TaskState.SCHEDULED
            ).limit(10).all()

            for task in pending_tasks:
                try:
                    LOG.info(f"⚙ Processing task: {task.task_id}")
                    # TODO: Implement actual task processing logic
                    task.state = TaskState.DONE
                    task.completed_at = datetime.datetime.utcnow()
                    session.commit()
                except Exception as e:
                    LOG.error(f"⧱ Error processing task {task.task_id}: {e}", exc_info=True)
                    task.state = TaskState.ERROR
                    task.failed_at = datetime.datetime.utcnow()
                    session.commit()
        finally:
            session.close()

    def run_periodic_check(self):
        """Run the periodic check cycle."""
        LOG.debug("🐛 Running periodic check cycle")

        # Attempt to become scheduler if we can acquire the lock
        if self.attempt_worker_lock():
            try:
                self.run_tasks()
            except Exception as e:
                LOG.error(f"⧱ Error scheduling new migrations: {e}", exc_info=True)
            finally:
                self.release_worker_lock()

    def attempt_worker_lock(self) -> bool:
        """Attempt to acquire scheduler role."""
        if not self.worker_lock:
            return False

        try:
            return self.worker_lock.acquire(blocking=False)
        except Exception as e:
            LOG.warning(f"⚠ Failed to acquire scheduler lock: {e}")
            return False

    def release_worker_lock(self):
        """Release scheduler role."""
        if self.worker_lock and self.worker_lock.acquired:
            try:
                self.worker_lock.release()
            except Exception as e:
                LOG.warning(f"⚠ Failed to release scheduler lock: {e}")


def main():
    """Main entry point for the worker service."""
    try:
        # Initialize configuration
        CONF = config.init_config()

        # Setup logging
        logging.setup(CONF, 'vigie-worker')
        LOG.info("⚙ Vigie worker is starting up...")

        # Create and start worker service
        worker = WorkerService(CONF)

        # Start the worker
        success = worker.start()

        if success:
            LOG.info("⚙ Vigie worker stopped normally")
            sys.exit(0)
        else:
            LOG.error("⧱ Vigie worker failed to start")
            sys.exit(1)

    except Exception as e:
        LOG.error(f"⧱ Fatal error in worker service: {e}", exc_info=True)
        sys.exit(1)
