# file: data_processor.py
import json
import re

def process_law_data_to_chunks(structured_data_input):
    flat_list = []

    if isinstance(structured_data_input, dict) and "article" in structured_data_input:
        articles_list = [structured_data_input]
    elif isinstance(structured_data_input, list):
        articles_list = structured_data_input
    else:
        print("Lỗi: Dữ liệu đầu vào không hợp lệ.")
        return flat_list

    for article_data in articles_list:
        if not isinstance(article_data, dict):
            continue

        article_metadata_base = {
            "source_document": article_data.get("source_document"),
            "article": article_data.get("article"),
            "article_title": article_data.get("article_title")
        }

        clauses = article_data.get("clauses", [])
        if not isinstance(clauses, list):
            continue

        for clause_data in clauses:
            if not isinstance(clause_data, dict):
                continue

            clause_metadata_base = article_metadata_base.copy()
            clause_metadata_base.update({
                "clause_number": clause_data.get("clause_number"),
                "clause_metadata_summary": clause_data.get("clause_metadata_summary")
            })

            points_in_clause = clause_data.get("points_in_clause", [])
            if not isinstance(points_in_clause, list):
                continue
            
            if points_in_clause:
                for point_data in points_in_clause:
                    if not isinstance(point_data, dict):
                        continue
                    
                    # <<< THAY ĐỔI BẮT ĐẦU: LÀM GIÀU VĂN BẢN >>>
                    
                    # 1. Thu thập các thành phần văn bản để làm giàu
                    article_title = article_metadata_base.get('article_title', '')
                    point_text = point_data.get("point_text_original") or point_data.get("violation_description_summary")
                    
                    # Lấy thông tin tóm tắt của Khoản (thường là mức phạt chung)
                    clause_summary_dict = clause_data.get("clause_metadata_summary", {})
                    clause_summary_text = ""
                    if clause_summary_dict:
                        # Lấy giá trị từ các key có thể có
                        summary_keys = ["overall_fine_note_for_clause", "overall_points_deduction_note_for_clause"]
                        for key in summary_keys:
                            if key in clause_summary_dict:
                                clause_summary_text = clause_summary_dict[key]
                                break
                    
                    # Nếu không có tóm tắt ở Khoản, thử lấy trực tiếp từ text gốc của Khoản
                    if not clause_summary_text:
                        clause_original_text = clause_data.get("clause_text_original", "")
                        # Chỉ lấy dòng đầu tiên làm tóm tắt (thường là dòng mức phạt)
                        clause_summary_text = clause_original_text.split('\n')[0]

                    # 2. Tạo chuỗi văn bản giàu ngữ cảnh
                    text_parts = [
                        part.strip() for part in [article_title, clause_summary_text, point_text] if part
                    ]
                    # Dùng ": " để nối các phần, giúp phân tách ngữ cảnh
                    enriched_text = ": ".join(text_parts)
                    
                    # <<< THAY ĐỔI KẾT THÚC >>>

                    if not enriched_text:
                        continue

                    current_point_metadata = clause_metadata_base.copy()
                    
                    # Giữ lại toàn bộ thông tin chi tiết trong metadata
                    point_specific_metadata = point_data.copy()
                    
                    # Xóa trường text gốc khỏi metadata để tránh trùng lặp không cần thiết
                    if "point_text_original" in point_specific_metadata:
                        del point_specific_metadata["point_text_original"]

                    current_point_metadata.update(point_specific_metadata)
                    
                    final_metadata_cleaned = {k: v for k, v in current_point_metadata.items() if v is not None}

                    flat_list.append({
                        # Sử dụng văn bản đã được làm giàu
                        "text": enriched_text,
                        "metadata": final_metadata_cleaned
                    })
            else:
                # Xử lý các Khoản không có Điểm
                chunk_text = clause_data.get("clause_text_original")
                if chunk_text:
                    current_clause_metadata = clause_metadata_base.copy()
                    additional_clause_info = {k: v for k, v in clause_data.items() if k not in ["clause_text_original", "points_in_clause", "clause_number", "clause_metadata_summary"]}
                    current_clause_metadata.update(additional_clause_info)
                    final_metadata_cleaned = {k: v for k, v in current_clause_metadata.items() if v is not None}
                    
                    # <<< THAY ĐỔI: Cũng làm giàu văn bản cho các Khoản đứng một mình >>>
                    article_title = article_metadata_base.get('article_title', '')
                    enriched_text = f"{article_title}: {chunk_text}" if article_title else chunk_text

                    flat_list.append({
                        "text": enriched_text,
                        "metadata": final_metadata_cleaned
                    })
    return flat_list