"""Retention (spec §2.1, §95–96): history is never deleted. The only prunable rows are bookkeeping:
- `observations` older than `retention_observations_days` that were `not_modified` / unchanged, are not referenced by a snapshot and are not
  the sensor's most recent observation (aggregate counts stay in `sensors.observation_count`; pruned totals are kept in settings_kv);
- `crawl_runs` older than `retention_crawl_runs_days`;
- finished `llm_jobs` older than `retention_llm_jobs_days`, summarised into `cost_ledger` (dimension `llm`, key `archived:<model>`) first.
Snapshots, changes, events, entities, metrics and series are never touched.
"""
from __future__ import annotations

import logging
from datetime import UTC, datetime, timedelta
from typing import Any

from companyatlas.config import settings
from companyatlas.db import execute, fetch_all, fetch_val, jsonb, transaction
from companyatlas.services.periodic import periodic

log = logging.getLogger(__name__)

PRUNED_KEY = "retention:pruned"
BATCH = 5000


async def prune_observations(*, dry_run: bool = False, batch: int = BATCH) -> int:
    cutoff = datetime.now(UTC) - timedelta(days=settings.retention_observations_days)
    total = 0
    async with transaction() as conn:
        while True:
            ids = [r["id"] for r in await fetch_all(conn, """
                select o.id from observations o
                where o.fetched_at < :cutoff and (o.not_modified or (o.changed = false and o.failure_class is null))
                  and not exists (select 1 from snapshots s where s.observation_id = o.id)
                  and o.id <> (select o2.id from observations o2 where o2.sensor_id = o.sensor_id order by o2.fetched_at desc limit 1)
                limit :batch""", cutoff=cutoff, batch=batch)]
            if not ids or dry_run:
                total += len(ids)
                break
            await execute(conn, "delete from observations where id = any(cast(:ids as text[]))", ids=ids)
            total += len(ids)
            if len(ids) < batch:
                break
        if total and not dry_run:
            prev = await fetch_val(conn, "select value from settings_kv where key = :k", k=PRUNED_KEY)
            state = dict(prev) if isinstance(prev, dict) else {}
            state["observations"] = int(state.get("observations", 0)) + total
            state["last_run_at"] = datetime.now(UTC).isoformat()
            await execute(conn, "insert into settings_kv (key, value, updated_at) values (:k, cast(:v as jsonb), now()) on conflict (key) do update set value = excluded.value, updated_at = now()",
                          k=PRUNED_KEY, v=jsonb(state))
    return total


async def prune_crawl_runs(*, dry_run: bool = False) -> int:
    cutoff = datetime.now(UTC) - timedelta(days=settings.retention_crawl_runs_days)
    async with transaction() as conn:
        n = int(await fetch_val(conn, "select count(*) from crawl_runs where started_at < :c", c=cutoff) or 0)
        if n and not dry_run:
            await execute(conn, "delete from crawl_runs where started_at < :c", c=cutoff)
    return n


async def archive_llm_jobs(*, dry_run: bool = False) -> int:
    cutoff = datetime.now(UTC) - timedelta(days=settings.retention_llm_jobs_days)
    async with transaction() as conn:
        rows = await fetch_all(conn, """select coalesce(model, 'unknown') as model, count(*) as n, coalesce(sum(request_tokens), 0) + coalesce(sum(response_tokens), 0) as tokens
                                        from llm_jobs where status in ('done', 'failed') and finished_at < :c group by 1""", c=cutoff)
        n = sum(int(r["n"]) for r in rows)
        if n and not dry_run:
            for r in rows:
                await execute(conn, """insert into cost_ledger (day, dimension, key, units, cost_estimate) values (current_date, 'llm', :key, :units, 0)
                                       on conflict (day, dimension, key) do update set units = cost_ledger.units + excluded.units""",
                              key=f"archived:{r['model']}", units=float(r["tokens"]))
                await execute(conn, """insert into cost_ledger (day, dimension, key, units, cost_estimate) values (current_date, 'llm', :key, :units, 0)
                                       on conflict (day, dimension, key) do update set units = cost_ledger.units + excluded.units""",
                              key=f"archived-jobs:{r['model']}", units=float(r["n"]))
            await execute(conn, "delete from llm_jobs where status in ('done', 'failed') and finished_at < :c", c=cutoff)
    return n


async def run_retention(*, dry_run: bool = False) -> dict[str, Any]:
    stats = {"observations": await prune_observations(dry_run=dry_run), "crawl_runs": await prune_crawl_runs(dry_run=dry_run),
             "llm_jobs": await archive_llm_jobs(dry_run=dry_run), "dry_run": dry_run}
    log.info("retention", extra=stats)
    return stats


@periodic("retention", cron="50 3 * * *")
async def retention_task() -> None:
    await run_retention()


__all__ = ["archive_llm_jobs", "prune_crawl_runs", "prune_observations", "run_retention"]