from __future__ import annotations

import logging
import re
import threading
import time
from argparse import ArgumentParser
from contextlib import contextmanager
from datetime import timedelta
from typing import Iterator, cast

from django.core.management.base import BaseCommand

from django_goals.models import (
    handle_unblocked_goals, handle_waiting_for_date,
    handle_waiting_for_failed_preconditions, handle_waiting_for_preconditions,
    handle_waiting_for_worker, remove_old_goals,
)
from django_goals.pickups import PickupMonitorThread

from .goals_busy_worker import stop_signal_handler


logger = logging.getLogger(__name__)


class Command(BaseCommand):
    help = 'Run the worker'

    def add_arguments(self, parser: ArgumentParser) -> None:
        parser.add_argument(
            '--threads',
            action='append',  # This allows multiple instances of --threads
            help='Number of threads and optional deadline horizon (e.g., "3", "2:30m")',
        )
        parser.add_argument(
            '--once',
            action='store_true',
            help='Exit when no work is available',
        )

    def handle(self, *args: object, **options: object) -> None:
        # Parse the threads parameter
        threads_specs = cast(list[str], options.get('threads') or ['1'])  # Default to 1 worker if not specified
        worker_specs = []
        for spec in threads_specs:
            try:
                if ':' in spec:
                    # Format is N:HORIZON
                    count_str, horizon_str = spec.split(':', 1)
                    count = int(count_str)
                    horizon = parse_duration(horizon_str)
                else:
                    # Format is just N
                    count = int(spec)
                    horizon = None

                if count <= 0:
                    raise ValueError("Thread count must be a positive integer")
                worker_specs.append((count, horizon))

            except ValueError as e:
                self.stderr.write(f"Error parsing thread count '{spec}': {e}")
                return

        with stop_signal_handler() as stop_event:
            threaded_worker(
                worker_specs=worker_specs,
                stop_event=stop_event,
                once=cast(bool, options['once']),
            )


def parse_duration(duration_str: str) -> timedelta | None:
    """
    Parse duration strings like "30m", "2h", "1d" into timedelta objects.
    Returns None if the string is "none", empty, or None.
    """
    if not duration_str or duration_str.lower() == 'none':
        return None

    units = {
        's': 'seconds',
        'm': 'minutes',
        'h': 'hours',
        'd': 'days',
        'w': 'weeks'
    }

    matches = re.match(r'^(\d+)([smhdw])$', duration_str.lower())
    if not matches:
        raise ValueError(f"Invalid duration format: {duration_str}. Use format like '30m', '2h', '1d'")

    value, unit = matches.groups()
    kwargs = {units[unit]: int(value)}
    return timedelta(**kwargs)


def threaded_worker(
    worker_specs: list[tuple[int, timedelta | None]],
    stop_event: threading.Event | None = None,
    once: bool = False,
) -> None:
    if stop_event is None:
        stop_event = threading.Event()

    pickup_monitor = PickupMonitorThread()
    pickup_monitor.start()

    total_workers = sum(count for count, _ in worker_specs)
    workers_state = WorkersState(total_workers + 1)  # +1 for transitions thread

    threads: list[threading.Thread] = [
        TransitionsThread(
            stop_event=stop_event,
            once=once,
            workers_state=workers_state,
            thread_id="transitions",
        ),
    ]

    worker_id = 0
    for count, horizon in worker_specs:
        for i in range(count):
            threads.append(
                HeavyLiftingThread(
                    stop_event=stop_event,
                    once=once,
                    workers_state=workers_state,
                    thread_id=f"worker_{worker_id}",
                    deadline_horizon=horizon,
                    pickup_monitor=pickup_monitor,
                )
            )
            worker_id += 1

    for thread in threads:
        thread.start()

    for thread in threads:
        thread.join()

    pickup_monitor.shutdown()
    pickup_monitor.join()


class HeavyLiftingThread(threading.Thread):
    def __init__(
        self,
        stop_event: threading.Event,
        once: bool,
        workers_state: WorkersState,
        thread_id: str,
        deadline_horizon: timedelta | None = None,
        pickup_monitor: PickupMonitorThread | None = None,
    ):
        super().__init__()
        self.stop_event = stop_event
        self.once = once
        self.workers_state = workers_state
        self.thread_id = thread_id
        self.deadline_horizon = deadline_horizon
        self.pickup_monitor = pickup_monitor

    def run(self) -> None:
        logger.info('Busy-wait worker started, deadline_horizon: %s', self.deadline_horizon)

        while not self.stop_event.is_set():
            with self.workers_state.work_session(self.thread_id):
                try:
                    progress = handle_waiting_for_worker(
                        deadline_horizon=self.deadline_horizon,
                        pickup_monitor=self.pickup_monitor,
                    )
                except Exception as e:
                    logger.exception(e)
                    # Treat exceptions as if we didn't do work
                    progress = None

                did_work = progress is not None
                self.workers_state.report_work(self.thread_id, did_work)

            if self.workers_state.all_idle and self.once:
                logger.info('All threads are idle. Exiting because of `once` flag.')
                break

            if not did_work:
                logger.debug('Nothing to do, sleeping for a bit')
                time.sleep(1)

        logger.info('Busy-wait worker exiting')


class TransitionsThread(threading.Thread):
    def __init__(
        self,
        stop_event: threading.Event,
        once: bool,
        workers_state: WorkersState,
        thread_id: str,
    ):
        super().__init__()
        self.stop_event = stop_event
        self.once = once
        self.workers_state = workers_state
        self.thread_id = thread_id

    def run(self) -> None:
        logger.info('Transitions worker started')

        while not self.stop_event.is_set():
            with self.workers_state.work_session(self.thread_id):
                did_work = self._run_handlers()
                self.workers_state.report_work(self.thread_id, did_work)

            if self.workers_state.all_idle and self.once:
                logger.info('All threads are idle. Exiting because of `once` flag.')
                break

            if not did_work:
                logger.debug('Nothing to do, sleeping for a bit')
                time.sleep(1)

        logger.info('Transitions worker exiting')

    def _run_handlers(self) -> bool:
        """Run all handlers and return True if any work was done"""
        try:
            results = [
                handle_waiting_for_preconditions(),
                handle_waiting_for_failed_preconditions(),
                handle_waiting_for_date(),
                handle_unblocked_goals(),
                remove_old_goals()
            ]
            return any(r for r in results if r)
        except Exception as e:
            logger.exception(e)
            return False  # No work done if there was an exception


class WorkersState:
    def __init__(self, thread_count: int):
        self.lock = threading.Lock()
        self.idle_threads: set[str] = set()  # Threads that are permanently idle (no more work)
        self.active_threads: set[str] = set()  # Threads currently registered as active
        self.total_threads = thread_count

    @contextmanager
    def work_session(self, thread_id: str) -> Iterator[None]:
        """
        Context manager to track a thread's work session.
        A thread enters a work session to check for and perform work.
        """
        with self.lock:
            self.active_threads.add(thread_id)

        yield

        with self.lock:
            assert thread_id not in self.active_threads, "You must call report_work before exiting the session"

    def report_work(self, thread_id: str, did_work: bool) -> None:
        with self.lock:
            assert thread_id in self.active_threads, "Thread must be in work session to report work"
            self.active_threads.discard(thread_id)

            if did_work:
                # reactivate all, because the work we did might have unblocked other threads
                self.idle_threads.clear()

            else:
                # This thread is permanently idle
                self.idle_threads.add(thread_id)

    @property
    def all_idle(self) -> bool:
        with self.lock:
            return len(self.idle_threads) == self.total_threads
