"""
State persistence utilities for flow manager.

Provides saving flow state to local files or S3 storage.
"""

import json
import os
from datetime import datetime, timezone
from typing import Any, Optional

from loguru import logger

from utils.s3_client import get_s3_client


async def save_flow_state(
    flow_manager: Any,
    recording_dir: Optional[str] = None,
    s3_key: Optional[str] = None,
    use_local: bool = True,
) -> bool:
    """
    Save flow manager state to either local file or S3.

    Args:
        flow_manager: The flow manager instance with state to save
        recording_dir: Directory for local recording (required if use_local=True)
        s3_key: S3 key for cloud storage (required if use_local=False)
        use_local: Whether to save locally or to S3

    Returns:
        True if save succeeded, False otherwise
    """
    if not flow_manager or not flow_manager.state:
        logger.debug("No flow manager state to save")
        return False

    logger.info("Saving flow manager state on cleanup")

    try:
        state_data = {
            "state": flow_manager.state,
            "current_node": flow_manager.current_node,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }
        state_json = json.dumps(state_data, indent=2)

        if use_local:
            return _save_to_local(state_json, recording_dir)
        else:
            return _save_to_s3(state_json, s3_key)

    except Exception as e:
        logger.error(f"Failed to save flow state: {e}")
        return False


def _save_to_local(state_json: str, recording_dir: Optional[str]) -> bool:
    """Save state to local file."""
    if not recording_dir:
        logger.error("Recording directory not provided for local save")
        return False

    state_file_path = os.path.join(recording_dir, "flow_state.json")

    try:
        os.makedirs(recording_dir, exist_ok=True)

        with open(state_file_path, "w") as f:
            f.write(state_json)

        logger.info(f"Flow state saved to {state_file_path}")
        return True

    except Exception as e:
        logger.error(f"Failed to save state to local file: {e}")
        return False


def _save_to_s3(state_json: str, s3_key: Optional[str]) -> bool:
    """Save state to S3 using centralized S3Client."""
    if not s3_key:
        logger.error("S3 key not provided for cloud save")
        return False

    try:
        s3_client = get_s3_client()
        return s3_client.upload_state(s3_key, state_json)

    except Exception as e:
        logger.error(f"Failed to save state to S3: {e}")
        return False