#!/usr/bin/env python3
import os
import time
import requests
from concurrent.futures import ThreadPoolExecutor

API_KEY = "sk-e14298f02fd847dcbb827796fbc14b86"

with open('original_10_pages.txt', 'r', encoding='utf-8') as f:
    text = f.read()
    
paragraphs = text.split('\n\n')
chunks = []
current_chunk = ""
for p in paragraphs:
    if len(current_chunk) + len(p) > 2000:
        chunks.append(current_chunk)
        current_chunk = p
    else:
        current_chunk += "\n\n" + p if current_chunk else p
if current_chunk:
    chunks.append(current_chunk)

print(f"Total chunks: {len(chunks)}")

prompt = """Bạn là một biên tập viên chuyên nghiệp am hiểu Phật giáo. Hãy hiệu đính đoạn văn bản sau.
Yêu cầu: văn phong điềm đạm, gần gũi, dễ hiểu, giữ trạng thái văn nói nhưng không xuồng sã.
Sửa lại các câu cú không phù hợp, lủng củng, diễn đạt thô.
Giữ nguyên ý nghĩa cốt lõi và các thuật ngữ Phật giáo (như Bhikkhācāravattena, kappiyakāraka...).
CHỈ TRẢ VỀ NỘI DUNG ĐÃ HIỆU ĐÍNH, KHÔNG BÌNH LUẬN, KHÔNG GIẢI THÍCH, KHÔNG THÊM CÁC THẺ MARKDOWN."""

headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {API_KEY}"
}

def process_chunk(index, chunk):
    print(f"Starting chunk {index+1}/{len(chunks)}...")
    data = {
        "model": "deepseek-chat",
        "messages": [
            {"role": "system", "content": prompt},
            {"role": "user", "content": chunk}
        ],
        "temperature": 0.3
    }
    
    retries = 3
    while retries > 0:
        try:
            resp = requests.post("https://api.deepseek.com/chat/completions", headers=headers, json=data, timeout=120)
            resp.raise_for_status()
            res_json = resp.json()
            edited_text = res_json['choices'][0]['message']['content'].strip()
            print(f"Finished chunk {index+1}")
            return (index, edited_text)
        except Exception as e:
            print(f"Error on chunk {index+1}, retries left: {retries-1}. Exception: {e}")
            retries -= 1
            time.sleep(2)
            if retries == 0:
                return (index, "[ERROR_PROCESSING_CHUNK]\n" + chunk)

with ThreadPoolExecutor(max_workers=5) as executor:
    results = list(executor.map(lambda arg: process_chunk(*arg), enumerate(chunks)))

results.sort(key=lambda x: x[0])
edited_chunks = [r[1] for r in results]

with open('edited_10_pages.txt', 'w', encoding='utf-8') as f:
    f.write('\n\n'.join(edited_chunks))

print("Saved to edited_10_pages.txt")