from __future__ import annotations

import logging
from datetime import UTC, datetime, timedelta

from host_agent.config import HostAgentConfig
from storage.task_metadata import TaskMetadataStore
from storage.timeline import Timeline

logger = logging.getLogger(__name__)


def prune_task_history(
    metadata_store: TaskMetadataStore,
    timeline: Timeline,
    *,
    config: HostAgentConfig,
    now: datetime | None = None,
) -> int:
    """Delete tasks beyond the configured retention window/count.

    Computes two candidate retained sets -- the last ``max_count`` tasks
    and tasks younger than ``max_age_days`` -- and keeps whichever set
    is smaller (i.e. the more restrictive bound always wins).

    Returns the number of tasks pruned.
    """
    reference = now or datetime.now(UTC)
    tasks = metadata_store.list_tasks()

    keep_by_count = {task["id"] for task in tasks[: config.task_retention_max_count]}
    cutoff = reference - timedelta(days=config.task_retention_max_age_days)
    keep_by_age = {
        task["id"] for task in tasks if _parse_timestamp(task["created_at"]) >= cutoff
    }
    keep_ids = keep_by_count if len(keep_by_count) <= len(keep_by_age) else keep_by_age

    pruned = 0
    for task in tasks:
        if task["id"] not in keep_ids:
            _safe_delete(timeline, metadata_store, task["id"])
            pruned += 1
    return pruned


def _safe_delete(
    timeline: Timeline, metadata_store: TaskMetadataStore, task_id: str
) -> None:
    try:
        timeline.delete_task(task_id)
    except Exception:
        logger.debug("timeline delete failed for task %s", task_id, exc_info=True)
    try:
        metadata_store.delete_task(task_id)
    except Exception:
        logger.debug("metadata delete failed for task %s", task_id, exc_info=True)


def _parse_timestamp(value: str) -> datetime:
    parsed = datetime.fromisoformat(value)
    if parsed.tzinfo is None:
        parsed = parsed.replace(tzinfo=UTC)
    return parsed