from __future__ import annotations

from typing import Annotated, Any, Dict, List, Literal

from pydantic import BaseModel, ConfigDict, Field, PlainSerializer, PlainValidator

from ..datasource.type_registry import register_datasource_config_model, validate_registered_datasource_config


class BaseDataSourceConfig(BaseModel):
    model_config = ConfigDict(extra="forbid", validate_assignment=True)

    type: str
    handler_configs: Dict[str, Dict[str, Any]] = Field(default_factory=dict)


class JsonlDataSourceConfig(BaseDataSourceConfig):

    type: Literal["jsonl"] = "jsonl"
    jsonl_dir: str
    read_only: bool = False


class ApiDataSourceConfig(BaseDataSourceConfig):

    type: Literal["api"] = "api"
    api_base_url: str = ""
    api_uid: str = ""
    api_secret: str = ""
    api_debug: bool = False
    poll_max_retries: int = 10
    poll_interval: float = 0.5
    api_rate_limit_max_requests: int = 10
    api_rate_limit_window_seconds: float = 10.0


DataSourceConfig = Annotated[
    BaseDataSourceConfig,
    PlainValidator(validate_registered_datasource_config),
    PlainSerializer(
        lambda value: value.model_dump(mode="python") if isinstance(value, BaseModel) else value,
        return_type=dict[str, Any],
    ),
]


register_datasource_config_model("jsonl", JsonlDataSourceConfig)
register_datasource_config_model("api", ApiDataSourceConfig)