# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Models for the diagrams-mcp-server."""

from enum import Enum
from pydantic import BaseModel, Field, field_validator
from typing import Dict, List, Literal, Optional


class DiagramType(str, Enum):
    """Enum for supported diagram types."""

    AWS = 'aws'
    GCP = 'gcp'
    AZURE = 'azure'
    K8S = 'k8s'
    ONPREM = 'onprem'
    HYBRID = 'hybrid'
    MULTICLOUD = 'multicloud'
    SEQUENCE = 'sequence'
    FLOW = 'flow'
    CLASS = 'class'
    CUSTOM = 'custom'
    ALL = 'all'


class DiagramGenerateRequest(BaseModel):
    """Request model for diagram generation."""

    code: str = Field(..., description='Python code string using the diagrams package DSL')
    filename: Optional[str] = Field(
        None,
        description='Output filename (without extension). If not provided, a random name will be generated.',
    )
    timeout: int = Field(90, description='Timeout in seconds for diagram generation', ge=1, le=300)
    workspace_dir: Optional[str] = Field(
        None,
        description='The user\'s current workspace directory. If provided, diagrams will be saved to a "generated-diagrams" subdirectory.',
    )

    @field_validator('code')
    @classmethod
    def validate_code(cls, v):
        """Validate that the code contains a Diagram definition."""
        if 'Diagram(' not in v:
            raise ValueError('Code must contain a Diagram definition')
        return v


class DiagramExampleRequest(BaseModel):
    """Request model for diagram examples."""

    diagram_type: DiagramType = Field(
        DiagramType.ALL,
        description='Type of diagram example to return',
    )


class DiagramGenerateResponse(BaseModel):
    """Response model for diagram generation."""

    status: Literal['success', 'error']
    path: Optional[str] = None
    message: str
    image_data: Optional[str] = Field(None, description='Base64-encoded PNG image data')
    mime_type: Optional[str] = Field(None, description='MIME type of the image')


class DiagramExampleResponse(BaseModel):
    """Response model for diagram examples."""

    examples: Dict[str, str]


class DiagramIconsRequest(BaseModel):
    """Request model for listing available diagram icons."""

    provider_filter: Optional[str] = Field(
        None, description='Filter icons by provider name (e.g., "aws", "gcp", "k8s")'
    )
    service_filter: Optional[str] = Field(
        None, description='Filter icons by service name (e.g., "compute", "database", "network")'
    )


class DiagramIconsResponse(BaseModel):
    """Response model for listing available diagram icons."""

    providers: Dict[str, Dict[str, List[str]]]
    filtered: bool = False
    filter_info: Optional[Dict[str, str]] = None
