from typing import Any, Dict, Type

from rm_gallery.core.model.base import BaseLLM
from rm_gallery.core.model.openai_llm import OpenaiLLM
from rm_gallery.core.reward.base import BaseReward
from rm_gallery.core.reward.composition import RouterComposition
from rm_gallery.core.reward.registry import RewardRegistry
from rm_gallery.gallery.rm.finance.event_interpretation.composition import EventInterpretationComposition
from rm_gallery.gallery.rm.finance.industry_research.composition import IndustryResearchComposition
from rm_gallery.gallery.rm.finance.macro_analysis.composition import MacroAnalysisComposition
from rm_gallery.gallery.rm.finance.stock_analysis.composition import StockAnalysisComposition
from rm_gallery.gallery.rm.finance.stock_search.composition import StockSearchComposition

@RewardRegistry.register("finance_composition")
class FinanceComposition(RouterComposition):

    rewards: Dict[str, Type[BaseReward]] = {
        "industry_research": IndustryResearchComposition,
        "event_interpretation": EventInterpretationComposition,
        "stock_analysis": StockAnalysisComposition,
        "macro_analysis": MacroAnalysisComposition,
        "stock_search": StockSearchComposition,
    }
    params: Dict[str, Any] = {
        "is_parallel": True,
    }
    
    
    def __init__(self, name: str | None, llm: BaseLLM | str | None, params: Dict[str, Any] = None,*args,**kwargs):
        # 显式将 llm 添加到 params 字典中，确保子组件能正确接收
        if params is None:
            params = {}
        if isinstance(llm, str):
            llm = OpenaiLLM(model=llm,**params)
            params["llm"] = llm
        elif isinstance(llm, BaseLLM) or issubclass(llm, BaseLLM):
            params["llm"] = llm
        if not params.get("llm"):
            raise ValueError("llm is required")
        super().__init__(name=name, params=params)

    def _condition(self, sample) -> str:
        return sample.metadata.get("domain", "default")