app/idempotency.py

import hashlib
import json
from uuid import UUID
from datetime import datetime, timezone, timedelta
from sqlalchemy.orm import Session
from sqlalchemy.dialects.postgresql import insert
from fastapi import Header, Response
from app.models import IdempotencyKey
from app.errors import AppError

def require_idempotency_key(x_idempotency_key: str | None = Header(default=None)) -> UUID:
    if x_idempotency_key is None:
        raise AppError(400, "IDEMPOTENCY_KEY_MISSING", "Header X-Idempotency-Key is required.")
    try:
        key = UUID(x_idempotency_key)
        if key.version != 4:
            raise ValueError
        return key
    except ValueError:
        raise AppError(422, "IDEMPOTENCY_KEY_INVALID", "X-Idempotency-Key must be a UUIDv4.")

def body_hash(data: dict) -> str:
    canonical = json.dumps(data, sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(canonical.encode()).hexdigest()

def purge_stale(db: Session, user_id: UUID):
    now = datetime.now(timezone.utc)
    # Delete expired keys OR in_progress keys older than 60 seconds
    db.query(IdempotencyKey).filter(
        IdempotencyKey.user_id == user_id,
        (IdempotencyKey.expires_at < now) | 
        ((IdempotencyKey.status == "in_progress") & (IdempotencyKey.created_at < now - timedelta(seconds=60)))
    ).delete()
    db.commit()

def claim_key(db: Session, user_id: UUID, key: UUID, req_hash: str, path: str) -> bool:
    stmt = insert(IdempotencyKey).values(
        user_id=user_id, key=key, request_method="POST", request_path=path,
        request_hash=req_hash, status="in_progress",
        expires_at=datetime.now(timezone.utc) + timedelta(hours=24)
    ).on_conflict_do_nothing().returning(IdempotencyKey.key)
    
    result = db.execute(stmt)
    db.commit()
    # fetchone() returns None if the insert was skipped (conflict)
    return result.fetchone() is not None

def complete_key(db: Session, user_id: UUID, key: UUID, status: int, body: dict):
    row = db.query(IdempotencyKey).filter_by(user_id=user_id, key=key).first()
    if row:
        row.status = "completed"
        row.response_status = status
        row.response_body = body
        db.commit()

def release_claim(db: Session, user_id: UUID, key: UUID):
    db.query(IdempotencyKey).filter_by(user_id=user_id, key=key).delete()
    db.commit()

def replay_or_conflict(db: Session, user_id: UUID, key: UUID, req_hash: str):
    row = db.query(IdempotencyKey).filter_by(user_id=user_id, key=key).first()
    if not row:
        raise AppError(500, "INTERNAL_ERROR", "Idempotency key vanished.")
    
    if row.request_hash != req_hash:
        raise AppError(422, "IDEMPOTENCY_KEY_REUSED", "This idempotency key was used with a different payload.")
    
    if row.status == "in_progress":
        raise AppError(409, "IDEMPOTENCY_REQUEST_IN_PROGRESS", "Original request still processing.", headers={"Retry-After": "2"})
    
    return Response(
        content=json.dumps(row.response_body),
        status_code=row.response_status,
        media_type="application/json",
        headers={"Idempotent-Replayed": "true"}
    )