#!/bin/bash
# 从Docker容器导出热点股数据

set -e

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

CONTAINER_NAME="auv"
DB_PATH="/app/backend/data/stock_data.db"

# 检查容器是否运行
if ! docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
    echo -e "${RED}❌ 容器 ${CONTAINER_NAME} 未运行${NC}"
    exit 1
fi

# 检查数据库是否存在
if ! docker exec "$CONTAINER_NAME" test -f "$DB_PATH"; then
    echo -e "${RED}❌ 数据库文件不存在：$DB_PATH${NC}"
    exit 1
fi

# 显示帮助
show_help() {
    echo "用法："
    echo "  $0 export [日期]    - 导出指定日期的数据（默认：2026-08-20）"
    echo "  $0 import <文件>    - 导入数据文件到容器"
    echo "  $0 list             - 列出所有有数据的日期"
    echo "  $0 backup           - 备份整个数据库"
    echo "  $0 shell            - 进入容器的Python shell"
    echo ""
    echo "示例："
    echo "  $0 export 2026-08-20"
    echo "  $0 import hotspot_2026-08-20.json"
    echo "  $0 backup"
}

# 导出数据
export_data() {
    local date="${1:-2026-08-20}"
    local output_file="hotspot_${date}.json"
    
    echo -e "${YELLOW}📤 从容器导出 $date 的热点股数据...${NC}"
    
    # 在容器内执行导出命令
    docker exec "$CONTAINER_NAME" python3 -c "
import sys
sys.path.insert(0, '/app')
from database import get_connection, dict_from_row
import json

date = '$date'
conn = get_connection()
try:
    themes = conn.execute('SELECT * FROM daily_top_themes WHERE trade_date = ? ORDER BY rank ASC', (date,)).fetchall()
    core_stocks = conn.execute('SELECT * FROM daily_core_stocks WHERE trade_date = ? ORDER BY rank ASC', (date,)).fetchall()
    stock_themes = conn.execute('SELECT * FROM daily_core_stock_themes WHERE trade_date = ?', (date,)).fetchall()
    
    data = {
        'version': '1.0',
        'trade_date': date,
        'themes': [dict_from_row(r) for r in themes],
        'core_stocks': [dict_from_row(r) for r in core_stocks],
        'stock_themes': [dict_from_row(r) for r in stock_themes],
    }
    data['stats'] = {
        'theme_count': len(data['themes']),
        'core_stock_count': len(data['core_stocks']),
        'stock_theme_count': len(data['stock_themes']),
    }
    
    print(json.dumps(data, ensure_ascii=False, indent=2))
finally:
    conn.close()
" > "$output_file"
    
    if [ -f "$output_file" ]; then
        echo -e "${GREEN}✅ 导出完成：$output_file${NC}"
        echo "文件大小：$(du -h "$output_file" | cut -f1)"
    else
        echo -e "${RED}❌ 导出失败${NC}"
        exit 1
    fi
}

# 导入数据
import_data() {
    local file="$1"
    
    if [ -z "$file" ]; then
        echo -e "${RED}❌ 请指定要导入的文件${NC}"
        exit 1
    fi
    
    if [ ! -f "$file" ]; then
        echo -e "${RED}❌ 文件不存在：$file${NC}"
        exit 1
    fi
    
    echo -e "${YELLOW}📥 导入数据到容器：$file${NC}"
    
    # 复制文件到容器
    docker cp "$file" "$CONTAINER_NAME:/tmp/import_data.json"
    
    # 在容器内执行导入
    docker exec "$CONTAINER_NAME" python3 -c "
import sys
import json
sys.path.insert(0, '/app')
from database import get_connection

with open('/tmp/import_data.json', 'r', encoding='utf-8') as f:
    data = json.load(f)

date = data['trade_date']
print(f'📅 导入日期：{date}')
print(f'   题材数量：{len(data[\"themes\"])}')
print(f'   核心股票：{len(data[\"core_stocks\"])}')
print(f'   股票-题材关联：{len(data[\"stock_themes\"])}')

conn = get_connection()
try:
    for item in data['themes']:
        item.pop('id', None)
        item.pop('created_at', None)
        conn.execute(
            'INSERT OR IGNORE INTO daily_top_themes (trade_date, theme_code, theme_name, bf3, hot_rank, rank) VALUES (?,?,?,?,?,?)',
            (item['trade_date'], item['theme_code'], item['theme_name'], item['bf3'], item['hot_rank'], item['rank']),
        )
    
    for item in data['core_stocks']:
        item.pop('id', None)
        item.pop('created_at', None)
        conn.execute(
            'INSERT OR IGNORE INTO daily_core_stocks (trade_date, stock_code, stock_name, f3, cover_count, rank) VALUES (?,?,?,?,?,?)',
            (item['trade_date'], item['stock_code'], item['stock_name'], item['f3'], item['cover_count'], item['rank']),
        )
    
    for item in data['stock_themes']:
        item.pop('id', None)
        item.pop('created_at', None)
        conn.execute(
            'INSERT OR IGNORE INTO daily_core_stock_themes (trade_date, stock_code, theme_code, theme_name) VALUES (?,?,?,?)',
            (item['trade_date'], item['stock_code'], item['theme_code'], item['theme_name']),
        )
    
    conn.commit()
    print(f'\\n✅ 导入成功：{date}')
except Exception as e:
    conn.rollback()
    print(f'\\n❌ 导入失败：{e}')
finally:
    conn.close()
"
}

# 列出日期
list_dates() {
    echo -e "${YELLOW}📅 列出容器中的日期...${NC}"
    docker exec "$CONTAINER_NAME" python3 -c "
import sys
sys.path.insert(0, '/app')
from database import get_connection

conn = get_connection()
try:
    rows = conn.execute('SELECT DISTINCT trade_date FROM daily_core_stocks ORDER BY trade_date DESC').fetchall()
    if rows:
        print('📅 数据库中的日期：')
        for row in rows:
            print(f'   {row[\"trade_date\"]}')
    else:
        print('📭 数据库中暂无数据')
finally:
    conn.close()
"
}

# 备份数据库
backup_database() {
    local timestamp=$(date +%Y%m%d_%H%M%S)
    local backup_file="stock_data_${timestamp}.db"
    
    echo -e "${YELLOW}💾 备份数据库...${NC}"
    docker cp "$CONTAINER_NAME:$DB_PATH" "$backup_file"
    
    if [ -f "$backup_file" ]; then
        echo -e "${GREEN}✅ 备份完成：$backup_file${NC}"
        echo "文件大小：$(du -h "$backup_file" | cut -f1)"
    else
        echo -e "${RED}❌ 备份失败${NC}"
        exit 1
    fi
}

# 进入Python shell
enter_shell() {
    echo -e "${YELLOW}🐍 进入容器Python shell...${NC}"
    docker exec -it "$CONTAINER_NAME" python3 -c "
import sys
sys.path.insert(0, '/app')
from database import get_connection
print('Python shell 已启动')
print('可用变量：')
print('  conn - 数据库连接')
print('  get_connection - 获取新连接函数')
print()
conn = get_connection()
"
}

# 主逻辑
case "${1:-help}" in
    export)
        export_data "$2"
        ;;
    import)
        import_data "$2"
        ;;
    list)
        list_dates
        ;;
    backup)
        backup_database
        ;;
    shell)
        enter_shell
        ;;
    *)
        show_help
        ;;
esac