from pydantic import BaseModel, create_model
from typing import Any, Dict, List, Optional, Union, get_args, get_origin
import json
from datetime import datetime, date

def generate_json_example(model: type[BaseModel], include_optional: bool = False,list_length = 2) -> str:
    """
    根据 Pydantic V2 模型生成示例 JSON 数据结构。
    """
    def _generate_example(field_type: Any) -> Any:
        origin = get_origin(field_type)
        args = get_args(field_type)

        if origin is list or origin is List:
            # 生成多个元素，这里生成 3 个
            result = [_generate_example(args[0]) for _ in range(list_length)] if args else []
            result.append("......")
            return result
        elif origin is dict or origin is Dict:
            if len(args) == 2:
                return {"key": _generate_example(args[1])}
            return {}
        elif origin is Union:
            # 处理 Optional 类型（Union[T, None]）
            non_none_args = [arg for arg in args if arg is not type(None)]
            return _generate_example(non_none_args[0]) if non_none_args else None
        elif field_type is str:
            return "string"
        elif field_type is int:
            return 0
        elif field_type is float:
            return 0.0
        elif field_type is bool:
            return True
        elif field_type is datetime:
            return datetime.now().isoformat()
        elif field_type is date:
            return date.today().isoformat()
        elif isinstance(field_type, type) and issubclass(field_type, BaseModel):
            return json.loads(generate_json_example(field_type, include_optional))
        else:
            # 处理直接类型注解（非泛型）
            if field_type is type(None):
                return None
            try:
                if issubclass(field_type, BaseModel):
                    return json.loads(generate_json_example(field_type, include_optional))
            except TypeError:
                pass
            return "unknown"

    example_data = {}
    for field_name, field in model.model_fields.items():
        if include_optional or not isinstance(field.default, type(None)):
            example_data[field_name] = _generate_example(field.annotation)

    return json.dumps(example_data, indent=2, default=str)

if __name__ == "__main__":
    import sys
    from pathlib import Path
    # 添加项目根目录到sys.path
    sys.path.append(str(Path(__file__).resolve().parent.parent))
    from schema import Dataset

    print("示例 JSON:")
    print(generate_json_example(Dataset))
