import os
import yaml
import tushare as ts
import pandas as pd
from datetime import datetime, timedelta
from sqlalchemy import create_engine

# 模块级单例
_config = None
_engine = None


def load_config():
    global _config
    if _config is not None:
        return _config

    config_path = 'config.yaml'
    if not os.path.exists(config_path):
        config = {
            'tushare_token': 'your_token_here',
            'mysql': {
                'user': 'user',
                'password': 'password',
                'host': '127.0.0.1',
                'port': 3306,
                'database': 'tushare',
                'charset': 'utf8'
            },
            'openai_api': {
                'api_key': 'sk-your-api-key-here',
                'base_url': 'https://api.tu-zi.com/v1',
                'model': 'gpt-4'
            }
        }
        with open(config_path, 'w') as f:
            yaml.dump(config, f)
        print(f"请在 {config_path} 中填入您的 tushare token、MySQL 连接信息以及第三方 API 信息")
        exit(1)

    with open(config_path, 'r') as f:
        _config = yaml.safe_load(f)

    return _config


def get_engine():
    """获取单例数据库引擎"""
    global _engine
    if _engine is None:
        config = load_config()
        mysql = config['mysql']
        connection_string = f"mysql+pymysql://{mysql['user']}:{mysql['password']}@{mysql['host']}:{mysql['port']}/{mysql['database']}?charset={mysql['charset']}&use_unicode=1"
        _engine = create_engine(connection_string)

    return _engine

def get_trade_cal(start_date=None, end_date=None):
    """
    获取指定时间段内的交易日历

    参数：
    start_date (str): 开始日期，格式'YYYYMMDD'
    end_date (str): 结束日期，格式'YYYYMMDD'

    返回：
    pandas.DataFrame: 交易日历
    """
    if start_date is None:
        start_date = (datetime.now() - timedelta(days=30)).strftime('%Y%m%d')
    if end_date is None:
        end_date = datetime.now().strftime('%Y%m%d')

    pro = ts.pro_api()
    trade_cal_df = pro.trade_cal(exchange='', start_date=start_date, end_date=end_date)
    return trade_cal_df[trade_cal_df['is_open'] == 1]['cal_date'].tolist()