import json
import logging
import time
from dataclasses import dataclass
from typing import Any, Union
from uuid import UUID

from pottery import Redlock

from config.redis_client import get_redis_client

logger = logging.getLogger(__name__)

_QUEUED_KEY_TEMPLATE = "agent-event-processing:queued:{agent_id}"
_DEFAULT_QUEUE_TTL_SECONDS = 3600
_HEARTBEAT_KEY_TEMPLATE = "agent-event-processing:heartbeat:{agent_id}"
_DEFAULT_HEARTBEAT_TTL_SECONDS = 600
_STOP_REQUESTED_KEY_TEMPLATE = "agent-event-processing:stop-requested:{agent_id}"
_DEFAULT_STOP_REQUESTED_TTL_SECONDS = 3600
_HUMAN_INBOUND_GENERATION_KEY_TEMPLATE = "agent-event-processing:human-inbound-generation:{agent_id}"
_HUMAN_INBOUND_CONSUMED_GENERATION_KEY_TEMPLATE = "agent-event-processing:human-inbound-consumed-generation:{agent_id}"
_QUEUED_AGENT_SET_KEY = "agent-event-processing:index:queued"
_HEARTBEAT_AGENT_SET_KEY = "agent-event-processing:index:heartbeat"
_LOCKED_AGENT_SET_KEY = "agent-event-processing:index:locked"
_PENDING_SET_KEY = "agent-event-processing:pending"
_PENDING_DRAIN_SCHEDULE_KEY = "agent-event-processing:pending:drain:schedule"
_DEFAULT_PENDING_SET_TTL_SECONDS = 3600
_DEFAULT_PENDING_DRAIN_SCHEDULE_TTL_SECONDS = 60


@dataclass(frozen=True)
class PendingDrainSettings:
    pending_set_ttl_seconds: int
    pending_drain_delay_seconds: int
    pending_drain_limit: int
    pending_drain_schedule_ttl_seconds: int


def _queued_key(agent_id: Union[str, UUID]) -> str:
    return _QUEUED_KEY_TEMPLATE.format(agent_id=agent_id)


def _heartbeat_key(agent_id: Union[str, UUID]) -> str:
    return _HEARTBEAT_KEY_TEMPLATE.format(agent_id=agent_id)


def _stop_requested_key(agent_id: Union[str, UUID]) -> str:
    return _STOP_REQUESTED_KEY_TEMPLATE.format(agent_id=agent_id)


def _human_inbound_generation_key(agent_id: Union[str, UUID]) -> str:
    return _HUMAN_INBOUND_GENERATION_KEY_TEMPLATE.format(agent_id=agent_id)


def _human_inbound_consumed_generation_key(agent_id: Union[str, UUID]) -> str:
    return _HUMAN_INBOUND_CONSUMED_GENERATION_KEY_TEMPLATE.format(agent_id=agent_id)


def processing_lock_storage_keys(agent_id: Union[str, UUID]) -> tuple[str, str]:
    normalized_agent_id = str(agent_id)
    prefix = getattr(Redlock, "_KEY_PREFIX", "redlock")
    return (
        f"{prefix}:agent-event-processing:{normalized_agent_id}",
        f"agent-event-processing:{normalized_agent_id}",
    )


def _smembers_as_strings(redis_client, key: str) -> list[str]:
    values = getattr(redis_client, "smembers", lambda _key: set())(key)
    normalized: list[str] = []
    for value in values:
        if isinstance(value, (bytes, bytearray)):
            normalized.append(value.decode("utf-8", "ignore"))
        else:
            normalized.append(str(value))
    return normalized


def get_processing_queued_agent_ids(*, client=None) -> list[str]:
    try:
        redis_client = client or get_redis_client()
        return _smembers_as_strings(redis_client, _QUEUED_AGENT_SET_KEY)
    except Exception:
        logger.exception("Failed to list queued processing agents")
        return []


def get_processing_heartbeat_agent_ids(*, client=None) -> list[str]:
    try:
        redis_client = client or get_redis_client()
        return _smembers_as_strings(redis_client, _HEARTBEAT_AGENT_SET_KEY)
    except Exception:
        logger.exception("Failed to list heartbeat processing agents")
        return []


def get_processing_locked_agent_ids(*, client=None) -> list[str]:
    try:
        redis_client = client or get_redis_client()
        return _smembers_as_strings(redis_client, _LOCKED_AGENT_SET_KEY)
    except Exception:
        logger.exception("Failed to list locked processing agents")
        return []


def set_processing_queued_flag(
    agent_id: Union[str, UUID],
    *,
    ttl: int = _DEFAULT_QUEUE_TTL_SECONDS,
    client=None,
) -> None:
    """Mark the agent as having queued processing work."""
    try:
        redis_client = client or get_redis_client()
        key = _queued_key(agent_id)
        pipeline = getattr(redis_client, "pipeline", None)
        if callable(pipeline):
            pipe = pipeline()
            pipe.set(key, "1")
            if ttl > 0:
                pipe.expire(key, ttl)
            pipe.sadd(_QUEUED_AGENT_SET_KEY, str(agent_id))
            pipe.execute()
            return

        redis_client.set(key, "1")
        if ttl > 0:
            redis_client.expire(key, ttl)
        redis_client.sadd(_QUEUED_AGENT_SET_KEY, str(agent_id))
    except Exception:
        logger.exception("Failed to set processing queued flag for agent %s", agent_id)


def clear_processing_queued_flag(agent_id: Union[str, UUID], *, client=None) -> None:
    """Clear the queued processing flag for the agent."""
    try:
        redis_client = client or get_redis_client()
        pipeline = getattr(redis_client, "pipeline", None)
        if callable(pipeline):
            pipe = pipeline()
            pipe.delete(_queued_key(agent_id))
            pipe.srem(_QUEUED_AGENT_SET_KEY, str(agent_id))
            pipe.execute()
            return

        redis_client.delete(_queued_key(agent_id))
        redis_client.srem(_QUEUED_AGENT_SET_KEY, str(agent_id))
    except Exception:
        logger.exception("Failed to clear processing queued flag for agent %s", agent_id)


def _coerce_generation(value: Any) -> int:
    if value is None:
        return 0
    if isinstance(value, (bytes, bytearray)):
        value = value.decode("utf-8", "ignore")
    try:
        return max(0, int(value))
    except (TypeError, ValueError):
        return 0


def bump_human_inbound_generation(
    agent_id: Union[str, UUID],
    *,
    ttl: int = _DEFAULT_QUEUE_TTL_SECONDS,
    client=None,
) -> int:
    """Increment and return the generation for human-authored inbound messages."""
    try:
        redis_client = client or get_redis_client()
        key = _human_inbound_generation_key(agent_id)
        consumed_key = _human_inbound_consumed_generation_key(agent_id)
        generation = int(redis_client.incr(key))
        consumed = get_consumed_human_inbound_generation(agent_id, client=redis_client)
        if generation <= consumed:
            generation = int(redis_client.incr(key, amount=consumed - generation + 1))
        if ttl > 0:
            redis_client.expire(key, ttl)
            redis_client.expire(consumed_key, ttl)
        return generation
    except Exception:
        logger.exception("Failed to bump human inbound generation for agent %s", agent_id)
        return 0


def get_human_inbound_generation(agent_id: Union[str, UUID], *, client=None) -> int:
    """Return the latest human-inbound generation for the agent."""
    try:
        redis_client = client or get_redis_client()
        return _coerce_generation(redis_client.get(_human_inbound_generation_key(agent_id)))
    except Exception:
        logger.exception("Failed to read human inbound generation for agent %s", agent_id)
        return 0


def get_consumed_human_inbound_generation(agent_id: Union[str, UUID], *, client=None) -> int:
    """Return the latest human-inbound generation accepted by the orchestrator."""
    try:
        redis_client = client or get_redis_client()
        return _coerce_generation(redis_client.get(_human_inbound_consumed_generation_key(agent_id)))
    except Exception:
        logger.exception("Failed to read consumed human inbound generation for agent %s", agent_id)
        return 0


def mark_human_inbound_generation_consumed(
    agent_id: Union[str, UUID],
    generation: int | str | None,
    *,
    ttl: int = _DEFAULT_QUEUE_TTL_SECONDS,
    client=None,
) -> int:
    """Record that the orchestrator accepted a prompt at this generation."""
    parsed_generation = _coerce_generation(generation)
    if parsed_generation <= 0:
        return 0

    try:
        redis_client = client or get_redis_client()
        key = _human_inbound_consumed_generation_key(agent_id)
        current_key = _human_inbound_generation_key(agent_id)
        previous = _coerce_generation(redis_client.get(key))
        next_generation = max(previous, parsed_generation)
        redis_client.set(key, str(next_generation))
        if ttl > 0:
            redis_client.expire(key, ttl)
            redis_client.expire(current_key, ttl)
        return next_generation
    except Exception:
        logger.exception("Failed to mark human inbound generation consumed for agent %s", agent_id)
        return parsed_generation


def is_human_inbound_generation_consumed(
    agent_id: Union[str, UUID],
    generation: int | str | None,
    *,
    client=None,
) -> bool:
    """Return True when a queued human-inbound task is redundant."""
    parsed_generation = _coerce_generation(generation)
    if parsed_generation <= 0:
        return False

    current = get_human_inbound_generation(agent_id, client=client)
    consumed = get_consumed_human_inbound_generation(agent_id, client=client)
    return consumed >= parsed_generation and current <= consumed


def clear_processing_work_state(agent_id: Union[str, UUID], client=None) -> None:
    """Clear queued and pending processing state for a single agent."""
    redis_client = client
    if redis_client is None:
        try:
            redis_client = get_redis_client()
        except Exception:
            logger.exception("Failed to acquire Redis client while clearing processing state for agent %s", agent_id)
            return

    try:
        pipeline = getattr(redis_client, "pipeline", None)
        if callable(pipeline):
            pipe = pipeline()
            pipe.delete(_queued_key(agent_id))
            pipe.srem(_QUEUED_AGENT_SET_KEY, str(agent_id))
            pipe.srem(_PENDING_SET_KEY, str(agent_id))
            pipe.srem(_LOCKED_AGENT_SET_KEY, str(agent_id))
            pipe.srem(_HEARTBEAT_AGENT_SET_KEY, str(agent_id))
            pipe.execute()
            return

        redis_client.delete(_queued_key(agent_id))
        redis_client.srem(_QUEUED_AGENT_SET_KEY, str(agent_id))
    except Exception:
        logger.exception("Failed to clear queued processing state for agent %s", agent_id)

    try:
        redis_client.srem(_PENDING_SET_KEY, str(agent_id))
        redis_client.srem(_LOCKED_AGENT_SET_KEY, str(agent_id))
        redis_client.srem(_HEARTBEAT_AGENT_SET_KEY, str(agent_id))
    except Exception:
        logger.exception("Failed to clear pending processing state for agent %s", agent_id)


def set_processing_stop_requested(
    agent_id: Union[str, UUID],
    *,
    ttl: int = _DEFAULT_STOP_REQUESTED_TTL_SECONDS,
    client=None,
) -> None:
    """Record a graceful stop request for the agent."""
    try:
        redis_client = client or get_redis_client()
        if ttl > 0:
            redis_client.set(_stop_requested_key(agent_id), "1", ex=ttl)
        else:
            redis_client.set(_stop_requested_key(agent_id), "1")
    except Exception:
        logger.exception("Failed to set stop request for agent %s", agent_id)


def is_processing_stop_requested(agent_id: Union[str, UUID], client=None) -> bool:
    """Check whether a graceful stop has been requested for the agent."""
    try:
        redis_client = client or get_redis_client()
        return bool(redis_client.exists(_stop_requested_key(agent_id)))
    except Exception:
        logger.exception("Failed to check stop request for agent %s", agent_id)
        return False


def clear_processing_stop_requested(agent_id: Union[str, UUID], client=None) -> None:
    """Clear the graceful stop request for the agent."""
    try:
        redis_client = client or get_redis_client()
        redis_client.delete(_stop_requested_key(agent_id))
    except Exception:
        logger.exception("Failed to clear stop request for agent %s", agent_id)


def is_processing_queued(agent_id: Union[str, UUID], client=None) -> bool:
    """Check whether the agent currently has queued processing work."""
    try:
        redis_client = client or get_redis_client()
        return bool(redis_client.exists(_queued_key(agent_id)))
    except Exception:
        logger.exception("Failed to check processing queued flag for agent %s", agent_id)
        return False


def set_processing_heartbeat(
    agent_id: Union[str, UUID],
    *,
    ttl: int = _DEFAULT_HEARTBEAT_TTL_SECONDS,
    run_id: str | None = None,
    worker_pid: int | None = None,
    stage: str | None = None,
    started_at: float | None = None,
    client=None,
) -> None:
    """Record a processing heartbeat for the agent."""
    if ttl <= 0:
        return
    now = time.time()
    payload = {
        "agent_id": str(agent_id),
        "run_id": run_id,
        "worker_pid": worker_pid,
        "stage": stage,
        "started_at": started_at if started_at is not None else now,
        "last_seen": now,
    }
    try:
        redis_client = client or get_redis_client()
        pipeline = getattr(redis_client, "pipeline", None)
        if callable(pipeline):
            pipe = pipeline()
            pipe.set(_heartbeat_key(agent_id), json.dumps(payload), ex=ttl)
            pipe.sadd(_HEARTBEAT_AGENT_SET_KEY, str(agent_id))
            pipe.execute()
            return

        redis_client.set(_heartbeat_key(agent_id), json.dumps(payload), ex=ttl)
        redis_client.sadd(_HEARTBEAT_AGENT_SET_KEY, str(agent_id))
    except Exception:
        logger.exception("Failed to set processing heartbeat for agent %s", agent_id)


def clear_processing_heartbeat(agent_id: Union[str, UUID], client=None) -> None:
    """Clear the processing heartbeat for the agent."""
    try:
        redis_client = client or get_redis_client()
        pipeline = getattr(redis_client, "pipeline", None)
        if callable(pipeline):
            pipe = pipeline()
            pipe.delete(_heartbeat_key(agent_id))
            pipe.srem(_HEARTBEAT_AGENT_SET_KEY, str(agent_id))
            pipe.execute()
            return

        redis_client.delete(_heartbeat_key(agent_id))
        redis_client.srem(_HEARTBEAT_AGENT_SET_KEY, str(agent_id))
    except Exception:
        logger.exception("Failed to clear processing heartbeat for agent %s", agent_id)


def get_processing_heartbeat(agent_id: Union[str, UUID], client=None) -> dict | None:
    """Fetch the last processing heartbeat payload for the agent."""
    try:
        redis_client = client or get_redis_client()
        raw = redis_client.get(_heartbeat_key(agent_id))
    except Exception:
        logger.exception("Failed to read processing heartbeat for agent %s", agent_id)
        return None
    if not raw:
        return None
    if isinstance(raw, (bytes, bytearray)):
        raw = raw.decode("utf-8", "ignore")
    try:
        return json.loads(raw)
    except (TypeError, ValueError):
        logger.exception("Failed to parse processing heartbeat for agent %s", agent_id)
        return None


def pending_set_key() -> str:
    return _PENDING_SET_KEY


def pending_drain_schedule_key() -> str:
    return _PENDING_DRAIN_SCHEDULE_KEY


def mark_processing_lock_active(agent_id: Union[str, UUID], *, client=None) -> None:
    try:
        redis_client = client or get_redis_client()
        redis_client.sadd(_LOCKED_AGENT_SET_KEY, str(agent_id))
    except Exception:
        logger.exception("Failed to mark processing lock active for agent %s", agent_id)


def clear_processing_lock_active(agent_id: Union[str, UUID], *, client=None) -> None:
    try:
        redis_client = client or get_redis_client()
        redis_client.srem(_LOCKED_AGENT_SET_KEY, str(agent_id))
    except Exception:
        logger.exception("Failed to clear processing lock active for agent %s", agent_id)


def _coerce_positive_int(value: Any, default: int) -> int:
    try:
        parsed = int(value)
    except (TypeError, ValueError):
        parsed = default
    return max(1, parsed)


def get_pending_drain_settings(settings_obj=None) -> PendingDrainSettings:
    if settings_obj is None:
        from django.conf import settings as django_settings

        settings_obj = django_settings

    pending_set_ttl_seconds = _coerce_positive_int(
        getattr(settings_obj, "AGENT_EVENT_PROCESSING_PENDING_SET_TTL_SECONDS", None),
        _DEFAULT_PENDING_SET_TTL_SECONDS,
    )
    pending_drain_delay_seconds = _coerce_positive_int(
        getattr(settings_obj, "AGENT_EVENT_PROCESSING_PENDING_DRAIN_DELAY_SECONDS", None),
        5,
    )
    pending_drain_limit = _coerce_positive_int(
        getattr(settings_obj, "AGENT_EVENT_PROCESSING_PENDING_DRAIN_LIMIT", None),
        50,
    )
    schedule_default = max(30, pending_drain_delay_seconds * 6)
    pending_drain_schedule_ttl_seconds = _coerce_positive_int(
        getattr(settings_obj, "AGENT_EVENT_PROCESSING_PENDING_DRAIN_SCHEDULE_TTL_SECONDS", None),
        schedule_default,
    )
    return PendingDrainSettings(
        pending_set_ttl_seconds=pending_set_ttl_seconds,
        pending_drain_delay_seconds=pending_drain_delay_seconds,
        pending_drain_limit=pending_drain_limit,
        pending_drain_schedule_ttl_seconds=pending_drain_schedule_ttl_seconds,
    )


def enqueue_pending_agent(
    agent_id: Union[str, UUID],
    *,
    ttl: int = _DEFAULT_PENDING_SET_TTL_SECONDS,
    client=None,
) -> bool:
    """Add an agent to the pending processing set. Returns True if newly added."""
    try:
        redis_client = client or get_redis_client()
        added = redis_client.sadd(_PENDING_SET_KEY, str(agent_id))
        if ttl > 0:
            redis_client.expire(_PENDING_SET_KEY, ttl)
        return bool(added)
    except Exception:
        logger.exception("Failed to enqueue pending processing for agent %s", agent_id)
        return False


def is_agent_pending(agent_id: Union[str, UUID], client=None) -> bool:
    """Check whether an agent is in the pending processing set."""
    try:
        redis_client = client or get_redis_client()
        return bool(redis_client.sismember(_PENDING_SET_KEY, str(agent_id)))
    except Exception:
        logger.exception("Failed to check pending processing for agent %s", agent_id)
        return False


def remove_pending_agent(agent_id: Union[str, UUID], client=None) -> None:
    """Remove an agent from the pending processing set."""
    try:
        redis_client = client or get_redis_client()
        redis_client.srem(_PENDING_SET_KEY, str(agent_id))
    except Exception:
        logger.exception("Failed to remove pending processing for agent %s", agent_id)


def pop_pending_agents(
    *,
    limit: int,
    client=None,
) -> list[str]:
    """Pop up to limit agent IDs from the pending processing set."""
    if limit <= 0:
        return []
    try:
        redis_client = client or get_redis_client()
        result = redis_client.spop(_PENDING_SET_KEY, count=limit)
    except Exception:
        logger.exception("Failed to pop pending agents")
        return []

    if not result:
        return []
    if isinstance(result, list):
        items = result
    else:
        items = [result]
    normalized: list[str] = []
    for item in items:
        if isinstance(item, (bytes, bytearray)):
            normalized.append(item.decode("utf-8", "ignore"))
        else:
            normalized.append(str(item))
    return normalized


def count_pending_agents(client=None) -> int:
    """Return the number of pending agents."""
    try:
        redis_client = client or get_redis_client()
        return int(redis_client.scard(_PENDING_SET_KEY))
    except Exception:
        logger.exception("Failed to count pending agents")
        return 0


def claim_pending_drain_slot(
    *,
    ttl: int = _DEFAULT_PENDING_DRAIN_SCHEDULE_TTL_SECONDS,
    client=None,
) -> bool:
    """Claim the pending-drain schedule slot. Returns True if claimed."""
    try:
        redis_client = client or get_redis_client()
        claimed = redis_client.set(
            _PENDING_DRAIN_SCHEDULE_KEY,
            "1",
            ex=ttl,
            nx=True,
        )
        return bool(claimed)
    except Exception:
        logger.exception("Failed to claim pending drain slot")
        return False


def clear_pending_drain_slot(client=None) -> None:
    """Clear the pending-drain schedule slot."""
    try:
        redis_client = client or get_redis_client()
        redis_client.delete(_PENDING_DRAIN_SCHEDULE_KEY)
    except Exception:
        logger.exception("Failed to clear pending drain slot")
