"""Pipeline execution state persistence — Redis hot state + PostgreSQL cold storage.

Architecture:
    PipelineStateMemory  — In-memory fallback (always available, for testing)
    PipelineStateRedis   — Redis hot state (low-latency reads/writes)
    PipelineStatePG      — PostgreSQL cold persistence (durable audit trail)
    PipelineStateManager — Unified manager (Redis + PG dual write, fallback chain)
"""

from __future__ import annotations

import json
import logging
import uuid
from datetime import datetime, timezone
from typing import Any, Callable, Coroutine

from agentkit.orchestrator.pipeline_models import (
    PipelineExecutionModel,
    PipelineStepHistoryModel,
)

logger = logging.getLogger(__name__)

# Redis key patterns
_EXEC_KEY_PREFIX = "agentkit:pipeline:exec:"
_INDEX_KEY = "agentkit:pipeline:index"
_TTL_SECONDS = 7 * 24 * 3600  # 7 days


class PipelineStateMemory:
    """In-memory pipeline state storage (testing / fallback)."""

    def __init__(self) -> None:
        self._executions: dict[str, dict[str, Any]] = {}
        self._step_history: dict[str, list[dict[str, Any]]] = {}

    async def create_execution(
        self,
        pipeline_name: str,
        steps: list[str],
        input_data: dict[str, Any] | None = None,
        tenant_id: str | None = None,
    ) -> str:
        execution_id = str(uuid.uuid4())
        now = datetime.now(timezone.utc).isoformat()
        self._executions[execution_id] = {
            "id": execution_id,
            "pipeline_name": pipeline_name,
            "status": "running",
            "current_step": steps[0] if steps else None,
            "completed_steps": [],
            "step_results": {},
            "input_data": input_data,
            "final_output": None,
            "error_message": None,
            "tenant_id": tenant_id,
            "created_at": now,
            "updated_at": now,
            "completed_at": None,
        }
        self._step_history[execution_id] = []
        return execution_id

    async def update_step(
        self,
        execution_id: str,
        step_name: str,
        status: str,
        output: dict[str, Any] | None = None,
        error: str | None = None,
        duration_ms: int | None = None,
    ) -> None:
        exec_state = self._executions.get(execution_id)
        if exec_state is None:
            logger.warning(f"Execution '{execution_id}' not found for step update")
            return

        exec_state["current_step"] = step_name
        exec_state["updated_at"] = datetime.now(timezone.utc).isoformat()

        if status == "completed":
            if step_name not in exec_state["completed_steps"]:
                exec_state["completed_steps"].append(step_name)
            if output is not None:
                exec_state["step_results"][step_name] = output
        elif status == "failed":
            exec_state["error_message"] = error

        # Record step history event
        step_event: dict[str, Any] = {
            "id": str(uuid.uuid4()),
            "execution_id": execution_id,
            "step_name": step_name,
            "status": status,
            "output_data": output,
            "error_message": error,
            "duration_ms": duration_ms,
            "started_at": datetime.now(timezone.utc).isoformat(),
            "completed_at": datetime.now(timezone.utc).isoformat() if status in ("completed", "failed") else None,
        }
        self._step_history[execution_id].append(step_event)

    async def complete_execution(
        self,
        execution_id: str,
        final_output: dict[str, Any] | None = None,
    ) -> None:
        exec_state = self._executions.get(execution_id)
        if exec_state is None:
            return
        now = datetime.now(timezone.utc).isoformat()
        exec_state["status"] = "completed"
        exec_state["final_output"] = final_output
        exec_state["updated_at"] = now
        exec_state["completed_at"] = now

    async def fail_execution(
        self,
        execution_id: str,
        step_name: str,
        error: str,
    ) -> None:
        exec_state = self._executions.get(execution_id)
        if exec_state is None:
            return
        now = datetime.now(timezone.utc).isoformat()
        exec_state["status"] = "failed"
        exec_state["error_message"] = f"Step '{step_name}' failed: {error}"
        exec_state["updated_at"] = now
        exec_state["completed_at"] = now

    async def get_execution(self, execution_id: str) -> dict[str, Any] | None:
        return self._executions.get(execution_id)

    async def list_executions(
        self,
        status: str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        results = list(self._executions.values())
        if status:
            results = [e for e in results if e.get("status") == status]
        results.sort(key=lambda e: e.get("created_at", ""), reverse=True)
        return results[offset : offset + limit]

    async def get_step_history(self, execution_id: str) -> list[dict[str, Any]]:
        return self._step_history.get(execution_id, [])

    def get_execution_sync(self, execution_id: str) -> dict[str, Any] | None:
        """Synchronous accessor for execution state (used by Redis dual-write)."""
        return self._executions.get(execution_id)


class PipelineStateRedis:
    """Redis-backed pipeline state storage (hot state).

    Uses Redis Hash for execution state and Sorted Set for indexing.
    Falls back to PipelineStateMemory if Redis is unavailable.
    Automatically retries Redis after a cooldown period.
    """

    _RECOVERY_COOLDOWN_SECONDS = 30

    def __init__(self, redis_url: str = "redis://localhost:6379/0") -> None:
        self._redis_url = redis_url
        self._redis: Any = None
        self._fallback = PipelineStateMemory()
        self._use_fallback = False
        self._fallback_since: float | None = None

    async def _get_redis(self):
        if self._redis is None:
            import redis.asyncio as aioredis

            self._redis = aioredis.from_url(
                self._redis_url,
                decode_responses=True,
            )
        return self._redis

    async def _safe_redis_call(
        self, fn: Callable[..., Coroutine[Any, Any, Any]], *args: Any, **kwargs: Any
    ) -> Any:
        """Execute a Redis call, falling back to memory on failure.

        After falling back, periodically retries Redis to enable recovery.
        On successful recovery, the original operation is executed immediately.
        """
        if self._use_fallback:
            # Check if enough time has passed to attempt recovery
            if self._fallback_since is not None:
                import time as _time
                elapsed = _time.monotonic() - self._fallback_since
                if elapsed >= self._RECOVERY_COOLDOWN_SECONDS:
                    try:
                        self._redis = None
                        redis = await self._get_redis()
                        await redis.ping()
                        # Recovery successful — continue to execute the operation
                        self._use_fallback = False
                        self._fallback_since = None
                        logger.info("Redis connection recovered, switching back from fallback")
                        # Fall through to execute the actual operation on Redis
                    except Exception:
                        # Still down, reset cooldown timer
                        self._fallback_since = _time.monotonic()
                        return None
                else:
                    return None
            else:
                return None
        try:
            redis = await self._get_redis()
            return await fn(redis, *args, **kwargs)
        except Exception as exc:
            logger.warning(f"Redis operation failed, switching to memory fallback: {exc}")
            self._use_fallback = True
            import time as _time
            self._fallback_since = _time.monotonic()
            self._redis = None
            return None

    def _key(self, execution_id: str) -> str:
        return f"{_EXEC_KEY_PREFIX}{execution_id}"

    async def create_execution(
        self,
        pipeline_name: str,
        steps: list[str],
        input_data: dict[str, Any] | None = None,
        tenant_id: str | None = None,
    ) -> str:
        # Always write to fallback first for consistency
        execution_id = await self._fallback.create_execution(
            pipeline_name, steps, input_data, tenant_id
        )

        # Try Redis
        async def _redis_create(redis: Any) -> None:
            state = self._fallback.get_execution_sync(execution_id)
            score = datetime.now(timezone.utc).timestamp()
            pipe = redis.pipeline()
            pipe.set(self._key(execution_id), json.dumps(state), ex=_TTL_SECONDS)
            pipe.zadd(_INDEX_KEY, {execution_id: score})
            await pipe.execute()

        await self._safe_redis_call(_redis_create)
        return execution_id

    async def update_step(
        self,
        execution_id: str,
        step_name: str,
        status: str,
        output: dict[str, Any] | None = None,
        error: str | None = None,
        duration_ms: int | None = None,
    ) -> None:
        await self._fallback.update_step(execution_id, step_name, status, output, error, duration_ms)

        async def _redis_update(redis: Any) -> None:
            state = self._fallback.get_execution_sync(execution_id)
            if state is None:
                return
            await redis.set(self._key(execution_id), json.dumps(state), ex=_TTL_SECONDS)

        await self._safe_redis_call(_redis_update)

    async def complete_execution(
        self,
        execution_id: str,
        final_output: dict[str, Any] | None = None,
    ) -> None:
        await self._fallback.complete_execution(execution_id, final_output)

        async def _redis_complete(redis: Any) -> None:
            state = self._fallback.get_execution_sync(execution_id)
            if state is None:
                return
            await redis.set(self._key(execution_id), json.dumps(state), ex=_TTL_SECONDS)

        await self._safe_redis_call(_redis_complete)

    async def fail_execution(
        self,
        execution_id: str,
        step_name: str,
        error: str,
    ) -> None:
        await self._fallback.fail_execution(execution_id, step_name, error)

        async def _redis_fail(redis: Any) -> None:
            state = self._fallback.get_execution_sync(execution_id)
            if state is None:
                return
            await redis.set(self._key(execution_id), json.dumps(state), ex=_TTL_SECONDS)

        await self._safe_redis_call(_redis_fail)

    async def get_execution(self, execution_id: str) -> dict[str, Any] | None:
        # Try Redis first
        if not self._use_fallback:
            try:
                redis = await self._get_redis()
                raw = await redis.get(self._key(execution_id))
                if raw is not None:
                    return json.loads(raw)
            except Exception:
                pass

        # Fallback to memory
        return await self._fallback.get_execution(execution_id)

    async def list_executions(
        self,
        status: str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        # Try Redis sorted set for efficient listing
        if not self._use_fallback:
            try:
                redis = await self._get_redis()
                # Get recent execution IDs from sorted set (newest first)
                ids = await redis.zrevrange(_INDEX_KEY, offset, offset + limit - 1)
                if ids:
                    keys = [self._key(eid) for eid in ids]
                    values = await redis.mget(keys)
                    results = []
                    for raw in values:
                        if raw is None:
                            continue
                        state = json.loads(raw)
                        if status is None or state.get("status") == status:
                            results.append(state)
                    return results
            except Exception:
                pass

        return await self._fallback.list_executions(status, limit, offset)

    async def get_step_history(self, execution_id: str) -> list[dict[str, Any]]:
        return await self._fallback.get_step_history(execution_id)

    async def health_check(self) -> bool:
        if self._use_fallback:
            return False
        try:
            redis = await self._get_redis()
            return await redis.ping()
        except Exception:
            return False

    @property
    def using_fallback(self) -> bool:
        return self._use_fallback


class PipelineStatePG:
    """PostgreSQL cold persistence for pipeline execution records.

    If session_factory is None, all methods are no-op.
    """

    def __init__(self, session_factory: Any = None) -> None:
        self._session_factory = session_factory

    @property
    def enabled(self) -> bool:
        return self._session_factory is not None

    async def persist_execution(self, state: dict[str, Any]) -> None:
        """Write a completed/failed execution to PostgreSQL."""
        if not self.enabled:
            return
        try:
            from sqlalchemy.ext.asyncio import AsyncSession

            async with self._session_factory() as session:
                model = PipelineExecutionModel(
                    id=state["id"],
                    pipeline_name=state["pipeline_name"],
                    status=state["status"],
                    current_step=state.get("current_step"),
                    completed_steps=state.get("completed_steps", []),
                    step_results=state.get("step_results", {}),
                    input_data=state.get("input_data"),
                    final_output=state.get("final_output"),
                    error_message=state.get("error_message"),
                    tenant_id=state.get("tenant_id"),
                    created_at=datetime.fromisoformat(state["created_at"]) if state.get("created_at") else None,
                    updated_at=datetime.fromisoformat(state["updated_at"]) if state.get("updated_at") else None,
                    completed_at=datetime.fromisoformat(state["completed_at"]) if state.get("completed_at") else None,
                )
                await session.merge(model)
                await session.commit()
        except Exception as exc:
            logger.error(f"Failed to persist execution to PG: {exc}")

    async def persist_step_history(
        self, execution_id: str, steps: list[dict[str, Any]]
    ) -> None:
        """Write step history to PostgreSQL."""
        if not self.enabled:
            return
        try:
            async with self._session_factory() as session:
                for idx, step in enumerate(steps):
                    model = PipelineStepHistoryModel(
                        id=step.get("id", str(uuid.uuid4())),
                        execution_id=execution_id,
                        step_name=step["step_name"],
                        step_index=idx,
                        status=step["status"],
                        input_data=step.get("input_data"),
                        output_data=step.get("output_data"),
                        error_message=step.get("error_message"),
                        duration_ms=step.get("duration_ms"),
                        retry_attempt=step.get("retry_attempt", 0),
                        started_at=datetime.fromisoformat(step["started_at"]) if step.get("started_at") else None,
                        completed_at=datetime.fromisoformat(step["completed_at"]) if step.get("completed_at") else None,
                    )
                    await session.merge(model)
                await session.commit()
        except Exception as exc:
            logger.error(f"Failed to persist step history to PG: {exc}")

    async def query_executions(
        self,
        pipeline_name: str | None = None,
        status: str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        """Query historical executions from PostgreSQL."""
        if not self.enabled:
            return []
        try:
            from sqlalchemy import select

            async with self._session_factory() as session:
                stmt = select(PipelineExecutionModel).order_by(
                    PipelineExecutionModel.created_at.desc()
                )
                if pipeline_name:
                    stmt = stmt.where(
                        PipelineExecutionModel.pipeline_name == pipeline_name
                    )
                if status:
                    stmt = stmt.where(PipelineExecutionModel.status == status)
                stmt = stmt.offset(offset).limit(limit)
                result = await session.execute(stmt)
                rows = result.scalars().all()
                return [self._model_to_dict(row) for row in rows]
        except Exception as exc:
            logger.error(f"Failed to query executions from PG: {exc}")
            return []

    async def get_execution(self, execution_id: str) -> dict[str, Any] | None:
        """Get a single execution from PostgreSQL (for Redis miss fallback)."""
        if not self.enabled:
            return None
        try:
            from sqlalchemy import select

            async with self._session_factory() as session:
                stmt = select(PipelineExecutionModel).where(
                    PipelineExecutionModel.id == execution_id
                )
                result = await session.execute(stmt)
                row = result.scalar_one_or_none()
                if row is None:
                    return None
                return self._model_to_dict(row)
        except Exception as exc:
            logger.error(f"Failed to get execution from PG: {exc}")
            return None

    @staticmethod
    def _model_to_dict(model: PipelineExecutionModel) -> dict[str, Any]:
        return {
            "id": model.id,
            "pipeline_name": model.pipeline_name,
            "status": model.status,
            "current_step": model.current_step,
            "completed_steps": model.completed_steps or [],
            "step_results": model.step_results or {},
            "input_data": model.input_data,
            "final_output": model.final_output,
            "error_message": model.error_message,
            "tenant_id": model.tenant_id,
            "created_at": model.created_at.isoformat() if model.created_at else None,
            "updated_at": model.updated_at.isoformat() if model.updated_at else None,
            "completed_at": model.completed_at.isoformat() if model.completed_at else None,
        }


class PipelineStateManager:
    """Unified pipeline state manager — Redis hot + PG cold.

    - create / update → Redis (with in-memory fallback)
    - complete / fail → Redis + async persist to PG
    - get → Redis first, PG fallback
    - list → Redis for recent, PG for historical
    """

    def __init__(
        self,
        redis_url: str | None = None,
        session_factory: Any = None,
    ) -> None:
        if redis_url:
            self._hot = PipelineStateRedis(redis_url=redis_url)
        else:
            self._hot = PipelineStateMemory()
        self._cold = PipelineStatePG(session_factory=session_factory)

    @property
    def hot_store(self) -> PipelineStateMemory | PipelineStateRedis:
        return self._hot

    @property
    def cold_store(self) -> PipelineStatePG:
        return self._cold

    async def create_execution(
        self,
        pipeline_name: str,
        steps: list[str],
        input_data: dict[str, Any] | None = None,
        tenant_id: str | None = None,
    ) -> str:
        return await self._hot.create_execution(pipeline_name, steps, input_data, tenant_id)

    async def update_step(
        self,
        execution_id: str,
        step_name: str,
        status: str,
        output: dict[str, Any] | None = None,
        error: str | None = None,
        duration_ms: int | None = None,
    ) -> None:
        await self._hot.update_step(execution_id, step_name, status, output, error, duration_ms)

    async def complete_execution(
        self,
        execution_id: str,
        final_output: dict[str, Any] | None = None,
    ) -> None:
        await self._hot.complete_execution(execution_id, final_output)
        # Persist to PG
        state = await self._hot.get_execution(execution_id)
        if state:
            await self._cold.persist_execution(state)
            step_history = await self._hot.get_step_history(execution_id)
            if step_history:
                await self._cold.persist_step_history(execution_id, step_history)

    async def fail_execution(
        self,
        execution_id: str,
        step_name: str,
        error: str,
    ) -> None:
        await self._hot.fail_execution(execution_id, step_name, error)
        # Persist to PG
        state = await self._hot.get_execution(execution_id)
        if state:
            await self._cold.persist_execution(state)
            step_history = await self._hot.get_step_history(execution_id)
            if step_history:
                await self._cold.persist_step_history(execution_id, step_history)

    async def get_execution(self, execution_id: str) -> dict[str, Any] | None:
        # Redis / memory first
        state = await self._hot.get_execution(execution_id)
        if state is not None:
            return state
        # PG fallback
        return await self._cold.get_execution(execution_id)

    async def list_executions(
        self,
        status: str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        # Hot store for recent executions
        results = await self._hot.list_executions(status, limit, offset)
        if results:
            return results
        # Cold store for historical queries
        return await self._cold.query_executions(status=status, limit=limit, offset=offset)

    async def get_step_history(self, execution_id: str) -> list[dict[str, Any]]:
        return await self._hot.get_step_history(execution_id)

    async def health_check(self) -> dict[str, bool]:
        """Check health of both stores."""
        hot_ok = True
        if isinstance(self._hot, PipelineStateRedis):
            hot_ok = await self._hot.health_check()
        cold_ok = self._cold.enabled
        return {"hot": hot_ok, "cold": cold_ok}
