from api_data_structure import HiveComputeAPI
from utils import logger, create_success_response, handle_exception
from typing import Dict, Any
from fastmcp import FastMCP
from constant import Constants

mcp = FastMCP(
    name="Compute with Hivenet MCP"
)

@mcp.tool()
async def create_compute_instance(name: str = "default", location: str = "uae", config: str = "2x RTX 4090") -> Dict[str, Any]:
    """
    Create a new compute instance with the specified configuration.

    Args:
        name: Name of the instance. Defaults to "default".
        location: Location where the instance will be deployed. Defaults to "uae".
                  Valid locations: france, uae, texas, uae-2
        config: Predefined configuration. Defaults to "1x RTX 4090".
                GPU configs: "1x RTX 4090", "2x RTX 4090", "4x RTX 4090", "8x RTX 4090",
                            "1x RTX 5090", "2x RTX 5090", "4x RTX 5090", "8x RTX 5090"
                vCPU configs: "2vCPU", "4vCPU", "8vCPU", "16vCPU", "32vCPU"

    Returns:
        Dict containing the created instance information.
    """
    token = Constants.HIVE_COMPUTE_DEFAULT_API_TOKEN
    api_handler = HiveComputeAPI(token=token)
    try:
        logger.info(f"Creating compute instance: name={name}, location={location}, config={config}")
        result = api_handler.create_instance(name=name, location=location, config=config)

        if result is None:
            logger.error(f"Failed to create instance - API returned None")
            return {
                "status": "error",
                "error": "Failed to create instance. Check API logs for details."
            }

        logger.info(f"Successfully created instance: {result.id if hasattr(result, 'id') else 'unknown'}")
        return create_success_response(result)
    except Exception as e:
        logger.error(f"Exception creating instance: {str(e)}")
        return handle_exception(e, "create_compute_instance")

@mcp.tool()
async def list_all_compute_instances(category: str = None) -> Dict[str, Any]:
    """
    List all instances belonging to the user, organized into 4 categories: RUNNING, STOPPED, ERRORED, and TERMINATED.
    Shows basic information about each instance, including ID, name, status, location, spending, and resource allocation.
    Spending information includes hourly price and total in EUR spent so far.

    Args:
        category: Optional filter to return only instances from a specific category.
                 Valid values: "RUNNING", "STOPPED", "ERRORED", "TERMINATED".
                 If not provided, returns all categories.

    Returns:
        Dict containing instances. If category is specified, returns only instances from that category.
        If category is not specified, returns all instances organized by status categories.
    """
    token = Constants.HIVE_COMPUTE_DEFAULT_API_TOKEN
    api_handler = HiveComputeAPI(token=token)
    try:
        logger.info(f"Listing all compute instances for token: {token}, category filter: {category}")
        all_instances = api_handler.get_all_instances()

        # Categorize instances into 4 groups
        categorized = {
            "RUNNING": [],
            "STOPPED": [],
            "ERRORED": [],
            "TERMINATED": []
        }

        for inst in all_instances:
            status = inst.status_string
            # Map statuses to categories
            if status in ["RUNNING"]:
                categorized["RUNNING"].append(inst)
            elif status in ["STOPPED"]:
                categorized["STOPPED"].append(inst)
            elif status == "ERRORED":
                categorized["ERRORED"].append(inst)
            elif status in ["TERMINATED"]:
                categorized["TERMINATED"].append(inst)

        # If category filter is specified, return only that category
        if category and category.upper() in categorized:
            return create_success_response(categorized[category.upper()])

        # Otherwise return all categories
        return create_success_response(categorized)
    except Exception as e:
        return handle_exception(e, "list_all_compute_instances")