import os
import csv
from decimal import Decimal
from datetime import datetime


def merge_portfolio_files():
    data_dir = "data"
    if not os.path.exists(data_dir):
        print(f"目录 {data_dir} 不存在")
        return

    portfolio_files = []
    for filename in os.listdir(data_dir):
        if filename.startswith("portfolio_") and filename.endswith(".csv"):
            portfolio_files.append(filename)

    if not portfolio_files:
        print("未找到portfolio_*.csv文件")
        return

    for input_filename in portfolio_files:
        output_filename = input_filename.replace("portfolio_", "merged_")
        input_path = os.path.join(data_dir, input_filename)
        output_path = os.path.join(data_dir, output_filename)

        try:
            process_single_file(input_path, output_path)
            print(f"成功处理: {input_filename} -> {output_filename}")
        except Exception as e:
            print(f"处理文件 {input_filename} 时出错: {e}")


def process_single_file(input_path, output_path):
    date_data = {}  # 格式: {日期: {'amount': 净额总和, 'last_time': 最后时间}}

    with open(input_path, "r", encoding="utf-8") as infile:
        reader = csv.DictReader(infile)

        if not all(field in reader.fieldnames for field in ["日期", "类型", "净额"]):
            raise ValueError("CSV文件格式不正确，需要的列：日期,类型,净额")

        for row in reader:
            full_datetime_str = row["日期"]
            try:
                full_datetime = datetime.fromisoformat(
                    full_datetime_str.replace("Z", "+00:00")
                )
            except ValueError:
                print(f"警告：跳过无效的日期格式: {full_datetime_str}")
                continue

            date_only = full_datetime.date().isoformat()

            try:
                amount = Decimal(row["净额"])
            except ValueError:
                print(f"警告：跳过无效的净额值: {row['净额']}")
                continue

            if date_only in date_data:
                date_data[date_only]["amount"] += amount
                if full_datetime > date_data[date_only]["last_time"]:
                    date_data[date_only]["last_time"] = full_datetime
            else:
                date_data[date_only] = {"amount": amount, "last_time": full_datetime}

    sorted_dates = sorted(date_data.keys())

    with open(output_path, "w", encoding="utf-8", newline="") as outfile:
        writer = csv.writer(outfile)
        writer.writerow(["日期", "类型", "净额"])

        for date in sorted_dates:
            data = date_data[date]
            last_datetime_str = data["last_time"].strftime("%Y-%m-%dT%H:%M")
            writer.writerow([last_datetime_str, "利息", data['amount']])


def main():
    """主函数"""
    print("开始处理portfolio文件...")
    merge_portfolio_files()
    print("处理完成！")


if __name__ == "__main__":
    main()