from typing import Any, Dict, List, Optional, Union
import json

import mcp.types as types
from prefect import get_client

from .server import mcp


@mcp.tool
async def get_variables(
    limit: Optional[int] = None,
    offset: Optional[int] = None,
    name: Optional[str] = None,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Get a list of variables with optional filtering.
    
    Args:
        limit: Maximum number of variables to return
        offset: Number of variables to skip
        name: Filter by name pattern
        
    Returns:
        A list of variables with their details
    """
    async with get_client() as client:
        # Build filter parameters
        filters = {}
        if name:
            filters["name"] = {"like_": f"%{name}%"}
        
        variables = await client.read_variables(
            limit=limit,
            # prefect 3.3.3 doesn't have these
            # offset=offset,
            # **filters
        )
        
        variables_result = {
            "variables": [variable.model_dump() for variable in variables]
        }
        
        return [types.TextContent(type="text", text=str(variables_result))]


@mcp.tool
async def get_variable(
    name: str,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Get a variable by name.
    
    Args:
        name: The variable name
        
    Returns:
        Variable details
    """
    async with get_client() as client:
        variable = await client.read_variable(name)
        
        return [types.TextContent(type="text", text=str(variable.model_dump()))]


@mcp.tool
async def create_variable(
    name: str,
    value: Any,  # Change type to Any to support different value types
    tags: Optional[List[str]] = None,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Create a variable.
    
    Args:
        name: The variable name
        value: The variable value (can be string, dict, list, etc.)
        tags: Optional tags
        
    Returns:
        Details of the created variable
    """
    try:
        async with get_client() as client:
            # Import the VariableCreate model
            from prefect.client.schemas.actions import VariableCreate
            
            # Create the proper variable object
            variable_create = VariableCreate(
                name=name,
                value=value,  # Pass value directly, no parsing needed
                tags=tags or []
            )
            
            # Use the variable object with the client
            variable = await client.create_variable(variable=variable_create)
            
            variable_result = {"variable": variable.model_dump()}
            return [types.TextContent(type="text", text=str(variable_result))]
    except Exception as e:
        return [types.TextContent(type="text", text=str({"error": str(e)}))]


@mcp.tool
async def update_variable(
    name: str,
    value: Optional[str] = None,
    tags: Optional[List[str]] = None,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Update a variable.
    
    Args:
        name: The variable name
        value: New value
        tags: New tags
        
    Returns:
        Details of the updated variable
    """
    async with get_client() as client:
        # Prepare update data
        update_data = {}
        if value is not None:
            # Parse value if it's a valid JSON
            try:
                parsed_value = json.loads(value)
            except json.JSONDecodeError:
                # If it's not valid JSON, use the string as-is
                parsed_value = value
                
            update_data["value"] = parsed_value
        if tags is not None:
            update_data["tags"] = tags
        
        updated_variable = await client.update_variable(
            name=name,
            **update_data
        )
        
        return [types.TextContent(type="text", text=str(updated_variable.model_dump()))]


@mcp.tool
async def delete_variable(
    name: str,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Delete a variable by name.
    
    Args:
        name: The variable name
        
    Returns:
        Confirmation message
    """
    async with get_client() as client:
        await client.delete_variable_by_name(name)
        
        return [types.TextContent(type="text", text=f"Variable '{name}' deleted successfully.")]
