"""Memory Segmentation & Lifecycle.

Groups memories into three temperature bands based on age and access
frequency, and rotates them periodically so that searches stay fast as
memory volume grows. This is the foundation for enterprise-scale
retention policies.

Segments:
    ACTIVE: last N months (default 3) — full, fast access
    WARM:   3-12 months — on-demand access, semantic-search only
    COLD:   12+ months — compressed, rarely searched

Keys:
    - Overlap window: during rotation, the last 2 weeks of the retiring
      segment are COPIED into the next so queries spanning the transition
      still find context.
    - Cascading search: active -> warm -> cold, with early termination
      if enough high-relevance results are found.
    - Importance scoring: memories with access_count > threshold are
      pinned to ACTIVE regardless of age.
"""

import logging
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Literal

logger = logging.getLogger("soagila.segmentation")

Segment = Literal["active", "warm", "cold"]


@dataclass
class SegmentationPolicy:
    active_days: int = 90          # 3 months
    warm_days: int = 365           # 12 months
    # Age cutoffs: items older than warm_days go to cold
    overlap_days: int = 14         # 2-week overlap on rotation
    importance_pin_threshold: int = 5  # access_count >= this pins to ACTIVE
    compress_cold: bool = True     # Compress cold segment with gzip


DEFAULT_POLICY = SegmentationPolicy()


def compute_segment(
    created_at: datetime | str,
    access_count: int = 0,
    last_accessed: datetime | str | None = None,
    policy: SegmentationPolicy = DEFAULT_POLICY,
) -> Segment:
    """Compute which segment a memory belongs to right now.

    Rules (in order):
        1. access_count >= importance_pin_threshold → ACTIVE (pinned)
        2. last_accessed within active_days → ACTIVE
        3. created within active_days → ACTIVE
        4. created within warm_days → WARM
        5. older → COLD
    """
    now = datetime.now(timezone.utc)

    # Rule 1: pinned by importance
    if access_count >= policy.importance_pin_threshold:
        return "active"

    created_dt = _parse_datetime(created_at)
    last_dt = _parse_datetime(last_accessed) if last_accessed else None

    # Rule 2: recently accessed
    if last_dt and (now - last_dt).days < policy.active_days:
        return "active"

    # Rule 3: recently created
    age_days = (now - created_dt).days
    if age_days < policy.active_days:
        return "active"

    # Rule 4: warm range
    if age_days < policy.warm_days:
        return "warm"

    # Rule 5: cold
    return "cold"


def _parse_datetime(value: datetime | str) -> datetime:
    if isinstance(value, datetime):
        return value if value.tzinfo else value.replace(tzinfo=timezone.utc)
    # String: try ISO first, then common formats
    try:
        dt = datetime.fromisoformat(value.replace("Z", "+00:00"))
        return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)
    except Exception:
        return datetime.now(timezone.utc)


async def cascading_search(
    query_text: str,
    query_embedding: list[float] | None,
    *,
    max_results: int = 10,
    min_active_results: int = 3,
    policy: SegmentationPolicy = DEFAULT_POLICY,
) -> list:
    """Search ACTIVE first; only fall back to WARM/COLD if insufficient.

    Returns memories in relevance order. Stops early if ACTIVE yields
    at least `min_active_results` high-confidence matches.
    """
    from soagila.core.tcmg.episodic_memory import query as ep_query

    now = datetime.now(timezone.utc)
    active_cutoff = (now - timedelta(days=policy.active_days)).isoformat()
    warm_cutoff = (now - timedelta(days=policy.warm_days)).isoformat()

    # Step 1: query ACTIVE
    active_results = await _query_segment(
        query_text,
        query_embedding,
        since=active_cutoff,
        limit=max_results,
    )

    if len(active_results) >= min_active_results:
        logger.debug(
            f"Segmentation: ACTIVE yielded {len(active_results)} — skipping warm/cold"
        )
        return active_results

    # Step 2: WARM fallback
    warm_results = await _query_segment(
        query_text,
        query_embedding,
        since=warm_cutoff,
        until=active_cutoff,
        limit=max_results - len(active_results),
    )

    results = active_results + warm_results
    if len(results) >= min_active_results:
        logger.debug(
            f"Segmentation: ACTIVE+WARM yielded {len(results)} — skipping cold"
        )
        return results

    # Step 3: COLD fallback (if user is on tier that allows it)
    cold_results = await _query_segment(
        query_text,
        query_embedding,
        until=warm_cutoff,
        limit=max_results - len(results),
    )

    results += cold_results
    logger.debug(
        f"Segmentation: full cascade yielded "
        f"active={len(active_results)} warm={len(warm_results)} cold={len(cold_results)}"
    )
    return results[:max_results]


async def _query_segment(
    query_text: str,
    query_embedding: list[float] | None,
    *,
    since: str | None = None,
    until: str | None = None,
    limit: int = 10,
):
    """Query episodic memories within an optional time window."""
    from soagila.core.tcmg.episodic_memory import query as ep_query

    # Delegate to existing query; post-filter by time window
    results = await ep_query("local", query_text, limit=limit * 3)

    def _in_window(mem) -> bool:
        created = getattr(mem, "created_at", None)
        if not created:
            return True
        if isinstance(created, str):
            try:
                created = datetime.fromisoformat(created.replace("Z", "+00:00"))
            except Exception:
                return True
        if not created.tzinfo:
            created = created.replace(tzinfo=timezone.utc)
        iso = created.isoformat()
        if since and iso < since:
            return False
        if until and iso >= until:
            return False
        return True

    filtered = [m for m in results if _in_window(m)]
    return filtered[:limit]


async def apply_retention_policy(
    policy: SegmentationPolicy = DEFAULT_POLICY,
    tier: str = "free",
) -> dict:
    """Enforce per-tier retention by marking old memories as decayed.

    This is the hook the background scheduler should call daily.
    """
    # Retention days by tier — memories older than this are soft-deleted
    tier_retention = {
        "free": 7,
        "pro": 365,
        "enterprise": 365 * 3,
        "enterprise_team": 365 * 7,  # 7 years — typical compliance window
        "enterprise_custom": 365 * 10,  # 10 years — healthcare
    }
    retention_days = tier_retention.get(tier, 7)
    cutoff = (
        datetime.now(timezone.utc) - timedelta(days=retention_days)
    ).isoformat()

    from soagila.db.sqlite_init import get_db
    db = await get_db()

    # Mark old memories as decayed (preserves audit chain)
    result = await db.execute(
        """
        UPDATE episodic_memories
        SET is_decayed = 1, updated_at = datetime('now')
        WHERE created_at < ?
          AND is_decayed = 0
          AND access_count < ?
        """,
        (cutoff, policy.importance_pin_threshold),
    )
    await db.commit()

    affected = getattr(result, "rowcount", 0) or 0
    logger.info(
        f"Retention: {affected} memories marked decayed "
        f"(tier={tier}, cutoff={retention_days}d)"
    )

    return {
        "tier": tier,
        "retention_days": retention_days,
        "cutoff": cutoff,
        "memories_decayed": affected,
    }
