"""
Pydantic models for ETL pipeline configuration.

Replaces JSON Schema files with type-safe Python models for:
- Extract configuration (HTTP, file, database, cloud storage)
- Transform configuration (simple ops, JSONata, custom function)
- Load configuration (Postgres, SQLite, file, cloud storage)
- Settings (shared pipeline config)
- Validation (source/target schema validation)
- DLQ (dead letter queue)
"""

from __future__ import annotations

from enum import Enum
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field

# =============================================================================
# ENUMS
# =============================================================================


class ExtractType(str, Enum):
    """Extract source types."""

    HTTP = "http"
    FILE = "file"
    DATABASE = "database"
    CLOUD_STORAGE = "cloud_storage"


class LoadType(str, Enum):
    """Load destination types."""

    POSTGRES = "postgres"
    POSTGRESQL = "postgresql"
    DATABASE = "database"
    SQLITE = "sqlite"
    FILE = "file"
    CLOUD_STORAGE = "cloud_storage"


class CloudProvider(str, Enum):
    """Cloud storage providers."""

    S3 = "s3"
    GCS = "gcs"
    AZURE = "azure"


class WriteMethod(str, Enum):
    """Database write methods."""

    INSERT = "insert"
    UPSERT = "upsert"
    REPLACE = "replace"
    UPDATE = "update"
    DELETE = "delete"
    TRUNCATE_AND_LOAD = "truncate_and_load"


class FileFormat(str, Enum):
    """File formats."""

    JSON = "json"
    JSONL = "jsonl"
    CSV = "csv"
    TSV = "tsv"
    PARQUET = "parquet"
    EXCEL = "excel"
    XML = "xml"


class WriteMode(str, Enum):
    """File write modes."""

    OVERWRITE = "overwrite"
    APPEND = "append"


class OnErrorAction(str, Enum):
    """Actions on validation error."""

    FAIL = "fail"
    WARN = "warn"
    SKIP = "skip"
    QUARANTINE = "quarantine"


class DLQBackend(str, Enum):
    """DLQ storage backends."""

    FILE = "file"
    DATABASE = "database"
    MEMORY = "memory"


class MetadataStoreType(str, Enum):
    """Metadata store types."""

    POSTGRES = "postgres"
    SQLITE = "sqlite"
    MONGODB = "mongodb"
    REDIS = "redis"
    MEMORY = "memory"


class FilterOperator(str, Enum):
    """Filter comparison operators."""

    EQ = "eq"
    NE = "ne"
    GT = "gt"
    GTE = "gte"
    LT = "lt"
    LTE = "lte"
    IN = "in"
    NOT_IN = "not_in"
    CONTAINS = "contains"
    NOT_CONTAINS = "not_contains"
    IS_NULL = "is_null"
    IS_NOT_NULL = "is_not_null"


class PaginationStrategy(str, Enum):
    """HTTP pagination strategies."""

    PAGE = "page"
    OFFSET = "offset"
    CURSOR = "cursor"
    LINK = "link"


class HTTPMethod(str, Enum):
    """HTTP methods."""

    GET = "GET"
    POST = "POST"
    PUT = "PUT"
    PATCH = "PATCH"
    DELETE = "DELETE"


class ResponseFormat(str, Enum):
    """HTTP response formats."""

    JSON = "json"
    TEXT = "text"
    XML = "xml"


class ConvertType(str, Enum):
    """Type conversion targets."""

    INT = "int"
    INTEGER = "integer"
    FLOAT = "float"
    NUMBER = "number"
    STR = "str"
    STRING = "string"
    BOOL = "bool"
    BOOLEAN = "boolean"
    DATETIME = "datetime"
    DATE = "date"


class JSONataMode(str, Enum):
    """JSONata application mode."""

    RECORD = "record"
    BATCH = "batch"


# =============================================================================
# SHARED CONFIG MODELS
# =============================================================================


class SSHTunnelConfig(BaseModel):
    """SSH tunnel configuration for database connections."""

    model_config = ConfigDict(extra="forbid")

    host: str
    port: int = 22
    username: str
    private_key_path: str | None = None
    password: str | None = None
    remote_host: str | None = None
    remote_port: int | None = None
    enabled: bool = False
    local_port: int = 5433


class DatabaseConfig(BaseModel):
    """Database connection configuration."""

    model_config = ConfigDict(extra="forbid")

    url: str


class CSVOptions(BaseModel):
    """CSV file options."""

    model_config = ConfigDict(extra="forbid")

    delimiter: str = ","
    quote_char: str = '"'
    escape_char: str | None = None
    has_header: bool = True
    include_header: bool = True  # for writing
    encoding: str = "utf-8"


class JSONOptions(BaseModel):
    """JSON file options."""

    model_config = ConfigDict(extra="forbid")

    encoding: str = "utf-8"
    lines: bool = False  # for reading JSONL
    indent: int | None = None  # for writing
    ensure_ascii: bool = False  # for writing


class CloudStorageConfig(BaseModel):
    """Cloud storage configuration."""

    model_config = ConfigDict(extra="forbid")

    provider: CloudProvider
    bucket: str | None = None  # S3, GCS
    container: str | None = None  # Azure
    path: str | None = None
    prefix: str | None = None
    credentials: dict[str, Any] | None = None


# =============================================================================
# DLQ CONFIG
# =============================================================================


class DLQConfig(BaseModel):
    """Dead Letter Queue configuration."""

    model_config = ConfigDict(extra="forbid")

    enabled: bool = True
    backend: DLQBackend = DLQBackend.FILE
    path: str | None = None  # for file backend
    connection_string: str | None = None  # for database backend
    schema_name: str | None = None  # DB schema
    table: str | None = None  # DB table name


# =============================================================================
# VALIDATION CONFIG
# =============================================================================


class ContractRef(BaseModel):
    """Reference to a data contract in database."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["database"] = "database"
    name: str
    version: str | None = None


class ExtractValidationConfig(BaseModel):
    """Validation config for extract stage (source schema)."""

    model_config = ConfigDict(extra="forbid")

    schema_path: str | None = None
    on_error: OnErrorAction = OnErrorAction.FAIL
    dlq: bool | DLQConfig | None = None


class LoadValidationConfig(BaseModel):
    """Validation config for load stage (target contract)."""

    model_config = ConfigDict(extra="forbid")

    contract: str | ContractRef | None = None  # file path or db ref
    use_contract: bool = False  # use settings.yaml contract
    on_error: OnErrorAction = OnErrorAction.FAIL
    dlq: bool | DLQConfig | None = None


# =============================================================================
# METADATA STORE CONFIG
# =============================================================================


class MetadataStoreConfig(BaseModel):
    """Metadata store configuration for database-based contracts."""

    model_config = ConfigDict(extra="forbid")

    type: MetadataStoreType
    connection_string: str


# =============================================================================
# SETTINGS CONFIG
# =============================================================================


class LineageConfig(BaseModel):
    """OpenLineage emission configuration (optional)."""

    model_config = ConfigDict(extra="forbid")

    enabled: bool = False
    namespace: str = "pycharter"
    backend_url: str | None = None


class PluginConfig(BaseModel):
    """Plugin configuration for settings.yaml.

    Maps plugin names to importable references (``module.path:ClassName``).
    """

    model_config = ConfigDict(extra="forbid")

    extractors: dict[str, str] | None = None
    loaders: dict[str, str] | None = None


class IncrementalConfig(BaseModel):
    """Incremental extraction configuration.

    When attached to an extract config, the pipeline will persist a watermark
    value between runs so that only new/updated records are extracted.
    """

    model_config = ConfigDict(extra="forbid")

    enabled: bool = True
    watermark_field: str
    initial_value: str | None = None


class StateStoreConfig(BaseModel):
    """State store configuration for cross-run watermark persistence."""

    model_config = ConfigDict(extra="forbid")

    backend: Literal["file", "sqlite"] = "file"
    path: str = "./.state"


class SettingsConfig(BaseModel):
    """Shared pipeline settings (settings.yaml)."""

    model_config = ConfigDict(extra="forbid")

    name: str | None = None
    description: str | None = None
    version: str | None = None
    metadata_store: MetadataStoreConfig | None = None
    dlq: DLQConfig | None = None
    contract: str | None = None  # default contract name
    lineage: LineageConfig | None = None
    plugins: PluginConfig | None = None
    state_store: StateStoreConfig | None = None


# =============================================================================
# EXTRACT CONFIGS
# =============================================================================


class PaginationConfig(BaseModel):
    """HTTP pagination configuration."""

    model_config = ConfigDict(extra="forbid")

    enabled: bool = False
    strategy: PaginationStrategy | None = None
    page_param: str | None = None
    offset_param: str | None = None
    limit_param: str | None = None
    cursor_param: str | None = None
    cursor_path: str | None = None
    next_link_path: str | None = None
    page_size: int | None = Field(None, ge=1)
    max_pages: int | None = Field(None, ge=1)


class RetryConfig(BaseModel):
    """HTTP retry configuration."""

    model_config = ConfigDict(extra="forbid")

    max_attempts: int = Field(3, ge=1)
    backoff_factor: float = Field(2.0, ge=0)
    retry_on_status: list[int] = Field(
        default_factory=lambda: [429, 500, 502, 503, 504]
    )


class TimeoutConfig(BaseModel):
    """HTTP timeout configuration."""

    model_config = ConfigDict(extra="ignore")

    connect: float | None = Field(None, ge=0)
    read: float | None = Field(None, ge=0)
    write: float | None = Field(None, ge=0)


class InputParamConfig(BaseModel):
    """Declaration for a single pipeline input parameter."""

    model_config = ConfigDict(extra="allow")

    type: str = "string"
    required: bool = False
    default: Any | None = None
    description: str = ""


class HTTPExtractConfig(BaseModel):
    """HTTP extraction configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["http"] = "http"
    url: str | None = None
    base_url: str | None = None
    api_endpoint: str | None = None
    method: HTTPMethod = HTTPMethod.GET
    headers: dict[str, str] | None = None
    params: dict[str, Any] | None = None
    body: dict[str, Any] | None = None
    response_format: ResponseFormat = ResponseFormat.JSON
    response_path: str | None = None
    pagination: PaginationConfig | None = None
    retry: RetryConfig | None = None
    timeout: TimeoutConfig | None = None
    rate_limit_delay: float | None = Field(None, ge=0)
    max_calls_per_minute: int | None = Field(None, ge=1)
    batch_size: int = Field(1000, ge=1)
    validation: ExtractValidationConfig | None = None
    input_params: dict[str, InputParamConfig] | None = Field(
        None,
        description="Declared input parameters with type, required, default, and description. "
        "Example: {type: {type: string, required: true, description: 'gainers or losers'}}",
    )
    default_param_values: dict[str, str | list[str]] | None = Field(
        None,
        description="Default values for path parameters. Scalar for single-run, list for run_bulk(). "
        "Example: {type: gainers} or {symbol: [AAPL, MSFT, GOOGL]}",
    )
    incremental: IncrementalConfig | None = None


class FileExtractConfig(BaseModel):
    """File extraction configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["file"] = "file"
    path: str
    format: FileFormat | None = None
    csv_options: CSVOptions | None = None
    json_options: JSONOptions | None = None
    batch_size: int = Field(1000, ge=1)
    validation: ExtractValidationConfig | None = None
    incremental: IncrementalConfig | None = None


class DatabaseExtractConfig(BaseModel):
    """Database extraction configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["database"] = "database"
    query: str
    connection_string: str | None = None
    database: DatabaseConfig | None = None
    query_params: dict[str, Any] | None = None
    ssh_tunnel: SSHTunnelConfig | None = None
    batch_size: int = Field(1000, ge=1)
    validation: ExtractValidationConfig | None = None
    incremental: IncrementalConfig | None = None


class CloudStorageExtractConfig(BaseModel):
    """Cloud storage extraction configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["cloud_storage"] = "cloud_storage"
    storage: CloudStorageConfig
    format: FileFormat | None = None
    batch_size: int = Field(1000, ge=1)
    validation: ExtractValidationConfig | None = None
    incremental: IncrementalConfig | None = None


class SSEExtractConfig(BaseModel):
    """SSE (Server-Sent Events) extraction configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["sse"] = "sse"
    url: str
    headers: dict[str, str] | None = None
    event_filter: str | None = None
    data_format: Literal["json", "text"] = "json"
    batch_size: int = Field(1000, ge=1)
    batch_timeout: float = Field(5.0, ge=0)
    reconnect: bool = True
    reconnect_delay: float = Field(3.0, ge=0)
    max_reconnects: int = Field(10, ge=0)
    max_records: int | None = None
    max_batches: int | None = None
    validation: ExtractValidationConfig | None = None


class WebSocketExtractConfig(BaseModel):
    """WebSocket extraction configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["websocket"] = "websocket"
    url: str
    headers: dict[str, str] | None = None
    subscribe_messages: list[dict[str, Any] | str] | None = None
    data_format: Literal["json", "text"] = "json"
    ping_interval: float | None = 20.0
    batch_size: int = Field(1000, ge=1)
    batch_timeout: float = Field(5.0, ge=0)
    reconnect: bool = True
    reconnect_delay: float = Field(3.0, ge=0)
    max_reconnects: int = Field(10, ge=0)
    max_records: int | None = None
    max_batches: int | None = None
    validation: ExtractValidationConfig | None = None


class FileWatcherExtractConfig(BaseModel):
    """File watcher extraction configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["file_watcher"] = "file_watcher"
    watch_dir: str
    patterns: list[str] | None = None
    file_format: str | None = None
    recursive: bool = False
    poll_interval: float = Field(1.0, ge=0.1)
    process_existing: bool = False
    move_after_process: str | None = None
    batch_size: int = Field(1000, ge=1)
    max_records: int | None = None
    validation: ExtractValidationConfig | None = None


class KafkaExtractConfig(BaseModel):
    """Kafka extraction configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["kafka"] = "kafka"
    topics: list[str]
    bootstrap_servers: str = "localhost:9092"
    consumer_group: str = "pycharter"
    data_format: str = "json"
    auto_offset_reset: str = "earliest"
    security_protocol: str = "PLAINTEXT"
    sasl_mechanism: str | None = None
    sasl_username: str | None = None
    sasl_password: str | None = None
    ssl_cafile: str | None = None
    batch_size: int = Field(1000, ge=1)
    batch_timeout: float = Field(5.0, ge=0)
    poll_timeout: float = Field(1.0, ge=0)
    max_records: int | None = None
    max_batches: int | None = None
    validation: ExtractValidationConfig | None = None


class RabbitMQExtractConfig(BaseModel):
    """RabbitMQ extraction configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["rabbitmq"] = "rabbitmq"
    queue: str
    url: str = "amqp://guest:guest@localhost/"
    exchange: str = ""
    routing_key: str = ""
    prefetch_count: int = Field(100, ge=1)
    data_format: str = "json"
    batch_size: int = Field(1000, ge=1)
    batch_timeout: float = Field(5.0, ge=0)
    max_records: int | None = None
    max_batches: int | None = None
    validation: ExtractValidationConfig | None = None


class SQSExtractConfig(BaseModel):
    """SQS extraction configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["sqs"] = "sqs"
    queue_url: str
    region: str = "us-east-1"
    max_messages_per_poll: int = Field(10, ge=1, le=10)
    wait_time_seconds: int = Field(20, ge=0)
    visibility_timeout: int = Field(300, ge=0)
    data_format: str = "json"
    batch_size: int = Field(1000, ge=1)
    batch_timeout: float = Field(5.0, ge=0)
    max_records: int | None = None
    max_batches: int | None = None
    validation: ExtractValidationConfig | None = None


# Union type for extract config - use discriminator on 'type'
ExtractConfig = (
    HTTPExtractConfig
    | FileExtractConfig
    | DatabaseExtractConfig
    | CloudStorageExtractConfig
    | SSEExtractConfig
    | WebSocketExtractConfig
    | FileWatcherExtractConfig
    | KafkaExtractConfig
    | RabbitMQExtractConfig
    | SQSExtractConfig
)

# =============================================================================
# LOAD CONFIGS
# =============================================================================


class PostgresLoadConfig(BaseModel):
    """PostgreSQL load configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["postgres", "postgresql", "database"]
    table: str
    schema_name: str | None = "public"
    write_method: WriteMethod = WriteMethod.INSERT
    primary_key: str | list[str] | None = None
    update_columns: list[str] | None = None
    connection_string: str | None = None
    database: DatabaseConfig | None = None
    ssh_tunnel: SSHTunnelConfig | None = None
    batch_size: int = Field(1000, ge=1)
    validation: LoadValidationConfig | None = None


class SQLiteLoadConfig(BaseModel):
    """SQLite load configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["sqlite"] = "sqlite"
    table: str
    path: str
    write_method: WriteMethod = WriteMethod.INSERT
    primary_key: str | list[str] | None = None
    batch_size: int = Field(1000, ge=1)
    validation: LoadValidationConfig | None = None


class FileLoadConfig(BaseModel):
    """File load configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["file"] = "file"
    path: str
    format: FileFormat = FileFormat.JSON
    write_mode: WriteMode = WriteMode.OVERWRITE
    csv_options: CSVOptions | None = None
    json_options: JSONOptions | None = None
    batch_size: int = Field(1000, ge=1)
    validation: LoadValidationConfig | None = None


class CloudStorageLoadConfig(BaseModel):
    """Cloud storage load configuration."""

    model_config = ConfigDict(extra="forbid")

    type: Literal["cloud_storage"] = "cloud_storage"
    storage: CloudStorageConfig
    format: FileFormat = FileFormat.JSON
    write_mode: WriteMode = WriteMode.OVERWRITE
    partition_by: list[str] | None = None
    batch_size: int = Field(1000, ge=1)
    validation: LoadValidationConfig | None = None


# Union type for load config
LoadConfig = (
    PostgresLoadConfig | SQLiteLoadConfig | FileLoadConfig | CloudStorageLoadConfig
)

# =============================================================================
# TRANSFORM CONFIG
# =============================================================================


class FilterConfig(BaseModel):
    """Filter operation configuration."""

    model_config = ConfigDict(extra="forbid")

    expression: str | None = None
    field: str | None = None
    operator: FilterOperator | None = None
    value: Any | None = None


class MapConfig(BaseModel):
    """Map (lookup) operation configuration."""

    model_config = ConfigDict(extra="forbid")

    field: str
    mapping: dict[str, Any] | None
    default: Any | None = None


class JSONataConfig(BaseModel):
    """JSONata transformation configuration."""

    model_config = ConfigDict(extra="forbid")

    expression: str
    mode: JSONataMode = JSONataMode.RECORD


class CustomFunctionConfig(BaseModel):
    """Custom Python function configuration."""

    model_config = ConfigDict(extra="forbid")

    module: str
    function: str
    kwargs: dict[str, Any] | None = None
    mode: JSONataMode = JSONataMode.BATCH


class SimpleOpsConfig(BaseModel):
    """Simple transform operations (object format)."""

    model_config = ConfigDict(extra="forbid")

    rename: dict[str, str] | None = None
    convert: dict[str, ConvertType] | None = None
    defaults: dict[str, Any] | None = None
    add: dict[str, Any] | None = None
    select: list[str] | None = None
    drop: list[str] | None = None
    filter: FilterConfig | None = None


class TransformConfig(BaseModel):
    """Transform configuration (simple ops, JSONata, custom function)."""

    model_config = ConfigDict(extra="forbid")

    # Simple ops: can be list (ordered) or dict (legacy)
    transform: SimpleOpsConfig | list[dict[str, Any]] | None = None
    # JSONata transformation
    jsonata: JSONataConfig | None = None
    # Custom Python function
    custom_function: CustomFunctionConfig | None = None


# =============================================================================
# PIPELINE CONFIG (single-file format)
# =============================================================================


class PipelineConfig(BaseModel):
    """Complete pipeline configuration (single-file format)."""

    model_config = ConfigDict(extra="forbid")

    name: str | None = None
    version: str | None = None
    description: str | None = None

    # Core ETL steps (validated separately based on type)
    extract: dict[str, Any] | None
    transform: dict[str, Any] | list[dict[str, Any]] | None = None
    jsonata: JSONataConfig | None = None
    custom_function: CustomFunctionConfig | None = None
    load: dict[str, Any] | None
    # Inline variables for ${VAR} substitution
    variables: dict[str, str] | None = None
    # Inline settings (for single-file format)
    metadata_store: MetadataStoreConfig | None = None
    dlq: DLQConfig | None = None
    contract: str | None = None
    lineage: LineageConfig | None = None


# =============================================================================
# HELPER FUNCTIONS
# =============================================================================


def parse_extract_config(config: dict[str, Any]) -> ExtractConfig:
    """Parse extract config dict into typed model based on 'type' field."""
    extract_type = config.get("type", "").lower()

    if extract_type == "http":
        return HTTPExtractConfig(**config)
    elif extract_type == "file":
        return FileExtractConfig(**config)
    elif extract_type == "database":
        return DatabaseExtractConfig(**config)
    elif extract_type == "cloud_storage":
        return CloudStorageExtractConfig(**config)
    elif extract_type == "sse":
        return SSEExtractConfig(**config)
    elif extract_type == "websocket":
        return WebSocketExtractConfig(**config)
    elif extract_type == "file_watcher":
        return FileWatcherExtractConfig(**config)
    elif extract_type == "kafka":
        return KafkaExtractConfig(**config)
    elif extract_type == "rabbitmq":
        return RabbitMQExtractConfig(**config)
    elif extract_type == "sqs":
        return SQSExtractConfig(**config)
    else:
        raise ValueError(
            f"Unknown extract type: '{extract_type}'. "
            f"Supported types: http, file, database, cloud_storage, "
            f"sse, websocket, file_watcher, kafka, rabbitmq, sqs"
        )


def parse_load_config(config: dict[str, Any]) -> LoadConfig:
    """Parse load config dict into typed model based on 'type' field."""
    load_type = config.get("type", "").lower()

    if load_type in ("postgres", "postgresql", "database"):
        return PostgresLoadConfig(**config)
    elif load_type == "sqlite":
        return SQLiteLoadConfig(**config)
    elif load_type == "file":
        return FileLoadConfig(**config)
    elif load_type == "cloud_storage":
        return CloudStorageLoadConfig(**config)
    else:
        raise ValueError(
            f"Unknown load type: '{load_type}'. "
            f"Supported types: postgres, postgresql, database, sqlite, file, cloud_storage"
        )


def parse_transform_config(
    config: dict[str, Any] | list[dict[str, Any]] | None,
) -> TransformConfig | None:
    """Parse transform config into typed model."""
    if config is None:
        return None

    # Handle top-level transform that's just a list or simple ops dict
    if isinstance(config, list):
        return TransformConfig(transform=config)

    return TransformConfig(**config)


def parse_settings_config(config: dict[str, Any]) -> SettingsConfig:
    """Parse settings config into typed model."""
    return SettingsConfig(**config)
