from __future__ import annotations

from typing import Any, Dict, List, Optional

from transformers import PretrainedConfig


DEFAULT_EVENT_CATEGORIES = [
    "earnings",
    "product",
    "macro",
    "regulation",
    "supply_chain",
    "competition",
    "other",
]

DEFAULT_COVARIATE_COLUMNS = [
    "cov_earnings_count",
    "cov_product_count",
    "cov_macro_count",
    "cov_regulation_count",
    "cov_supply_chain_count",
    "cov_competition_count",
    "cov_other_count",
    "cov_sentiment_pos_count",
    "cov_sentiment_neg_count",
    "cov_sentiment_neu_count",
    "cov_news_count",
    "cov_sentiment_mean",
    "cov_confidence_mean",
    "cov_event_score",
]


class JNUTSBConfig(PretrainedConfig):
    """Configuration for the JNU-TSB router wrapper.

    The repository stores lightweight code and metadata only. The upstream
    models, amazon/chronos-2 and EleutherAI/polyglot-ko-1.3b, are loaded lazily
    at runtime when the corresponding route is used.
    """

    model_type = "jnu_tsb"

    def __init__(
        self,
        repo_id: str = "HONGRIZON/JNU-TSB",
        project_name: str = "JNU-TSB",
        project_full_name: str = "Jeju National University Time-Series Bridge",
        chronos_model_id: str = "amazon/chronos-2",
        llm_model_id: str = "EleutherAI/polyglot-ko-1.3b",
        timestamp_column: str = "timestamp",
        target_column: str = "target",
        id_column: str = "item_id",
        default_item_id: str = "series_0",
        prediction_length: int = 5,
        quantile_levels: Optional[List[float]] = None,
        event_categories: Optional[List[str]] = None,
        covariate_columns: Optional[List[str]] = None,
        use_llm_extractor: bool = True,
        allow_naive_fallback: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.repo_id = repo_id
        self.project_name = project_name
        self.project_full_name = project_full_name
        self.chronos_model_id = chronos_model_id
        self.llm_model_id = llm_model_id
        self.timestamp_column = timestamp_column
        self.target_column = target_column
        self.id_column = id_column
        self.default_item_id = default_item_id
        self.prediction_length = int(prediction_length)
        self.quantile_levels = quantile_levels or [0.1, 0.5, 0.9]
        self.event_categories = event_categories or list(DEFAULT_EVENT_CATEGORIES)
        self.covariate_columns = covariate_columns or list(DEFAULT_COVARIATE_COLUMNS)
        self.use_llm_extractor = bool(use_llm_extractor)
        self.allow_naive_fallback = bool(allow_naive_fallback)

    def to_router_dict(self) -> Dict[str, Any]:
        return {
            "repo_id": self.repo_id,
            "project_name": self.project_name,
            "project_full_name": self.project_full_name,
            "chronos_model_id": self.chronos_model_id,
            "llm_model_id": self.llm_model_id,
            "timestamp_column": self.timestamp_column,
            "target_column": self.target_column,
            "id_column": self.id_column,
            "default_item_id": self.default_item_id,
            "prediction_length": self.prediction_length,
            "quantile_levels": self.quantile_levels,
            "event_categories": self.event_categories,
            "covariate_columns": self.covariate_columns,
            "use_llm_extractor": self.use_llm_extractor,
            "allow_naive_fallback": self.allow_naive_fallback,
        }