# coding: utf-8
# Copyright (c) 2025 inclusionAI.
import os
import traceback
import uuid
from collections import OrderedDict
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Callable, Union, Iterable, Type, TYPE_CHECKING, ClassVar

import yaml
from pydantic import BaseModel, Field

from aworld.logs.util import logger

if TYPE_CHECKING:
    from aworld.dataset.trajectory_strategy import TrajectoryStrategy


def load_config(file_name: str, dir_name: str = None) -> Dict[str, Any]:
    """Dynamically load config file form current path.

    Args:
        file_name: Config file name.
        dir_name: Config file directory.

    Returns:
        Config dict.
    """

    if dir_name:
        file_path = os.path.join(dir_name, file_name)
    else:
        # load conf form current path
        current_dir = Path(__file__).parent.absolute()
        file_path = os.path.join(current_dir, file_name)
    if not os.path.exists(file_path):
        logger.debug(f"{file_path} not exists, please check it.")

    configs = dict()
    try:
        with open(file_path, "r") as file:
            yaml_data = yaml.safe_load(file)
        configs.update(yaml_data)
    except FileNotFoundError:
        logger.debug(f"Can not find the file: {file_path}")
    except Exception:
        logger.warning(f"{file_name} read fail.\n", traceback.format_exc())
    return configs


def wipe_secret_info(config: Dict[str, Any], keys: List[str]) -> Dict[str, Any]:
    """Return a deep copy of this config as a plain Dict as well ass wipe up secret info, used to log."""

    def _wipe_secret(conf):
        def _wipe_secret_plain_value(v):
            if isinstance(v, List):
                return [_wipe_secret_plain_value(e) for e in v]
            elif isinstance(v, Dict):
                return _wipe_secret(v)
            else:
                return v

        key_list = []
        for key in conf.keys():
            key_list.append(key)
        for key in key_list:
            if key.strip('"') in keys:
                conf[key] = '-^_^-'
            else:
                _wipe_secret_plain_value(conf[key])
        return conf

    if not config:
        return config
    return _wipe_secret(config)


class ClientType(Enum):
    SDK = "sdk"
    HTTP = "http"


class HistoryWriteStrategy(Enum):
    """History write strategy for memory operations."""
    EVENT_DRIVEN = "event_driven"  # Write through message system (default)
    DIRECT = "direct"  # Direct call to memory handler


class ConfigDict(dict):
    """Object mode operates dict, can read non-existent attributes through `get` method."""
    __setattr__ = dict.__setitem__
    __getattr__ = dict.__getitem__

    def __init__(self, seq: dict = None, **kwargs):
        if seq is None:
            seq = OrderedDict()
        super(ConfigDict, self).__init__(seq, **kwargs)
        self.nested(self)

    def nested(self, seq: dict):
        """Nested recursive processing dict.

        Args:
            seq: Python original format dict
        """
        for k, v in seq.items():
            if isinstance(v, dict):
                seq[k] = ConfigDict(v)
                self.nested(v)


class BaseConfig(BaseModel):
    def to_dict(self) -> ConfigDict:
        return ConfigDict(self.model_dump())


class ModelConfig(BaseConfig):
    model_config = ConfigDict(extra='allow')
    llm_provider: str = "openai"
    llm_model_name: str = None
    llm_temperature: float = 1.
    llm_base_url: str = None
    llm_api_key: str = None
    llm_client_type: ClientType = ClientType.SDK
    llm_sync_enabled: bool = True
    llm_async_enabled: bool = True
    llm_stream_call: bool = False
    max_retries: int = 3
    max_model_len: Optional[int] = None  # Maximum model context length
    model_type: Optional[str] = 'qwen'  # Model type determines tokenizer and maximum length
    params: Optional[Dict[str, Any]] = {}
    ext_config: Optional[Dict[str, Any]] = {}

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        for key, value in kwargs.items():
            if hasattr(self, key):
                setattr(self, key, value)

        # init max_model_len
        if self.max_model_len is None:
            # qwen or other default model_type
            self.max_model_len = 128000 if self.model_type != 'claude' else 200000


class LlmCompressionConfig(BaseConfig):
    enabled: bool = False
    compress_type: str = 'llm'  # llm, llmlingua
    trigger_compress_token_length: int = 10000  # Trigger compression when exceeding this length
    compress_model: Optional[ModelConfig] = Field(default=None, description="Compression model configuration")


class OptimizationConfig(BaseConfig):
    enabled: bool = False
    max_token_budget_ratio: float = 0.5  # Maximum context length ratio


class SummaryPromptConfig(BaseConfig):
    """Configuration for summary prompt templates."""
    
    template: str = Field(description="Base template, such as AWORLD_MEMORY_EXTRACT_NEW_SUMMARY")
    summary_rule: str = Field(description="Summary rule, used to guide how to generate summaries")
    summary_schema: str = Field(description="Summary schema, defines output format and structure")
    memory_type: str = Field(default="summary", description="Memory type, used to distinguish different types of summaries")


class ContextRuleConfig(BaseConfig):
    """Context interference rule configuration"""

    # ===== Performance optimization configuration =====
    optimization_config: OptimizationConfig = OptimizationConfig()

    # ===== LLM conversation compression configuration =====
    llm_compression_config: LlmCompressionConfig = LlmCompressionConfig()


class AgentMemoryConfig(BaseConfig):
    """Configuration for procedural memory."""

    model_config = ConfigDict(
        from_attributes=True, validate_default=True, revalidate_instances='always', validate_assignment=True,
        arbitrary_types_allowed=True
    )
    # short-term config
    history_rounds: int = Field(default=100,
                                description="rounds of message msg; when the number of messages is greater than the history_rounds, the memory will be trimmed")
    history_write_strategy: HistoryWriteStrategy = Field(default=HistoryWriteStrategy.EVENT_DRIVEN,
                                                         description="History write strategy: event_driven (through message system) or direct (direct call to handler)")
    
    enable_summary: bool = Field(default=False,
                                 description="enable_summary use llm to create summary short-term memory")
    summary_model: Optional[str] = Field(default=None, description="short-term summary model")
    summary_rounds: Optional[int] = Field(default=5,
                                          description="rounds of message msg; when the number of messages is greater than the summary_rounds, the summary will be created")
    summary_context_length: Optional[int] = Field(default=40960,
                                                  description=" when the content length is greater than the summary_context_length, the summary will be created")
    summary_prompts: Optional[List[SummaryPromptConfig]] = Field(default=[])
    summary_summaried: Optional[bool] = Field(default=True, description="to summary summary message when summary triggered")
    summary_role: Optional[str] = Field(default="assistant", description="role for summary memory items")

    # Long-term memory config
    enable_long_term: bool = Field(default=False, description="enable_long_term use to store long-term memory")
    long_term_model: Optional[str] = Field(default=None, description="long-term extract model")
    # LongTermConfig
    long_term_config: Optional[BaseModel] = Field(default=None, description="long_term_config")


class AgentConfig(BaseConfig):
    llm_config: ModelConfig = ModelConfig()
    memory_config: AgentMemoryConfig = AgentMemoryConfig()
    context_rule: ContextRuleConfig = ContextRuleConfig()

    # default reset init in first
    need_reset: bool = True
    # use vision model
    use_vision: bool = True
    max_steps: int = 10
    max_input_tokens: int = 128000
    max_actions_per_step: int = 10
    system_prompt: Optional[str] = None
    system_prompt_template: Optional[str] = None
    working_dir: Optional[str] = None
    enable_recording: bool = False
    use_tools_in_prompt: bool = False
    exit_on_failure: bool = False
    human_tools: List[str] = []
    skill_configs: Dict[str, Any] = None
    ext: dict = {}

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # Initialize llm_config with relevant kwargs
        llm_config_kwargs = {}
        llm_config_ext = {}
        for k, v in kwargs.items():
            if k in ModelConfig.model_fields:
                llm_config_kwargs[k] = v
            elif k not in self.model_fields:
                llm_config_ext[k] = v

        # Reassignment if it has llm config args
        if llm_config_kwargs or not self.llm_config:
            self.llm_config = ModelConfig(**llm_config_kwargs)

        self.llm_config.ext_config.update(llm_config_ext)

    @property
    def llm_model_name(self) -> str:
        return self.llm_config.llm_model_name

    @property
    def llm_provider(self) -> str:
        return self.llm_config.llm_provider


class TaskRunMode(Enum):
    INTERACTIVE = "INTERACTIVE"
    ONE_WAY = "ONE_WAY"


class TaskConfig(BaseConfig):
    model_config = {"arbitrary_types_allowed": True}
    task_id: str = str(uuid.uuid4())
    task_name: str | None = None
    max_steps: int = 100
    trajectory_strategy: Optional[Type['TrajectoryStrategy']] = None
    stream: bool = False
    resp_carry_context: bool = True
    resp_carry_raw_llm_resp: bool = False
    exit_on_failure: bool = False
    ext: dict = {}
    run_mode: TaskRunMode = TaskRunMode.ONE_WAY


class ToolConfig(BaseConfig):
    name: str = None
    custom_executor: bool = False
    enable_recording: bool = False
    working_dir: str = ""
    max_retry: int = 3
    llm_config: ModelConfig = None
    reuse: bool = False
    use_async: bool = False
    exit_on_failure: bool = False
    ext: dict = {}


class EngineName:
    # Use asyncio or MultiProcess run in local
    LOCAL = "local"
    # Stateless(task) run in ray. Ray actor will use a new name
    RAY = "ray"
    SPARK = "spark"


class RunConfig(BaseConfig):
    job_name: str = "aworld_job"
    engine_name: str = EngineName.LOCAL
    worker_num: int = 1
    # engine whether to run in local
    in_local: bool = True
    # run in local whether to use the same process
    reuse_process: bool = True
    # Is the task sequence dependent
    sequence_dependent: bool = False
    # The custom implement of RuntimeEngine
    cls: Optional[str] = None
    event_bus: Optional[Dict[str, Any]] = None
    tracer: Optional[Dict[str, Any]] = None


class StorageConfig(BaseConfig):
    name: str = "inmemory"


class DataLoaderConfig(BaseConfig):
    batch_size: Optional[int] = 1
    sampler: Any = None
    shuffle: bool = False
    drop_last: bool = False
    seed: Optional[int] = None
    batch_sampler: Optional[Iterable[List[int]]] = None
    collate_fn: Optional[Callable[..., Any]] = None


class DatasetConfig(BaseConfig):
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    name: str
    metadata: Dict[str, Any] = Field(default_factory=dict)
    transforms: List[Callable[..., Any]] = Field(default_factory=list)

    # Config for loading dataset from source
    format: Optional[str] = None
    split: Optional[str] = None
    subset: Optional[str] = None
    json_field: Optional[str] = None
    parquet_columns: Optional[List[str]] = None
    encoding: str = "utf-8"
    limit: Optional[int] = None
    preload_transform: Optional[Callable[..., Any]] = None

    # Config for dataloader
    dataloader_config: DataLoaderConfig = DataLoaderConfig()


class EvaluationConfig(BaseConfig):
    '''
    Evaluation run config.
    '''
    # full class name of eval target, e.g. aworld.evaluations.base.EvalTarget
    eval_target: Any = None
    eval_target_full_class_name: str = None
    eval_target_config: dict = None
    eval_criterias: List[Union[dict]] = None
    # eval dataset id or file path, file path should be a jsonl file
    eval_dataset_id_or_file_path: str = None
    eval_dataset_load_config: Optional[DataLoaderConfig] = DataLoaderConfig()
    # preload transform function or function name, e.g. aworld.evaluations.base.preload_transform
    eval_dataset_preload_transform: Optional[Union[Callable[[any], Any], str]] = None
    eval_dataset_query_column: Optional[str] = "query"
    eval_dataset_answer_column: Optional[str] = "answer"
    eval_output_answer_column: Optional[str] = "answer"
    repeat_times: int = 1
    parallel_num: int = 1
    skip_passed_cases: bool = False
    skip_passed_on_metrics: List[str] = []
