"""Durable workflow server with request checkpointing and recovery.

This module provides :class:`DurableWorkflowServer`, which extends
:class:`WorkflowServer` with request durability features: automatic
checkpointing, idempotent deduplication, event-sourced audit trail,
and recovery from the latest checkpoint via :class:`DurableRequest`.
"""

import asyncio
import json
import logging
from datetime import UTC, datetime
from typing import Any, Callable, Dict, List, Optional

from fastapi import HTTPException, Request, Response
from fastapi.responses import JSONResponse

from ..middleware.gateway.checkpoint_manager import CheckpointManager
from ..middleware.gateway.deduplicator import RequestDeduplicator
from ..middleware.gateway.durable_request import (
    DurableRequest,
    RequestMetadata,
    RequestState,
)
from ..middleware.gateway.event_store import (
    EventStore,
    EventType,
    performance_metrics_projection,
    request_state_projection,
)
from .workflow_server import WorkflowServer

logger = logging.getLogger(__name__)


class DurableWorkflowServer(WorkflowServer):
    """Workflow server with durable request handling.

    Extends the basic WorkflowServer with:
    - Request durability and checkpointing
    - Automatic deduplication
    - Event sourcing for audit trail
    - Long-running request support
    - Recovery mechanisms

    This server provides reliability features for production deployments
    where request durability is important. For full enterprise features,
    consider using EnterpriseWorkflowServer.
    """

    def __init__(
        self,
        title: str = "Kailash Durable Workflow Server",
        description: str = "Durable workflow server with checkpointing",
        version: str = "1.0.0",
        max_workers: int = 10,
        cors_origins: Optional[list[str]] = None,
        # Durability configuration
        enable_durability: bool = True,
        checkpoint_manager: Optional[CheckpointManager] = None,
        deduplicator: Optional[RequestDeduplicator] = None,
        event_store: Optional[EventStore] = None,
        durability_opt_in: bool = True,  # If True, durability is opt-in per endpoint
        **kwargs,
    ):
        """Initialize durable workflow server."""
        super().__init__(
            title=title,
            description=description,
            version=version,
            max_workers=max_workers,
            cors_origins=cors_origins,
            **kwargs,
        )

        # Durability components
        self.enable_durability = enable_durability
        self.durability_opt_in = durability_opt_in
        self.checkpoint_manager = checkpoint_manager or CheckpointManager()
        # Initialize deduplicator lazily to avoid event loop issues
        self._deduplicator = deduplicator
        self._event_store = event_store

        # Track active requests
        self.active_requests: Dict[str, DurableRequest] = {}

        # Track background tasks
        self._background_tasks: List[asyncio.Task] = []

        # Initialize durability components lazily
        self._durability_initialized = False

        # Add durability middleware if enabled
        if self.enable_durability:
            self._add_durability_middleware()

        # Register durability endpoints
        self._register_durability_endpoints()

    @property
    def deduplicator(self) -> RequestDeduplicator:
        """Get deduplicator instance, initializing if needed."""
        if self._deduplicator is None:
            self._deduplicator = RequestDeduplicator()
        return self._deduplicator

    @property
    def event_store(self) -> EventStore:
        """Get event store instance, initializing if needed."""
        if self._event_store is None:
            self._event_store = EventStore()
            # Register event projections
            self._event_store.register_projection(
                "request_states",
                request_state_projection,
            )
            self._event_store.register_projection(
                "performance_metrics",
                performance_metrics_projection,
            )
        return self._event_store

    def _add_durability_middleware(self):
        """Add middleware for durable request handling."""

        @self.app.middleware("http")
        async def durability_middleware(request: Request, call_next):
            """Middleware to handle request durability."""
            # Check if this endpoint should use durability
            should_be_durable = self._should_use_durability(request)

            if not should_be_durable:
                # Pass through without durability
                return await call_next(request)

            # Extract request metadata.
            # Extract body BEFORE creating metadata so dedup fingerprinting
            # includes the actual request body content.
            # See: https://github.com/terrene-foundation/kailash-py/issues/175
            body_content = None
            if request.method in ["POST", "PUT", "PATCH"]:
                try:
                    raw_body = await request.body()
                    if raw_body:
                        # Try JSON parse first for structured fingerprinting
                        try:
                            body_content = json.loads(raw_body)
                        except (json.JSONDecodeError, ValueError):
                            # Non-JSON body: use raw string for fingerprinting
                            body_content = raw_body.decode("utf-8", errors="replace")
                except Exception:
                    logger.warning(
                        "Failed to extract request body for dedup fingerprint "
                        "on %s %s -- dedup will use path-only fingerprint for "
                        "this request",
                        request.method,
                        str(request.url.path),
                    )

            request_id = (
                request.headers.get("X-Request-ID")
                or f"req_{datetime.now(UTC).timestamp()}"
            )
            current_time = datetime.now(UTC)
            metadata = RequestMetadata(
                request_id=request_id,
                method=request.method,
                path=str(request.url.path),
                headers=dict(request.headers),
                query_params=dict(request.query_params),
                body=body_content,
                client_ip=request.client.host if request.client else "unknown",
                user_id=None,  # Will be set from auth if available
                tenant_id=None,  # Will be set from auth if available
                idempotency_key=request.headers.get("Idempotency-Key"),
                created_at=current_time,
                updated_at=current_time,
            )

            try:
                # Check for duplicate request
                cached_response = await self.deduplicator.check_duplicate(
                    method=request.method,
                    path=str(request.url.path),
                    query_params=dict(request.query_params),
                    body=metadata.body,
                    headers=dict(request.headers),
                    idempotency_key=metadata.idempotency_key,
                )
                if cached_response:
                    logger.info(f"Duplicate request detected: {request_id}")
                    return JSONResponse(content=cached_response)

                # Create durable request
                durable_request = DurableRequest(
                    metadata=metadata,
                )
                self.active_requests[request_id] = durable_request

                # Emit start event
                await self.event_store.append(
                    EventType.REQUEST_STARTED,
                    request_id,
                    {
                        "path": metadata.path,
                        "method": metadata.method,
                        "timestamp": metadata.created_at.isoformat(),
                    },
                )

                # Create checkpoint before processing
                from ..middleware.gateway.durable_request import Checkpoint

                checkpoint = Checkpoint(
                    checkpoint_id=f"ckpt_{request_id}_init",
                    request_id=request_id,
                    sequence=0,
                    name="request_initialized",
                    state=RequestState.INITIALIZED,
                    data={
                        "metadata": {
                            "request_id": metadata.request_id,
                            "method": metadata.method,
                            "path": metadata.path,
                            "client_ip": metadata.client_ip,
                            "created_at": metadata.created_at.isoformat(),
                        },
                        "created_at": datetime.now(UTC).isoformat(),
                    },
                    workflow_state=None,
                    created_at=datetime.now(UTC),
                    size_bytes=0,
                )
                await self.checkpoint_manager.save_checkpoint(checkpoint)

                # Process request
                response = await call_next(request)

                # Update state to completed
                durable_request.state = RequestState.COMPLETED

                # Cache response for deduplication
                if response.status_code < 400:
                    # Only cache successful responses
                    response_body = b"".join(
                        [chunk async for chunk in response.body_iterator]
                    )
                    try:
                        response_data = {"content": response_body.decode()}
                    except UnicodeDecodeError:
                        response_data = {"content": response_body.hex()}

                    await self.deduplicator.cache_response(
                        method=metadata.method,
                        path=metadata.path,
                        query_params=metadata.query_params,
                        body=metadata.body,
                        headers=metadata.headers,
                        idempotency_key=metadata.idempotency_key,
                        response_data=response_data,
                        status_code=response.status_code,
                        response_headers=dict(response.headers),
                    )

                    # Recreate response with new body
                    response = Response(
                        content=response_body,
                        status_code=response.status_code,
                        headers=dict(response.headers),
                        media_type=response.media_type,
                    )

                # Emit completion event
                await self.event_store.append(
                    EventType.REQUEST_COMPLETED,
                    request_id,
                    {
                        "status_code": response.status_code,
                        "timestamp": datetime.now(UTC).isoformat(),
                    },
                )

                return response

            except Exception as e:
                # Update state to failed
                if request_id in self.active_requests:
                    self.active_requests[request_id].state = RequestState.FAILED

                # Emit failure event
                await self.event_store.append(
                    EventType.REQUEST_FAILED,
                    request_id,
                    {
                        "error": str(e),
                        "timestamp": datetime.now(UTC).isoformat(),
                    },
                )

                logger.error(f"Request {request_id} failed: {e}")
                raise

            finally:
                # Clean up active request
                if request_id in self.active_requests:
                    del self.active_requests[request_id]

    def _should_use_durability(self, request: Request) -> bool:
        """Determine if request should use durability features."""
        if not self.enable_durability:
            return False

        if self.durability_opt_in:
            # Check for durability header
            return request.headers.get("X-Durable", "").lower() == "true"
        else:
            # Durability enabled by default, check for opt-out
            return request.headers.get("X-Durable", "").lower() != "false"

    def _register_durability_endpoints(self):
        """Register durability-specific endpoints."""

        @self.app.get("/durability/status")
        async def durability_status():
            """Get durability system status."""
            return {
                "enabled": self.enable_durability,
                "opt_in": self.durability_opt_in,
                "active_requests": len(self.active_requests),
                "checkpoint_count": len(
                    getattr(self.checkpoint_manager, "_memory_checkpoints", [])
                ),
                "event_count": self.event_store.event_count,
            }

        @self.app.get("/durability/requests")
        async def list_active_requests():
            """List currently active durable requests."""
            return {
                request_id: {
                    "state": req.state.value,
                    "metadata": {
                        "request_id": req.metadata.request_id,
                        "method": req.metadata.method,
                        "path": req.metadata.path,
                        "client_ip": req.metadata.client_ip,
                        "created_at": req.metadata.created_at.isoformat(),
                    },
                    "created_at": req.metadata.created_at.isoformat(),
                }
                for request_id, req in self.active_requests.items()
            }

        @self.app.get("/durability/requests/{request_id}")
        async def get_request_status(request_id: str):
            """Get status of a specific request."""
            if request_id in self.active_requests:
                req = self.active_requests[request_id]
                return {
                    "request_id": request_id,
                    "state": req.state.value,
                    "metadata": {
                        "request_id": req.metadata.request_id,
                        "method": req.metadata.method,
                        "path": req.metadata.path,
                        "created_at": req.metadata.created_at.isoformat(),
                        "updated_at": req.metadata.updated_at.isoformat(),
                    },
                    "active": True,
                }

            # Check checkpoint storage
            checkpoint = await self.checkpoint_manager.load_latest_checkpoint(
                request_id
            )
            if checkpoint:
                return {
                    "request_id": request_id,
                    "state": checkpoint.state.value,
                    "metadata": checkpoint.data.get("metadata", {}),
                    "active": False,
                }

            raise HTTPException(status_code=404, detail="Request not found")

        @self.app.post("/durability/requests/{request_id}/recover")
        async def recover_request(request_id: str):
            """Recover a failed request by resuming from its latest checkpoint.

            Creates a new :class:`DurableRequest`, restores the checkpoint
            state (including the :class:`ExecutionTracker` so already-completed
            nodes are skipped), and re-executes.
            """
            checkpoint = await self.checkpoint_manager.load_latest_checkpoint(
                request_id
            )
            if not checkpoint:
                raise HTTPException(
                    status_code=404, detail="Request checkpoint not found"
                )

            # Build a DurableRequest and attach the checkpoint manager
            durable_request = DurableRequest(
                request_id=request_id,
                checkpoint_manager=self.checkpoint_manager,
            )
            self.active_requests[request_id] = durable_request

            try:
                result = await durable_request.resume(
                    checkpoint_id=checkpoint.checkpoint_id
                )
                return {
                    "message": f"Recovery completed for request {request_id}",
                    "checkpoint": checkpoint.to_dict(),
                    "result": result,
                }
            except Exception as e:
                logger.error("Recovery failed for request %s: %s", request_id, e)
                raise HTTPException(
                    status_code=500,
                    detail=f"Recovery failed: {e}",
                )
            finally:
                self.active_requests.pop(request_id, None)

        @self.app.get("/durability/events")
        async def list_events(limit: int = 100, offset: int = 0):
            """List recent durability events."""
            # Access in-memory event stream directly for listing
            async with self.event_store._stream_lock:
                all_events = list(self.event_store._event_stream)
            events = all_events[offset : offset + limit]
            return {
                "events": [
                    {
                        "type": event.event_type.value,
                        "data": event.data,
                        "timestamp": event.timestamp.isoformat(),
                        "event_id": event.event_id,
                    }
                    for event in events
                ],
                "total": len(all_events),
                "limit": limit,
                "offset": offset,
            }

    async def cleanup_completed_requests(self, max_age_hours: int = 24):
        """Clean up old completed request data."""
        cutoff_time = datetime.now(UTC).timestamp() - (max_age_hours * 3600)

        # Clean up checkpoints - using garbage collection method
        await self.checkpoint_manager._garbage_collection()

        # Clean up cached responses - using internal cleanup
        await self.deduplicator._cleanup_expired()

        logger.info(f"Cleaned up durability data older than {max_age_hours} hours")

    def _register_root_endpoints(self):
        """Override to add durability info to root endpoint."""
        super()._register_root_endpoints()

        # Override the root endpoint to include durability info
        @self.app.get("/")
        async def root():
            """Server information with durability details."""
            base_info = {
                "name": self.app.title,
                "version": self.app.version,
                "workflows": list(self.workflows.keys()),
                "mcp_servers": list(self.mcp_servers.keys()),
                "type": "durable_workflow_server",
            }

            # Add durability info
            base_info["durability"] = {
                "enabled": self.enable_durability,
                "opt_in": self.durability_opt_in,
                "features": [
                    "request_checkpointing",
                    "automatic_deduplication",
                    "event_sourcing",
                    "request_recovery",
                ],
            }

            return base_info
