"""Mock factories and test data generators for ms-fabric-mcp-server tests."""

from typing import Dict, Any, List
from unittest.mock import Mock
import base64


class MockResponseFactory:
    """Factory for creating mock HTTP responses."""
    
    @staticmethod
    def success(data: Dict[str, Any] = None, status_code: int = 200):
        """Create successful response."""
        response = Mock()
        response.status_code = status_code
        response.ok = True
        response.json.return_value = data or {}
        response.text = str(data) if data else ""
        response.headers = {}
        return response
    
    @staticmethod
    def error(status_code: int, message: str = "Error"):
        """Create error response."""
        response = Mock()
        response.status_code = status_code
        response.ok = False
        response.json.return_value = {
            "error": {
                "code": f"Error{status_code}",
                "message": message
            }
        }
        response.text = message
        response.headers = {}
        return response
    
    @staticmethod
    def rate_limit(retry_after: int = 60):
        """Create rate limit response."""
        response = Mock()
        response.status_code = 429
        response.ok = False
        response.headers = {"Retry-After": str(retry_after)}
        response.json.return_value = {
            "error": {
                "code": "TooManyRequests",
                "message": "Rate limit exceeded"
            }
        }
        response.text = "Rate limit exceeded"
        return response


class FabricDataFactory:
    """Factory for generating test data matching Fabric API schemas."""
    
    @staticmethod
    def workspace(
        workspace_id: str = "ws-123",
        name: str = "Test Workspace",
        capacity_id: str = None
    ) -> Dict[str, Any]:
        """Generate workspace data."""
        data = {
            "id": workspace_id,
            "displayName": name,
            "type": "Workspace",
            "description": f"Test workspace {name}"
        }
        if capacity_id:
            data["capacityId"] = capacity_id
        return data
    
    @staticmethod
    def workspace_list(count: int = 3) -> Dict[str, Any]:
        """Generate workspace list response."""
        return {
            "value": [
                FabricDataFactory.workspace(f"ws-{i}", f"Workspace {i}")
                for i in range(count)
            ]
        }
    
    @staticmethod
    def item(
        item_id: str = "item-123",
        name: str = "Test Item",
        item_type: str = "Notebook",
        workspace_id: str = "ws-123",
        folder_id: str | None = None,
    ) -> Dict[str, Any]:
        """Generate item data."""
        data = {
            "id": item_id,
            "displayName": name,
            "type": item_type,
            "workspaceId": workspace_id,
            "description": f"Test {item_type}"
        }
        if folder_id:
            data["folderId"] = folder_id
        return data
    
    @staticmethod
    def item_list(count: int = 5, item_type: str = None) -> Dict[str, Any]:
        """Generate item list response."""
        items = []
        types = [item_type] if item_type else ["Notebook", "Lakehouse", "Warehouse", "Report", "Dashboard"]
        for i in range(count):
            item_type_val = types[i % len(types)]
            items.append(FabricDataFactory.item(f"item-{i}", f"Item {i}", item_type_val))
        return {"value": items}
    
    @staticmethod
    def notebook_definition(content: str = '{"cells":[]}') -> Dict[str, Any]:
        """Generate notebook definition."""
        encoded_content = base64.b64encode(content.encode()).decode()
        return {
            "format": "ipynb",
            "definition": {
                "parts": [
                    {
                        "path": "notebook-content.py",
                        "payload": encoded_content,
                        "payloadType": "InlineBase64"
                    }
                ]
            }
        }
    
    @staticmethod
    def livy_session(
        session_id: int = 1,
        state: str = "idle",
        kind: str = "pyspark"
    ) -> Dict[str, Any]:
        """Generate Livy session data."""
        return {
            "id": session_id,
            "name": None,
            "appId": f"app-{session_id}" if state not in ["not_started", "starting"] else None,
            "owner": None,
            "proxyUser": None,
            "state": state,
            "kind": kind,
            "appInfo": {
                "driverLogUrl": None,
                "sparkUiUrl": None
            },
            "log": []
        }
    
    @staticmethod
    def livy_session_list(count: int = 3) -> Dict[str, Any]:
        """Generate Livy session list."""
        return {
            "from": 0,
            "total": count,
            "sessions": [
                FabricDataFactory.livy_session(i, "idle")
                for i in range(count)
            ]
        }
    
    @staticmethod
    def livy_statement(
        statement_id: int = 0,
        state: str = "available",
        output: Dict[str, Any] = None
    ) -> Dict[str, Any]:
        """Generate Livy statement data."""
        if output is None:
            output = {
                "status": "ok",
                "execution_count": statement_id,
                "data": {"text/plain": "result"}
            }
        return {
            "id": statement_id,
            "code": "print('test')",
            "state": state,
            "output": output,
            "progress": 1.0 if state == "available" else 0.0
        }
    
    @staticmethod
    def job_instance(
        instance_id: str = "job-123",
        status: str = "Completed",
        item_id: str = "item-123"
    ) -> Dict[str, Any]:
        """Generate job instance data."""
        return {
            "id": instance_id,
            "itemId": item_id,
            "jobType": "RunNotebook",
            "invokeType": "Manual",
            "status": status,
            "startTimeUtc": "2025-10-14T10:00:00Z",
            "endTimeUtc": "2025-10-14T10:05:00Z" if status in ["Completed", "Failed"] else None
        }
    
    @staticmethod
    def sql_query_result(columns: List[str] = None, rows: List[List[Any]] = None) -> tuple:
        """Generate SQL query result data."""
        if columns is None:
            columns = ["id", "name", "value"]
        if rows is None:
            rows = [
                [1, "item1", 100],
                [2, "item2", 200],
            ]
        return columns, rows


class ServiceMockFactory:
    """Factory for creating mock service instances."""
    
    @staticmethod
    def workspace_service(mock_client=None):
        """Create mock workspace service."""
        from unittest.mock import Mock
        service = Mock()
        service.client = mock_client
        service.list_workspaces = Mock()
        service.get_workspace_by_id = Mock()
        service.get_workspace_by_name = Mock()
        service.create_workspace = Mock()
        service.delete_workspace = Mock()
        return service
    
    @staticmethod
    def item_service(mock_client=None, mock_workspace_service=None):
        """Create mock item service."""
        from unittest.mock import Mock
        service = Mock()
        service.client = mock_client
        service.workspace_service = mock_workspace_service
        service.list_items = Mock()
        service.get_item_by_name = Mock()
        service.delete_item = Mock()
        return service
    
    @staticmethod
    def notebook_service(mock_client=None, mock_workspace_service=None, mock_item_service=None):
        """Create mock notebook service."""
        from unittest.mock import Mock
        service = Mock()
        service.client = mock_client
        service.workspace_service = mock_workspace_service
        service.item_service = mock_item_service
        service.create_notebook = Mock()
        service.get_notebook_definition = Mock()
        service.update_notebook_definition = Mock()
        service.get_notebook_run_details = Mock()
        service.list_notebook_runs = Mock()
        service.get_notebook_driver_logs = Mock()
        return service
    
    @staticmethod
    def job_service(mock_client=None, mock_workspace_service=None):
        """Create mock job service."""
        from unittest.mock import Mock
        service = Mock()
        service.client = mock_client
        service.workspace_service = mock_workspace_service
        service.run_on_demand_job = Mock()
        service.get_job_instance = Mock()
        service.get_operation_result = Mock()
        return service
    
    @staticmethod
    def sql_service(mock_client=None, mock_workspace_service=None):
        """Create mock SQL service."""
        from unittest.mock import Mock
        service = Mock()
        service.client = mock_client
        service.workspace_service = mock_workspace_service
        service.get_sql_endpoint = Mock()
        service.execute_query = Mock()
        service.execute_statement = Mock()
        return service
    
    @staticmethod
    def livy_service(mock_client=None):
        """Create mock Livy service."""
        from unittest.mock import Mock
        service = Mock()
        service.client = mock_client
        service.create_session = Mock()
        service.list_sessions = Mock()
        service.get_session = Mock()
        service.delete_session = Mock()
        service.run_statement = Mock()
        service.get_statement = Mock()
        service.cancel_statement = Mock()
        service.get_session_log = Mock()
        return service
