import argparse
import json
import logging

from langchain_core.documents import Document
from langchain_text_splitters import CharacterTextSplitter

logger = logging.getLogger(__name__)

def chunk_documents(input_file_path: str, output_file_path: str):
    """
    JSONファイルをロードし、ドキュメントをチャンクに分割して、新しいJSONファイルに保存します。
    """
    # テキスト分割を初期化します
    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=50)

    # 入力JSONファイルからドキュメントを読み込みます
    with open(input_file_path, 'r', encoding='utf-8') as f:
        documents_data = json.load(f)

    chunked_documents = []
    for doc_data in documents_data:
        page_content = doc_data.get("page_content", "")
        metadata = doc_data.get("metadata", {})

        # page_contentをチャンクに分割します
        splits_texts = text_splitter.split_text(page_content)

        # 分割ごとに新しいドキュメントチャンクを作成します
        for i, text in enumerate(splits_texts):
            chunk_metadata = metadata.copy()
            chunk_metadata['chunk_index'] = i
            
            chunked_doc = {
                "metadata": chunk_metadata,
                "page_content": text
            }
            chunked_documents.append(chunked_doc)

    # チャンク化されたドキュメントを出力JSONファイルに保存します
    with open(output_file_path, 'w', encoding='utf-8') as f:
        json.dump(chunked_documents, f, ensure_ascii=False, indent=2)

    logger.info(f"{len(documents_data)}個のドキュメントを{len(chunked_documents)}個のチャンクに正常に分割しました。")
    logger.info(f"出力は {output_file_path} に保存されました")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="JSONファイルからドキュメントをチャンク化します。")
    parser.add_argument("input_file", help="入力JSONファイルのパス (例: documents.json)。")
    parser.add_argument("output_file", help="出力JSONファイルのパス (例: chunks.json)。")
    
    args = parser.parse_args()
    
    chunk_documents(args.input_file, args.output_file)