"""
cortex/capture/extraction_queue.py — Failed extraction retry queue.

When the extraction pipeline fails (LLM timeout, provider 503, etc.),
the conversation payload is queued here for automatic retry during the
next dream cycle.

Schema: cortex/storage/migrations/v16_extraction_queue.py
Used by: CortexPlugin.remember() (on failure), SleepConsolidator.dream() (retry)
"""
from __future__ import annotations

import json
import logging
import uuid
from datetime import datetime, timezone
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


@dataclass
class QueueEntry:
    """A single extraction queue entry."""
    id: str
    session_id: Optional[str]
    message_id: Optional[str]
    owner_id: str
    payload: str  # JSON
    failed_at: str
    retry_count: int = 0
    max_retries: int = 3
    status: str = "pending"
    error_message: Optional[str] = None
    created_at: Optional[str] = None
    updated_at: Optional[str] = None


async def enqueue_failed_extraction(
    storage: Any,
    *,
    session_id: Optional[str],
    owner_id: str,
    conversation: List[Dict[str, Any]],
    error_message: str,
    source_session_id: Optional[str] = None,
    source_message_id: Optional[str] = None,
    source_channel: Optional[str] = None,
    source_type: Optional[str] = None,
) -> Optional[str]:
    """Queue a failed extraction for retry.

    Stores the full conversation payload so it can be replayed during
    the next dream cycle.

    Returns the queue entry ID, or None if queueing failed.
    """
    entry_id = uuid.uuid4().hex
    now = datetime.now(timezone.utc).isoformat()

    payload = json.dumps({
        "conversation": conversation,
        "session_id": session_id,
        "owner_id": owner_id,
        "source_session_id": source_session_id,
        "source_message_id": source_message_id,
        "source_channel": source_channel,
        "source_type": source_type,
    })

    try:
        conn = await storage._get_conn()
        await conn.execute(
            """
            INSERT INTO extraction_queue (id, session_id, message_id, owner_id, payload,
                                          failed_at, retry_count, status, error_message,
                                          created_at, updated_at)
            VALUES (?, ?, ?, ?, ?, ?, 0, 'pending', ?, ?, ?)
            """,
            (entry_id, session_id, source_message_id, owner_id, payload,
             now, error_message, now, now),
        )
        await conn.commit()
        logger.info(
            "Queued failed extraction for retry: id=%s owner=%s session=%s error=%s",
            entry_id, owner_id, session_id, error_message[:100],
        )
        return entry_id
    except Exception as exc:
        logger.error("Failed to queue extraction for retry: %s", exc)
        return None


async def get_pending_entries(
    storage: Any,
    limit: int = 20,
) -> List[QueueEntry]:
    """Retrieve pending queue entries for retry processing."""
    try:
        conn = await storage._get_conn()
        async with conn.execute(
            """
            SELECT * FROM extraction_queue
            WHERE status = 'pending' AND retry_count < max_retries
            ORDER BY created_at ASC
            LIMIT ?
            """,
            (limit,),
        ) as cur:
            rows = await cur.fetchall()
        return [QueueEntry(**dict(r)) for r in rows]
    except Exception as exc:
        logger.error("Failed to fetch pending extraction queue entries: %s", exc)
        return []


async def mark_processing(storage: Any, entry_id: str) -> None:
    """Mark a queue entry as currently being processed."""
    try:
        conn = await storage._get_conn()
        await conn.execute(
            "UPDATE extraction_queue SET status = 'processing', updated_at = ? WHERE id = ?",
            (datetime.now(timezone.utc).isoformat(), entry_id),
        )
        await conn.commit()
    except Exception as exc:
        logger.error("Failed to mark queue entry %s as processing: %s", entry_id, exc)


async def mark_done(storage: Any, entry_id: str) -> None:
    """Mark a queue entry as successfully processed."""
    try:
        conn = await storage._get_conn()
        await conn.execute(
            "UPDATE extraction_queue SET status = 'done', updated_at = ? WHERE id = ?",
            (datetime.now(timezone.utc).isoformat(), entry_id),
        )
        await conn.commit()
    except Exception as exc:
        logger.error("Failed to mark queue entry %s as done: %s", entry_id, exc)


async def mark_failed(storage: Any, entry_id: str, error_message: str) -> None:
    """Increment retry count and update error. If max retries exceeded, set status='failed'."""
    try:
        conn = await storage._get_conn()
        now = datetime.now(timezone.utc).isoformat()
        await conn.execute(
            """
            UPDATE extraction_queue
            SET retry_count = retry_count + 1,
                error_message = ?,
                failed_at = ?,
                updated_at = ?,
                status = CASE
                    WHEN retry_count + 1 >= max_retries THEN 'failed'
                    ELSE 'pending'
                END
            WHERE id = ?
            """,
            (error_message, now, now, entry_id),
        )
        await conn.commit()
    except Exception as exc:
        logger.error("Failed to mark queue entry %s as failed: %s", entry_id, exc)


async def get_queue_stats(storage: Any) -> Dict[str, int]:
    """Return queue statistics by status."""
    try:
        conn = await storage._get_conn()
        async with conn.execute(
            "SELECT status, COUNT(*) as cnt FROM extraction_queue GROUP BY status"
        ) as cur:
            rows = await cur.fetchall()
        return {r["status"]: r["cnt"] for r in rows}
    except Exception as exc:
        logger.error("Failed to get extraction queue stats: %s", exc)
        return {}
