from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, Any
from enum import Enum
from pydantic import BaseModel, Field
from pydantic import model_validator

class SshAuthType(str, Enum):
    KEYFILE  = "keyfile"
    PASSWORD = "password"

class CommonSshAuthConfig(BaseModel):
    type: SshAuthType = Field(..., description="Type of SSH authentication to use.")
    username: str = Field(..., description="Username for the SSH connection.")

class SshKeyfileAuthConfig(CommonSshAuthConfig):
    type: Literal[SshAuthType.KEYFILE]
    keyfile: str = Field(..., description="Path to the private key file for SSH authentication.")

class SshPasswordAuthConfig(CommonSshAuthConfig):
    type: Literal[SshAuthType.PASSWORD]
    password: str = Field(..., description="Password for SSH authentication.")

SshAuthConfig = Annotated[
    Union[ 
        SshKeyfileAuthConfig,
        SshPasswordAuthConfig,
    ],
    Field(discriminator="type")
]

class SshConnectionConfig(BaseModel):
    host: str = Field(..., description="Host address of the SSH server to connect to.")
    port: int = Field(default=22, ge=1, le=65535, description="Port number used to connect to the SSH server.")
    auth: SshAuthConfig = Field(..., description="SSH authentication configuration.")
    keepalive_interval: Union[str, int, float] = Field(default="10s", description="SSH keepalive interval. Set to '0s' to disable keepalive.")
    watch_interval: Union[str, int, float] = Field(default="5s", description="Interval between SSH connection health checks.")
    retry_interval: Union[str, int, float] = Field(default="5s", description="Interval between reconnection attempts.")
    max_retry_count: int = Field(default=0, ge=0, description="Maximum number of reconnection attempts. 0 for unlimited.")
