import pandas as pd
import os
from typing import List

# --- 配置参数 ---
INPUT_CSV_PATH = "/home/hsichen/part_time/BERT_finetune/dataset_pretrain/Experiment_sentences_training_filtered_part1.csv"
# 输出的纯文本文件路径 (用于 DAPT 脚本中的 DOMAIN_TEXT_FILE)
OUTPUT_TXT_PATH = "/home/hsichen/part_time/BERT_finetune/dataset_pretrain/domain_corpus.txt"
# 纯文本文件的编码
ENCODING = 'utf-8'

def prepare_dapt_data(input_csv_path: str, output_txt_path: str, encoding: str):
    """
    从 CSV 文件中提取 'sentence' 列，并保存为纯文本文件，每行一个句子。
    
    Args:
        input_csv_path: 原始 CSV 文件的路径。
        output_txt_path: 目标纯文本文件的路径。
        encoding: 文件编码。
    """
    print(f"--- 1. 读取数据: {input_csv_path} ---")
    
    try:
        # 尝试读取 CSV 文件
        # 假设文件不含复杂编码问题，使用默认读取
        df = pd.read_csv(input_csv_path)
    except FileNotFoundError:
        print(f"错误：输入文件未找到在路径: {input_csv_path}")
        return
    except Exception as e:
        print(f"读取 CSV 文件时发生错误: {e}")
        return
        
    # --- 2. 数据处理与清洗 ---
    
    # 检查 'sentence' 列是否存在
    if 'sentence' not in df.columns:
        print("错误：CSV 文件中未找到 'sentence' 列。请检查列名是否正确。")
        return
        
    # 提取 'sentence' 列，并去除 NaN 值
    sentences: List[str] = df['sentence'].dropna().astype(str).tolist()
    
    if not sentences:
        print("警告：'sentence' 列中没有有效数据，无法生成语料库。")
        return
    
    # 简单清洗：去除多余的空格或换行符（如果有的话）
    sentences = [s.strip() for s in sentences]
    
    print(f"提取到 {len(sentences)} 条有效句子。")

    # --- 3. 保存为纯文本文件 ---
    print(f"--- 3. 保存至纯文本文件: {output_txt_path} ---")
    
    # 将句子列表写入文件，每行一个句子
    try:
        with open(output_txt_path, 'w', encoding=encoding) as f:
            f.write('\n'.join(sentences))
        
        print(f"数据成功保存！")
        
    except Exception as e:
        print(f"写入文件时发生错误: {e}")
        
# --- 运行主函数 ---
if __name__ == "__main__":
    # 请确保您已安装 pandas: pip install pandas
    prepare_dapt_data(INPUT_CSV_PATH, OUTPUT_TXT_PATH, ENCODING)