"""
In-memory storage implementation for pyergon.

Design Pattern: Adapter Pattern
InMemoryExecutionLog adapts in-memory dictionaries to ExecutionLog interface.

RUST COMPLIANCE: Matches Rust ergon InMemoryExecutionLog exactly.

From Dave Cheney:
"Make the zero value useful" - Instance is immediately usable after __init__.
"""

from __future__ import annotations

import asyncio
from datetime import datetime, timedelta

from uuid_extensions import uuid7

from pyergon.models import (
    Invocation,
    InvocationStatus,
    RetryPolicy,
    ScheduledFlow,
    TaskStatus,
    TimerInfo,
)
from pyergon.storage.base import ExecutionLog, StorageError


class InMemoryExecutionLog(ExecutionLog):
    """
    In-memory storage for testing.

    Design: Liskov Substitution (SOLID)
    Can be substituted for SqliteExecutionLog without changing client code.

    Usage:
        storage = InMemoryExecutionLog()  # Ready to use!
        await storage.log_invocation_start(...)
    """

    def __init__(self):
        """
        Initialize in-memory storage with notification support.

        **Rust Reference**: `src/storage/memory.rs` lines 49-59

        Creates notification events for:
        - work_notify: Wake workers when work becomes available
        - timer_notify: Wake timer processors when timer state changes
        - status_notify: Notify flow status changes (for future use)
        """
        # Storage: {(flow_id, step): Invocation}
        self._invocations: dict[tuple[str, int], Invocation] = {}

        # Storage: {task_id: ScheduledFlow}
        self._scheduled_flows: dict[str, ScheduledFlow] = {}

        # Storage: {(flow_id, step, suspension_key): bytes}
        # For storing suspension results (timer/signal completions)
        self._suspension_results: dict[tuple[str, int, str], bytes] = {}

        # Lock for thread-safety
        self._lock = asyncio.Lock()

        # Notification events (implements WorkNotificationSource and TimerNotificationSource)
        # From Rust: work_notify: Arc<Notify>, timer_notify: Arc<Notify>
        self._work_notify = asyncio.Event()
        self._timer_notify = asyncio.Event()
        self._status_notify = asyncio.Event()

    def __repr__(self) -> str:
        """Return string representation of storage instance."""
        return "InMemoryExecutionLog"

    # ========================================================================
    # Invocation Operations
    # ========================================================================

    async def log_invocation_start(
        self,
        flow_id: str,
        step: int,
        class_name: str,
        method_name: str,
        parameters: bytes,
        params_hash: int,
        delay: int | None = None,
        retry_policy: RetryPolicy | None = None,
    ) -> Invocation:
        """
        Record the start of a step execution.

        **Rust Reference**: `src/storage/memory.rs` lines 95-134

        This method checks if an invocation already exists:
        - If COMPLETE: skip, don't overwrite cached result
        - If WAITING_FOR_SIGNAL: skip, don't overwrite signal state
        - Otherwise: overwrite with new invocation
        """
        async with self._lock:
            key = (flow_id, step)

            # Check if invocation already exists and skip if Complete,
            # WaitingForSignal, or WaitingForTimer
            # From Rust lines 123-130
            if key in self._invocations:
                existing = self._invocations[key]
                if existing.status == InvocationStatus.COMPLETE:
                    # Don't overwrite cached results during replay
                    return existing
                if existing.status == InvocationStatus.WAITING_FOR_SIGNAL:
                    # Don't overwrite signal state during replay
                    return existing
                if existing.status == InvocationStatus.WAITING_FOR_TIMER:
                    # Don't overwrite timer state during replay
                    return existing

            # Create new invocation (or overwrite non-Complete/non-WaitingForSignal)
            now = datetime.now()
            inv_id = str(uuid7())

            invocation = Invocation(
                id=inv_id,
                flow_id=flow_id,
                step=step,
                timestamp=now,
                class_name=class_name,
                method_name=method_name,
                status=InvocationStatus.PENDING,
                attempts=0,
                parameters=parameters,
                params_hash=params_hash,
                return_value=None,
                delay=delay,
                retry_policy=retry_policy,
                is_retryable=None,
                timer_fire_at=None,
                timer_name=None,
                updated_at=now,
            )

            self._invocations[key] = invocation
            return invocation

    async def log_invocation_completion(
        self, flow_id: str, step: int, return_value: bytes, is_retryable: bool | None = None
    ) -> Invocation:
        """Record the completion of a step execution."""
        async with self._lock:
            key = (flow_id, step)

            if key not in self._invocations:
                raise StorageError(f"Invocation not found: flow_id={flow_id}, step={step}")

            old = self._invocations[key]

            # Create updated invocation
            updated = Invocation(
                id=old.id,
                flow_id=old.flow_id,
                step=old.step,
                timestamp=old.timestamp,
                class_name=old.class_name,
                method_name=old.method_name,
                status=InvocationStatus.COMPLETE,
                attempts=old.attempts + 1,
                parameters=old.parameters,
                params_hash=old.params_hash,
                return_value=return_value,
                delay=old.delay,
                retry_policy=old.retry_policy,
                is_retryable=is_retryable,
                timer_fire_at=old.timer_fire_at,
                timer_name=old.timer_name,
                updated_at=datetime.now(),
            )

            self._invocations[key] = updated
            return updated

    async def get_invocation(self, flow_id: str, step: int) -> Invocation | None:
        """Retrieve a specific step invocation."""
        async with self._lock:
            return self._invocations.get((flow_id, step))

    async def get_invocations_for_flow(self, flow_id: str) -> list[Invocation]:
        """Get all invocations for a flow."""
        async with self._lock:
            return [inv for (fid, _), inv in self._invocations.items() if fid == flow_id]

    async def get_incomplete_flows(self) -> list[Invocation]:
        """Get all flows with incomplete invocations."""
        async with self._lock:
            # Return invocations that are not complete
            return [inv for inv in self._invocations.values() if not inv.is_complete]

    async def has_non_retryable_error(self, flow_id: str) -> bool:
        """Check if flow has non-retryable error."""
        async with self._lock:
            for (fid, _), inv in self._invocations.items():
                if fid == flow_id and not inv.is_retryable:
                    return True
            return False

    async def update_is_retryable(self, flow_id: str, step: int, is_retryable: bool) -> None:
        """Update is_retryable field."""
        async with self._lock:
            key = (flow_id, step)
            if key in self._invocations:
                old = self._invocations[key]
                updated = Invocation(
                    id=old.id,
                    flow_id=old.flow_id,
                    step=old.step,
                    timestamp=old.timestamp,
                    class_name=old.class_name,
                    method_name=old.method_name,
                    status=old.status,
                    attempts=old.attempts,
                    parameters=old.parameters,
                    params_hash=old.params_hash,
                    return_value=old.return_value,
                    delay=old.delay,
                    retry_policy=old.retry_policy,
                    is_retryable=is_retryable,
                    timer_fire_at=old.timer_fire_at,
                    timer_name=old.timer_name,
                    updated_at=datetime.now(),
                )
                self._invocations[key] = updated

    # ========================================================================
    # Queue Operations
    # ========================================================================

    async def enqueue_flow(self, flow: ScheduledFlow) -> str:
        """
        Add a flow to the distributed work queue.

        **Rust Reference**: `src/storage/memory.rs` lines 259-311

        Accepts a complete ScheduledFlow object which includes parent_metadata
        for child flow invocation (Level 3 API).
        """
        async with self._lock:
            # Store the flow object directly
            # The flow object already contains all necessary fields including parent_metadata
            self._scheduled_flows[flow.task_id] = flow

            # Wake up one waiting worker if this is an immediate (non-delayed) flow
            # From Rust lines 307-309
            if flow.scheduled_for is None:
                self._work_notify.set()
                self._work_notify.clear()  # Reset for next notification

            return flow.task_id

    async def dequeue_flow(self, worker_id: str) -> ScheduledFlow | None:
        """
        Claim and retrieve a pending flow from queue.

        **Rust Reference**: `src/storage/memory.rs` lines 314-361

        Only dequeues flows that are:
        1. Status == PENDING
        2. scheduled_for is None OR scheduled_for <= now (ready to execute)
        """
        async with self._lock:
            now = datetime.now()

            # Find first pending flow that's ready to execute
            for task_id, flow in self._scheduled_flows.items():
                if flow.status == TaskStatus.PENDING:
                    # Check if flow is scheduled for the future
                    if flow.scheduled_for is not None and flow.scheduled_for > now:
                        # Not ready yet - skip this flow
                        continue

                    # Flow is ready - claim it
                    updated = ScheduledFlow(
                        task_id=flow.task_id,
                        flow_id=flow.flow_id,
                        flow_type=flow.flow_type,
                        flow_data=flow.flow_data,
                        status=TaskStatus.RUNNING,
                        locked_by=worker_id,
                        retry_count=flow.retry_count,
                        created_at=flow.created_at,
                        updated_at=now,
                        error_message=flow.error_message,
                        claimed_at=now,
                        completed_at=flow.completed_at,
                        scheduled_for=flow.scheduled_for,
                        parent_metadata=flow.parent_metadata,
                        retry_policy=flow.retry_policy,
                    )
                    self._scheduled_flows[task_id] = updated
                    return updated

            return None

    async def complete_flow(
        self, task_id: str, status: TaskStatus, error_message: str | None = None
    ) -> None:
        """
        Update flow task status.

        **Rust Reference**: `src/storage/memory.rs` lines 346-358

        Despite the name "complete_flow", this method updates the status to ANY value,
        including SUSPENDED (not just terminal states). The name comes from the most
        common use case (marking flows COMPLETE/FAILED), but it's actually a general
        status update method.

        From Rust line 348: entry.status = status  (no validation)

        Args:
            task_id: The task identifier
            status: The final status (COMPLETE, FAILED, or SUSPENDED)
            error_message: Optional error message if status is FAILED
        """
        async with self._lock:
            if task_id not in self._scheduled_flows:
                raise StorageError(f"Flow not found: task_id={task_id}")

            flow = self._scheduled_flows[task_id]
            now = datetime.now()

            # Set completed_at only for terminal statuses
            completed_at = now if status.is_terminal else flow.completed_at

            updated = ScheduledFlow(
                task_id=flow.task_id,
                flow_id=flow.flow_id,
                flow_type=flow.flow_type,
                flow_data=flow.flow_data,
                status=status,  # Accept ANY status (including SUSPENDED)
                locked_by=flow.locked_by,
                retry_count=flow.retry_count,
                created_at=flow.created_at,
                updated_at=now,
                error_message=error_message if error_message is not None else flow.error_message,
                claimed_at=flow.claimed_at,
                completed_at=completed_at,
                scheduled_for=flow.scheduled_for,
                parent_metadata=flow.parent_metadata,
                retry_policy=flow.retry_policy,
            )

            self._scheduled_flows[task_id] = updated

            # Notify any waiters that a flow status changed
            # From Rust lines 354-356
            self._status_notify.set()
            self._status_notify.clear()  # Reset for next notification

    async def retry_flow(self, task_id: str, error_message: str, delay: timedelta) -> None:
        """
        Reschedule a failed flow for retry after a delay.

        **Rust Reference**: `src/storage/memory.rs` lines 364-401

        This method:
        1. Increments retry_count
        2. Sets error_message
        3. Sets status back to PENDING
        4. Clears locked_by
        5. Sets scheduled_for (current time + delay)
        6. Re-enqueues the flow
        """
        async with self._lock:
            if task_id not in self._scheduled_flows:
                raise StorageError(f"Flow not found: task_id={task_id}")

            flow = self._scheduled_flows[task_id]
            now = datetime.now()
            scheduled_for = now + delay

            # Update flow for retry
            updated = ScheduledFlow(
                task_id=flow.task_id,
                flow_id=flow.flow_id,
                flow_type=flow.flow_type,
                flow_data=flow.flow_data,
                status=TaskStatus.PENDING,  # Back to pending
                locked_by=None,  # Clear lock
                retry_count=flow.retry_count + 1,  # Increment
                created_at=flow.created_at,
                updated_at=now,
                error_message=error_message,  # Set error
                claimed_at=None,
                completed_at=None,
                scheduled_for=scheduled_for,  # Schedule for future
                parent_metadata=flow.parent_metadata,
                retry_policy=flow.retry_policy,
            )

            self._scheduled_flows[task_id] = updated

            # Notify workers that new work is available (or will be available soon)
            # From Rust lines 388-395
            self._work_notify.set()
            self._work_notify.clear()  # Reset for next notification

    async def get_scheduled_flow(self, task_id: str) -> ScheduledFlow | None:
        """Retrieve a scheduled flow by task ID."""
        async with self._lock:
            return self._scheduled_flows.get(task_id)

    # ========================================================================
    # Timer Operations
    # ========================================================================

    async def log_timer(
        self, flow_id: str, step: int, timer_fire_at: datetime, timer_name: str | None = None
    ) -> None:
        """Schedule a durable timer."""
        async with self._lock:
            key = (flow_id, step)
            if key in self._invocations:
                old = self._invocations[key]
                updated = Invocation(
                    id=old.id,
                    flow_id=old.flow_id,
                    step=old.step,
                    timestamp=old.timestamp,
                    class_name=old.class_name,
                    method_name=old.method_name,
                    status=InvocationStatus.WAITING_FOR_TIMER,
                    attempts=old.attempts,
                    parameters=old.parameters,
                    params_hash=old.params_hash,
                    return_value=old.return_value,
                    delay=old.delay,
                    retry_policy=old.retry_policy,
                    is_retryable=old.is_retryable,
                    timer_fire_at=timer_fire_at,
                    timer_name=timer_name,
                    updated_at=datetime.now(),
                )
                self._invocations[key] = updated

                # Notify timer processor that a new timer was scheduled
                # From Rust lines 481-483
                self._timer_notify.set()
                self._timer_notify.clear()  # Reset for next notification

    async def get_expired_timers(self, now: datetime) -> list[TimerInfo]:
        """
        Get all timers that have expired.

        **Rust Reference**: storage/mod.rs lines 66-71 (TimerInfo struct)
        """

        async with self._lock:
            expired = []
            for (flow_id, step), inv in self._invocations.items():
                if (
                    inv.status == InvocationStatus.WAITING_FOR_TIMER
                    and inv.timer_fire_at is not None
                    and inv.timer_fire_at <= now
                ):
                    expired.append(
                        TimerInfo(
                            flow_id=flow_id,
                            step=step,
                            fire_at=inv.timer_fire_at,
                            timer_name=inv.timer_name,
                        )
                    )
            return expired

    async def claim_timer(
        self,
        flow_id: str,
        step: int,
        worker_id: str,  # noqa: ARG002 - Required by protocol signature
    ) -> bool:
        """Claim an expired timer (optimistic concurrency)."""
        async with self._lock:
            key = (flow_id, step)
            if key not in self._invocations:
                return False

            inv = self._invocations[key]
            if inv.status != InvocationStatus.WAITING_FOR_TIMER:
                return False

            # Claim it by changing status to PENDING
            updated = Invocation(
                id=inv.id,
                flow_id=inv.flow_id,
                step=inv.step,
                timestamp=inv.timestamp,
                class_name=inv.class_name,
                method_name=inv.method_name,
                status=InvocationStatus.PENDING,
                attempts=inv.attempts,
                parameters=inv.parameters,
                params_hash=inv.params_hash,
                return_value=inv.return_value,
                delay=inv.delay,
                retry_policy=inv.retry_policy,
                is_retryable=inv.is_retryable,
                timer_fire_at=inv.timer_fire_at,
                timer_name=inv.timer_name,
                updated_at=datetime.now(),
            )
            self._invocations[key] = updated

            # Notify timer processor that timer was claimed (may need to recalculate next wake time)
            # From Rust lines 495-497
            self._timer_notify.set()
            self._timer_notify.clear()  # Reset for next notification

            return True

    # ========================================================================
    # Signal Operations - External event coordination
    # ========================================================================

    async def store_suspension_result(
        self, flow_id: str, step: int, suspension_key: str, result: bytes
    ) -> None:
        """Store suspension result data (timer or signal completion)."""
        async with self._lock:
            key = (flow_id, step, suspension_key)
            self._suspension_results[key] = result

    async def get_suspension_result(
        self, flow_id: str, step: int, suspension_key: str
    ) -> bytes | None:
        """Get suspension result data if available."""
        async with self._lock:
            key = (flow_id, step, suspension_key)
            return self._suspension_results.get(key)

    async def remove_suspension_result(self, flow_id: str, step: int, suspension_key: str) -> None:
        """Remove suspension result data after consuming it."""
        async with self._lock:
            key = (flow_id, step, suspension_key)
            self._suspension_results.pop(key, None)

    async def log_signal(self, flow_id: str, step: int, signal_name: str) -> None:
        """
        Mark an invocation as waiting for an external signal.

        CRITICAL: This must be called BEFORE scheduling the child flow to prevent
        race conditions where the child completes before the parent is marked as waiting.
        """
        async with self._lock:
            key = (flow_id, step)
            if key not in self._invocations:
                raise StorageError(f"Invocation not found: flow_id={flow_id}, step={step}")

            inv = self._invocations[key]

            # Update invocation status to WAITING_FOR_SIGNAL
            # From Rust: storage/memory.rs lines 502-530
            updated = Invocation(
                id=inv.id,
                flow_id=inv.flow_id,
                step=inv.step,
                timestamp=inv.timestamp,
                class_name=inv.class_name,
                method_name=inv.method_name,
                status=InvocationStatus.WAITING_FOR_SIGNAL,
                attempts=inv.attempts,
                parameters=inv.parameters,
                params_hash=inv.params_hash,
                return_value=inv.return_value,
                delay=inv.delay,
                retry_policy=inv.retry_policy,
                is_retryable=inv.is_retryable,
                timer_fire_at=inv.timer_fire_at,
                timer_name=signal_name,  # Store signal name in timer_name field
                updated_at=datetime.now(),
            )
            self._invocations[key] = updated

    # ========================================================================
    # Flow Resume Operations
    # ========================================================================

    async def resume_flow(self, flow_id: str) -> bool:
        """
        Resume a suspended flow by changing status SUSPENDED → PENDING.

        **Rust Reference**: `src/storage/memory.rs` lines 626-662

        This method atomically:
        1. Checks if flow is SUSPENDED
        2. Changes status to PENDING
        3. Re-enqueues to work queue (adds back to scheduled_flows list)

        Returns:
            True if flow was resumed, False if not suspended
        """
        async with self._lock:
            # Find the flow's task by flow_id
            task_to_resume = None
            task_id_to_resume = None
            for tid, task in self._scheduled_flows.items():
                if task.flow_id == flow_id and task.status == TaskStatus.SUSPENDED:
                    task_to_resume = task
                    task_id_to_resume = tid
                    break

            if task_to_resume is None:
                # Flow not found or not suspended
                return False

            # Update status to PENDING
            # From Rust lines 635-638
            updated_task = ScheduledFlow(
                task_id=task_to_resume.task_id,
                flow_id=task_to_resume.flow_id,
                flow_type=task_to_resume.flow_type,
                flow_data=task_to_resume.flow_data,
                status=TaskStatus.PENDING,  # Change from SUSPENDED to PENDING
                locked_by=None,  # Clear lock
                scheduled_for=task_to_resume.scheduled_for,
                created_at=task_to_resume.created_at,
                updated_at=datetime.now(),
                retry_count=task_to_resume.retry_count,
                retry_policy=task_to_resume.retry_policy,
                parent_metadata=task_to_resume.parent_metadata,
            )

            # Replace in scheduled_flows dict
            self._scheduled_flows[task_id_to_resume] = updated_task

            # Wake up one waiting worker since we just made a flow available
            # From Rust lines 654-656
            self._work_notify.set()
            self._work_notify.clear()  # Reset for next notification

            return True

    async def get_next_timer_fire_time(self) -> datetime | None:
        """
        Get the earliest timer fire time across all flows.

        **Rust Reference**: `src/storage/memory.rs` lines 456-476

        Returns:
            datetime of next timer, or None if no timers pending
        """
        async with self._lock:
            next_fire_time = None

            # Iterate all invocations looking for timers
            # From Rust lines 459-471
            for inv in self._invocations.values():
                if inv.status == InvocationStatus.WAITING_FOR_TIMER:
                    if inv.timer_fire_at is not None:
                        if next_fire_time is None:
                            next_fire_time = inv.timer_fire_at
                        elif inv.timer_fire_at < next_fire_time:
                            next_fire_time = inv.timer_fire_at

            return next_fire_time

    # ========================================================================
    # Utility Operations
    # ========================================================================

    async def reset(self) -> None:
        """Clear all data (for testing)."""
        async with self._lock:
            self._invocations.clear()
            self._scheduled_flows.clear()
            self._suspension_results.clear()

    async def close(self) -> None:
        """Close storage (no-op for in-memory)."""
        pass

    # ========================================================================
    # Notification Protocol Methods
    # ========================================================================

    def work_notify(self) -> asyncio.Event:
        """
        Return event for work notifications (WorkNotificationSource protocol).

        **Rust Reference**: `src/storage/memory.rs` lines 1009-1012

        Workers wait on this event to be notified when work becomes available,
        instead of polling with sleep().

        Returns:
            asyncio.Event that is set when work is enqueued or resumed
        """
        return self._work_notify

    def timer_notify(self) -> asyncio.Event:
        """
        Return event for timer notifications (TimerNotificationSource protocol).

        **Rust Reference**: `src/storage/memory.rs` lines 1014-1017

        Timer processors wait on this event to be notified when timer state changes,
        instead of polling every N seconds.

        Returns:
            asyncio.Event that is set when timers are scheduled or claimed
        """
        return self._timer_notify

    def status_notify(self) -> asyncio.Event:
        """
        Return event for flow status change notifications.

        **Rust Reference**: `src/storage/memory.rs` lines 69-71

        Callers can use this to wait for flow status changes (completion, failure, etc.)
        instead of polling. The notification is triggered whenever any flow status changes.

        Returns:
            asyncio.Event that is set when any flow status changes
        """
        return self._status_notify
