# ABOUTME: Service for Fabric Livy/Spark operations.
# ABOUTME: Handles Spark session and statement management.
"""Service for Microsoft Fabric Livy integration.

This service provides Spark session and statement management through
the Fabric Livy API for interactive PySpark, Scala, and SparkR execution.
"""

import json
import logging
import time
from typing import Any, Dict, Optional

from ..client.exceptions import (
    FabricAPIError,
    FabricError,
    FabricLivyError,
    FabricLivySessionError,
    FabricLivyStatementError,
    FabricLivyTimeoutError,
)
from ..client.http_client import FabricClient

logger = logging.getLogger(__name__)


class FabricLivyService:
    """Service for Microsoft Fabric Livy operations.
    
    This service manages Spark sessions and code execution through the Livy API.
    Sessions can take 6+ minutes to start on first creation as Spark initializes.
    
    Example:
        ```python
        from ms_fabric_mcp_server import FabricConfig, FabricClient
        from ms_fabric_mcp_server.services import FabricLivyService
        
        config = FabricConfig.from_environment()
        client = FabricClient(config)
        livy_service = FabricLivyService(client, config)
        
        # Create a session (waits for availability)
        session = livy_service.create_session(
            workspace_id="workspace-123",
            lakehouse_id="lakehouse-456",
            kind="pyspark",
            with_wait=True
        )
        
        # Run code
        result = livy_service.run_statement(
            workspace_id="workspace-123",
            lakehouse_id="lakehouse-456",
            session_id=session["id"],
            code="df = spark.read.parquet('Files/data.parquet'); df.count()",
            wait=True
        )
        
        # Clean up
        livy_service.close_session(
            workspace_id="workspace-123",
            lakehouse_id="lakehouse-456",
            session_id=session["id"]
        )
        ```
    """
    
    def __init__(
        self,
        client: FabricClient,
        session_wait_timeout: int = 600,
        statement_wait_timeout: int = 600,
        poll_interval: float = 5.0,
    ):
        """Initialize the Livy service.
        
        Args:
            client: FabricClient instance for API communication
            session_wait_timeout: Maximum seconds to wait for session availability (default: 600)
            statement_wait_timeout: Maximum seconds to wait for statement completion (default: 600)
            poll_interval: Seconds between status polls (default: 5.0)
        """
        self.client = client
        self.session_wait_timeout = session_wait_timeout
        self.statement_wait_timeout = statement_wait_timeout
        self.poll_interval = poll_interval
        
        logger.debug("FabricLivyService initialized")

    @staticmethod
    def _extract_fallback_info(session_data: Dict[str, Any]) -> Dict[str, Any]:
        """Extract fallback warnings from Livy session tags."""
        tags = session_data.get("tags") or {}
        fallback_info: Dict[str, Any] = {}
        if "FallbackReasons" in tags:
            fallback_info["fallback_reasons"] = tags["FallbackReasons"]
        if "FallbackMessages" in tags:
            fallback_info["fallback_messages"] = tags["FallbackMessages"]
        return fallback_info
    
    def create_session(
        self,
        workspace_id: str,
        lakehouse_id: str,
        environment_id: Optional[str] = None,
        kind: str = "pyspark",
        conf: Optional[Dict[str, Any]] = None,
        with_wait: bool = True,
        timeout_seconds: Optional[int] = None
    ) -> Dict[str, Any]:
        """Create a new Livy session.
        
        Args:
            workspace_id: Fabric workspace ID
            lakehouse_id: Fabric lakehouse ID
            environment_id: Optional Fabric environment ID for pre-installed libraries
            kind: Session kind - 'pyspark' (default), 'scala', or 'sparkr'
            conf: Additional Spark configuration as key-value pairs
            with_wait: If True, wait for session to become available (default: True)
            timeout_seconds: Maximum time to wait for session availability 
                (default: session_wait_timeout)
            
        Returns:
            Session creation response (with final state if with_wait=True)
            
        Raises:
            FabricLivySessionError: If session creation fails
            FabricLivyTimeoutError: If session doesn't become available within timeout
            FabricAPIError: If API request fails
            
        Note:
            Session creation can take 6+ minutes on first startup as Spark initializes.
            It's recommended to keep with_wait=True to ensure the session is ready for use.
            
        Example:
            ```python
            session = livy_service.create_session(
                workspace_id="abc-123",
                lakehouse_id="def-456",
                environment_id="env-789",  # Optional
                kind="pyspark",
                conf={"spark.executor.memory": "4g"},
                with_wait=True
            )
            
            print(f"Session ID: {session['id']}")
            print(f"State: {session['state']}")
            ```
        """
        logger.info(
            f"Creating Livy session in workspace {workspace_id}, lakehouse {lakehouse_id}"
        )
        
        # Prepare session payload
        payload = {
            "kind": kind,
            "conf": conf or {}
        }
        
        # Add environment configuration if provided
        if environment_id:
            payload["conf"]["spark.fabric.environmentDetails"] = json.dumps({
                "id": environment_id
            })
            logger.debug(f"Using environment ID: {environment_id}")
        
        # Build endpoint URL
        endpoint = (
            f"workspaces/{workspace_id}/lakehouses/{lakehouse_id}/"
            f"livyapi/versions/2023-12-01/sessions"
        )
        
        try:
            response = self.client.make_api_request(
                method="POST",
                endpoint=endpoint,
                payload=payload,
                timeout=60
            )
            
            session_data = response.json()
            session_id = session_data.get("id")
            
            logger.info(f"Successfully created Livy session: {session_id}")
            fallback_info = self._extract_fallback_info(session_data)
            
            # If with_wait is True, poll until session becomes available
            if with_wait:
                logger.debug(f"Waiting for session {session_id} to become available...")
                final_session_data = self.wait_for_session(
                    workspace_id=workspace_id,
                    lakehouse_id=lakehouse_id,
                    session_id=session_id,
                    timeout_seconds=timeout_seconds
                )
                final_fallback_info = self._extract_fallback_info(final_session_data)
                if final_fallback_info:
                    return {**final_session_data, **final_fallback_info}
                if fallback_info:
                    return {**final_session_data, **fallback_info}
                return final_session_data
            
            if fallback_info:
                return {**session_data, **fallback_info}
            return session_data
            
        except FabricAPIError as exc:
            logger.error(f"API error creating Livy session: {exc}")
            raise FabricLivySessionError("new", f"Failed to create session: {exc}")
        except Exception as exc:
            logger.error(f"Unexpected error creating Livy session: {exc}")
            raise FabricLivyError(f"Unexpected error: {exc}")
    
    def list_sessions(
        self,
        workspace_id: str,
        lakehouse_id: str,
        top: Optional[int] = None,
        continuation_token: Optional[str] = None
    ) -> Dict[str, Any]:
        """List all Livy sessions in a workspace/lakehouse.
        
        Args:
            workspace_id: Fabric workspace ID
            lakehouse_id: Fabric lakehouse ID
            top: Maximum number of sessions to return (pagination)
            continuation_token: Token for pagination continuation
            
        Returns:
            Dictionary containing sessions list and pagination info
            
        Raises:
            FabricLivyError: If listing sessions fails
            FabricAPIError: If API request fails
            
        Example:
            ```python
            sessions = livy_service.list_sessions(
                workspace_id="abc-123",
                lakehouse_id="def-456"
            )
            
            for session in sessions.get("sessions", []):
                print(f"Session {session['id']}: {session['state']}")
            ```
        """
        logger.debug(
            f"Listing Livy sessions for workspace {workspace_id}, lakehouse {lakehouse_id}"
        )
        
        # Build endpoint URL
        endpoint = (
            f"workspaces/{workspace_id}/lakehouses/{lakehouse_id}/"
            f"livyapi/versions/2023-12-01/sessions"
        )
        
        # Build query parameters
        params = {}
        if top is not None:
            params["$top"] = top
        if continuation_token:
            params["continuationToken"] = continuation_token
        
        # Add query parameters to endpoint if any
        if params:
            param_str = "&".join([f"{k}={v}" for k, v in params.items()])
            endpoint = f"{endpoint}?{param_str}"
        
        try:
            response = self.client.make_api_request(
                method="GET",
                endpoint=endpoint,
                timeout=30
            )
            
            sessions_data = response.json()
            session_count = len(sessions_data.get("sessions", []))
            logger.debug(f"Retrieved {session_count} Livy sessions")
            
            return sessions_data
            
        except FabricAPIError as exc:
            logger.error(f"API error listing Livy sessions: {exc}")
            raise FabricLivyError(f"Failed to list sessions: {exc}")
        except Exception as exc:
            logger.error(f"Unexpected error listing Livy sessions: {exc}")
            raise FabricLivyError(f"Unexpected error: {exc}")
    
    def get_session_status(
        self,
        workspace_id: str,
        lakehouse_id: str,
        session_id: str
    ) -> Dict[str, Any]:
        """Get session details and status.
        
        Args:
            workspace_id: Fabric workspace ID
            lakehouse_id: Fabric lakehouse ID
            session_id: Livy session ID
            
        Returns:
            Session details including state, Spark info, and configuration
            
        Raises:
            FabricLivySessionError: If session retrieval fails
            FabricAPIError: If API request fails
            
        Note:
            Session states: 'not_started', 'starting', 'idle', 'busy', 
            'shutting_down', 'error', 'dead', 'killed', 'success'
            
        Example:
            ```python
            status = livy_service.get_session_status(
                workspace_id="abc-123",
                lakehouse_id="def-456",
                session_id="session-789"
            )
            
            print(f"State: {status['state']}")
            if status['state'] == 'idle':
                print("Session ready for statements")
            ```
        """
        logger.debug(f"Getting Livy session {session_id}")
        
        # Build endpoint URL
        endpoint = (
            f"workspaces/{workspace_id}/lakehouses/{lakehouse_id}/"
            f"livyapi/versions/2023-12-01/sessions/{session_id}"
        )
        
        try:
            response = self.client.make_api_request(
                method="GET",
                endpoint=endpoint,
                timeout=30
            )
            
            session_data = response.json()
            logger.debug(f"Session {session_id} state: {session_data.get('state', 'unknown')}")
            return session_data
            
        except FabricAPIError as exc:
            logger.error(f"API error getting Livy session: {exc}")
            raise FabricLivySessionError(session_id, f"Failed to get session: {exc}")
        except Exception as exc:
            logger.error(f"Unexpected error getting Livy session: {exc}")
            raise FabricLivyError(f"Unexpected error: {exc}")
    
    def close_session(
        self,
        workspace_id: str,
        lakehouse_id: str,
        session_id: str
    ) -> Dict[str, Any]:
        """Close (terminate) a Livy session.
        
        Args:
            workspace_id: Fabric workspace ID
            lakehouse_id: Fabric lakehouse ID
            session_id: Livy session ID
            
        Returns:
            Session closure response
            
        Raises:
            FabricLivySessionError: If session closure fails
            FabricAPIError: If API request fails
            
        Example:
            ```python
            livy_service.close_session(
                workspace_id="abc-123",
                lakehouse_id="def-456",
                session_id="session-789"
            )
            ```
        """
        logger.info(f"Closing Livy session {session_id}")
        
        # Build endpoint URL
        endpoint = (
            f"workspaces/{workspace_id}/lakehouses/{lakehouse_id}/"
            f"livyapi/versions/2023-12-01/sessions/{session_id}"
        )
        
        try:
            response = self.client.make_api_request(
                method="DELETE",
                endpoint=endpoint,
                timeout=30
            )
            
            # Handle successful deletion (may not return JSON)
            try:
                result = response.json()
            except Exception:
                result = {
                    "status": response.status_code,
                    "message": "Session closed successfully"
                }
            
            logger.info(f"Successfully closed Livy session {session_id}")
            return result
            
        except FabricAPIError as exc:
            logger.error(f"API error closing Livy session: {exc}")
            raise FabricLivySessionError(session_id, f"Failed to close session: {exc}")
        except Exception as exc:
            logger.error(f"Unexpected error closing Livy session: {exc}")
            raise FabricLivyError(f"Unexpected error: {exc}")
    
    def get_session_log(
        self,
        workspace_id: str,
        lakehouse_id: str,
        session_id: str,
        start: int = 0,
        size: int = 500
    ) -> Dict[str, Any]:
        """Fetch Livy session logs via Spark monitoring APIs.
        
        Args:
            workspace_id: Fabric workspace ID
            lakehouse_id: Fabric lakehouse ID
            session_id: Livy session ID
            start: Byte offset for partial log download (default: 0)
            size: Number of bytes to retrieve when partial download is enabled (default: 500)
            
        Returns:
            Dictionary with log content and metadata.
            
        Raises:
            FabricLivySessionError: If log retrieval fails
            FabricAPIError: If API request fails
            
        Note:
            Use this for debugging session startup issues or Spark driver problems.
            Supports incremental reads with start/size parameters for paging.
            
        Example:
            ```python
            # Get first 100 log lines
            logs = livy_service.get_session_log(
                workspace_id="abc-123",
                lakehouse_id="def-456",
                session_id="session-789",
                start=0,
                size=100
            )
            
            for line in logs.get("log", []):
                print(line)
            ```
        """
        logger.debug(
            f"Getting logs for Livy session {session_id} (start={start}, size={size})"
        )
        
        # Build endpoint URL using Spark monitoring API for Livy logs
        endpoint = (
            f"workspaces/{workspace_id}/lakehouses/{lakehouse_id}/"
            f"livySessions/{session_id}/applications/none/logs?type=livy"
        )
        if size and size > 0:
            endpoint += f"&isDownload=true&isPartial=true&offset={start}&size={size}"
        else:
            endpoint += "&isDownload=true"
        
        try:
            response = self.client.make_api_request(
                method="GET",
                endpoint=endpoint,
                timeout=30
            )
            
            log_content = response.text
            log_size = len(response.content or b"")
            logger.debug(f"Retrieved {log_size} bytes of log content")
            return {
                "status": "success",
                "log_content": log_content,
                "log_size_bytes": log_size,
                "offset": start,
                "size": size,
                "content_type": response.headers.get("Content-Type"),
            }
            
        except FabricAPIError as exc:
            logger.error(f"API error getting session logs: {exc}")
            raise FabricLivySessionError(session_id, f"Failed to get logs: {exc}")
        except Exception as exc:
            logger.error(f"Unexpected error getting session logs: {exc}")
            raise FabricLivyError(f"Unexpected error: {exc}")
    
    def run_statement(
        self,
        workspace_id: str,
        lakehouse_id: str,
        session_id: str,
        code: str,
        kind: str = "pyspark",
        with_wait: bool = True,
        timeout_seconds: Optional[int] = None
    ) -> Dict[str, Any]:
        """Execute code in a Livy session.
        
        Args:
            workspace_id: Fabric workspace ID
            lakehouse_id: Fabric lakehouse ID
            session_id: Livy session ID
            code: Code to execute (PySpark, Scala, or SparkR)
            kind: Statement kind - 'pyspark' (default), 'scala', or 'sparkr'
            with_wait: If True, wait for statement completion (default: True)
            timeout_seconds: Maximum time to wait for statement completion
                (default: statement_wait_timeout)
            
        Returns:
            Statement execution response (with final state if with_wait=True)
            
        Raises:
            FabricLivyStatementError: If statement execution fails
            FabricAPIError: If API request fails
            FabricLivyTimeoutError: If statement doesn't complete within timeout
            
        Note:
            - Use df.show() or df.printSchema() to inspect DataFrames before accessing columns
            - SHOW TABLES returns 'namespace' column, not 'database' in Fabric
            - Avoid direct Row attribute access without schema verification
            - When with_wait=False, returns immediately with statement ID
            
        Example:
            ```python
            result = livy_service.run_statement(
                workspace_id="abc-123",
                lakehouse_id="def-456",
                session_id="session-789",
                code=\"\"\"
                df = spark.read.parquet('Files/data.parquet')
                df.count()
                \"\"\",
                with_wait=True
            )
            
            if result.get("state") == "available":
                output = result.get("output", {})
                print(f"Result: {output.get('data', {}).get('text/plain')}")
            ```
        """
        logger.info(f"Running statement in Livy session {session_id}")
        logger.debug(f"Code to execute: {code[:200]}...")  # Log first 200 chars
        
        # Prepare statement payload
        payload = {
            "code": code,
            "kind": kind
        }
        
        # Build endpoint URL
        endpoint = (
            f"workspaces/{workspace_id}/lakehouses/{lakehouse_id}/"
            f"livyapi/versions/2023-12-01/sessions/{session_id}/statements"
        )
        
        try:
            response = self.client.make_api_request(
                method="POST",
                endpoint=endpoint,
                payload=payload,
                timeout=60
            )
            
            statement_data = response.json()
            statement_id = statement_data.get("id")
            
            logger.info(f"Successfully submitted statement: {statement_id}")
            
            # If with_wait is True, wait for completion
            if with_wait:
                final_statement_data = self.wait_for_statement(
                    workspace_id=workspace_id,
                    lakehouse_id=lakehouse_id,
                    session_id=session_id,
                    statement_id=statement_id,
                    timeout_seconds=timeout_seconds
                )
                return final_statement_data
            
            return statement_data
            
        except FabricLivyTimeoutError:
            raise
        except FabricLivyError:
            raise
        except FabricAPIError as exc:
            logger.error(f"API error running statement: {exc}")
            raise FabricLivyStatementError(
                session_id, "new", f"Failed to run statement: {exc}"
            )
        except Exception as exc:
            logger.error(f"Unexpected error running statement: {exc}")
            raise FabricLivyError(f"Unexpected error: {exc}")
    
    def get_statement_status(
        self,
        workspace_id: str,
        lakehouse_id: str,
        session_id: str,
        statement_id: str
    ) -> Dict[str, Any]:
        """Get the current status and output of a Livy statement.
        
        Args:
            workspace_id: Fabric workspace ID
            lakehouse_id: Fabric lakehouse ID
            session_id: Livy session ID
            statement_id: Statement ID to check
            
        Returns:
            Statement details including state, output, and execution details
            
        Raises:
            FabricLivyStatementError: If statement retrieval fails
            FabricAPIError: If API request fails
            
        Note:
            Statement states: 'waiting', 'running', 'available', 'error', 
            'cancelling', 'cancelled'
            
        Example:
            ```python
            status = livy_service.get_statement_status(
                workspace_id="abc-123",
                lakehouse_id="def-456",
                session_id="session-789",
                statement_id="stmt-123"
            )
            
            print(f"State: {status['state']}")
            if status['state'] == 'available':
                print(f"Output: {status.get('output')}")
            ```
        """
        logger.debug(f"Getting statement {statement_id} in session {session_id}")
        
        # Build endpoint URL
        endpoint = (
            f"workspaces/{workspace_id}/lakehouses/{lakehouse_id}/"
            f"livyapi/versions/2023-12-01/sessions/{session_id}/"
            f"statements/{statement_id}"
        )
        
        try:
            response = self.client.make_api_request(
                method="GET",
                endpoint=endpoint,
                timeout=30
            )
            
            statement_data = response.json()
            state = statement_data.get("state", "unknown")
            logger.debug(f"Statement {statement_id} state: {state}")
            return statement_data
            
        except FabricAPIError as exc:
            logger.error(f"API error getting statement: {exc}")
            raise FabricLivyStatementError(
                session_id, statement_id, f"Failed to get statement: {exc}"
            )
        except Exception as exc:
            logger.error(f"Unexpected error getting statement: {exc}")
            raise FabricLivyError(f"Unexpected error: {exc}")
    
    def cancel_statement(
        self,
        workspace_id: str,
        lakehouse_id: str,
        session_id: str,
        statement_id: str
    ) -> Dict[str, Any]:
        """Cancel a running Livy statement without killing the session.
        
        Args:
            workspace_id: Fabric workspace ID
            lakehouse_id: Fabric lakehouse ID
            session_id: Livy session ID
            statement_id: Statement ID to cancel
            
        Returns:
            Cancellation result (typically {"msg": "canceled"})
            
        Raises:
            FabricLivyStatementError: If statement cancellation fails
            FabricAPIError: If API request fails
            
        Note:
            - Use this to stop long-running statements without terminating the session
            - Only works on statements in 'waiting' or 'running' state
            - Statement will transition to 'cancelling' then 'cancelled' state
            - Session remains available for new statements
            
        Example:
            ```python
            result = livy_service.cancel_statement(
                workspace_id="abc-123",
                lakehouse_id="def-456",
                session_id="session-789",
                statement_id="stmt-123"
            )
            
            print(f"Cancellation: {result.get('msg')}")
            ```
        """
        logger.info(f"Cancelling statement {statement_id} in session {session_id}")
        
        # Build endpoint URL
        endpoint = (
            f"workspaces/{workspace_id}/lakehouses/{lakehouse_id}/"
            f"livyapi/versions/2023-12-01/sessions/{session_id}/"
            f"statements/{statement_id}/cancel"
        )
        
        try:
            response = self.client.make_api_request(
                method="POST",
                endpoint=endpoint,
                payload={},
                timeout=30
            )
            
            cancel_data = response.json()
            logger.info(f"Successfully cancelled statement {statement_id}")
            return cancel_data
            
        except FabricAPIError as exc:
            logger.error(f"API error cancelling statement: {exc}")
            raise FabricLivyStatementError(
                session_id, statement_id, f"Failed to cancel statement: {exc}"
            )
        except Exception as exc:
            logger.error(f"Unexpected error cancelling statement: {exc}")
            raise FabricLivyError(f"Unexpected error: {exc}")
    
    def wait_for_session(
        self,
        workspace_id: str,
        lakehouse_id: str,
        session_id: str,
        timeout_seconds: Optional[int] = None,
        poll_interval: Optional[float] = None
    ) -> Dict[str, Any]:
        """Poll session until it becomes available.
        
        Args:
            workspace_id: Fabric workspace ID
            lakehouse_id: Fabric lakehouse ID
            session_id: Livy session ID
            timeout_seconds: Maximum time to wait (default: session_wait_timeout)
            poll_interval: Polling interval in seconds (default: self.poll_interval)
            
        Returns:
            Final session data when available
            
        Raises:
            FabricLivyTimeoutError: If session doesn't become available within timeout
            FabricLivySessionError: If session fails or polling fails
            
        Example:
            ```python
            session = livy_service.wait_for_session(
                workspace_id="abc-123",
                lakehouse_id="def-456",
                session_id="session-789",
                timeout_seconds=900  # 15 minutes
            )
            
            print(f"Session ready: {session['state']}")
            ```
        """
        timeout = timeout_seconds or self.session_wait_timeout
        interval = poll_interval or self.poll_interval
        
        logger.info(
            f"Waiting for session {session_id} to become available (timeout={timeout}s)"
        )
        
        start_time = time.time()
        # Allow transient 404s during initial session creation (Fabric API lag)
        max_404_retries = 5
        consecutive_404_count = 0
        
        while True:
            elapsed = time.time() - start_time
            
            # Check timeout
            if elapsed > timeout:
                logger.error(f"Session {session_id} timed out after {timeout}s")
                raise FabricLivyTimeoutError("session creation", timeout)
            
            try:
                # Get current session status
                session_data = self.get_session_status(
                    workspace_id, lakehouse_id, session_id
                )
                
                # Reset 404 counter on successful response
                consecutive_404_count = 0
                
                state = session_data.get("state", "unknown")
                
                # Check if session is available
                if state == "idle":
                    logger.info(f"Session {session_id} is now available")
                    return session_data
                elif state in ["error", "dead", "killed"]:
                    error_msg = session_data.get("log", ["Session failed"])
                    if isinstance(error_msg, list) and error_msg:
                        error_msg = error_msg[-1]  # Get last log entry
                    logger.error(f"Session {session_id} failed: {error_msg}")
                    raise FabricLivySessionError(
                        session_id, f"Session {state}: {error_msg}"
                    )
                
                # Log progress for long-running session creation
                if elapsed % 30 < interval:  # Log every 30 seconds
                    logger.debug(
                        f"Session {session_id} still starting "
                        f"(state={state}, elapsed={elapsed:.1f}s)"
                    )
                
            except FabricLivyTimeoutError:
                raise
            except FabricLivySessionError as exc:
                # Check if this is a transient 404 (session not yet visible)
                if "404" in str(exc) and consecutive_404_count < max_404_retries:
                    consecutive_404_count += 1
                    logger.warning(
                        f"Session {session_id} returned 404, retrying "
                        f"({consecutive_404_count}/{max_404_retries})"
                    )
                else:
                    raise
            except Exception as exc:
                logger.error(f"Error polling session {session_id}: {exc}")
                raise FabricLivySessionError(session_id, f"Polling error: {exc}")
            
            # Wait before next poll
            time.sleep(interval)
    
    def wait_for_statement(
        self,
        workspace_id: str,
        lakehouse_id: str,
        session_id: str,
        statement_id: str,
        timeout_seconds: Optional[int] = None,
        poll_interval: Optional[float] = None
    ) -> Dict[str, Any]:
        """Poll statement until completion.
        
        Args:
            workspace_id: Fabric workspace ID
            lakehouse_id: Fabric lakehouse ID
            session_id: Livy session ID
            statement_id: Statement ID
            timeout_seconds: Maximum time to wait (default: statement_wait_timeout)
            poll_interval: Polling interval in seconds (default: self.poll_interval)
            
        Returns:
            Final statement data when completed
            
        Raises:
            FabricLivyTimeoutError: If statement doesn't complete within timeout
            FabricLivyStatementError: If statement fails or polling fails
            
        Example:
            ```python
            result = livy_service.wait_for_statement(
                workspace_id="abc-123",
                lakehouse_id="def-456",
                session_id="session-789",
                statement_id="stmt-123",
                timeout_seconds=300
            )
            
            if result['state'] == 'available':
                print("Statement completed successfully")
            ```
        """
        timeout = timeout_seconds or self.statement_wait_timeout
        interval = poll_interval or self.poll_interval
        
        logger.info(
            f"Waiting for statement {statement_id} to complete (timeout={timeout}s)"
        )
        
        start_time = time.time()
        
        while True:
            elapsed = time.time() - start_time
            
            # Check timeout
            if elapsed > timeout:
                logger.error(f"Statement {statement_id} timed out after {timeout}s")
                raise FabricLivyTimeoutError("statement execution", timeout)
            
            try:
                # Get current statement status
                statement_data = self.get_statement_status(
                    workspace_id, lakehouse_id, session_id, statement_id
                )
                
                state = statement_data.get("state", "unknown")
                
                # Check if statement is complete
                if state == "available":
                    logger.info(f"Statement {statement_id} completed successfully")
                    return statement_data
                elif state in ["error", "cancelled"]:
                    error_msg = statement_data.get("output", {}).get(
                        "evalue", "Statement failed"
                    )
                    logger.error(f"Statement {statement_id} failed: {error_msg}")
                    raise FabricLivyStatementError(
                        session_id, statement_id, f"Statement {state}: {error_msg}"
                    )
                
                # Log progress
                if elapsed % 30 < interval:  # Log every 30 seconds
                    logger.debug(
                        f"Statement {statement_id} still running "
                        f"(state={state}, elapsed={elapsed:.1f}s)"
                    )
                
            except (FabricLivyTimeoutError, FabricLivyStatementError):
                raise
            except Exception as exc:
                logger.error(f"Error polling statement {statement_id}: {exc}")
                raise FabricLivyStatementError(
                    session_id, statement_id, f"Polling error: {exc}"
                )
            
            # Wait before next poll
            time.sleep(interval)
