import json
import base64
import time
import urllib.request
import subprocess
import os
import re

def base64url_encode(data):
    return base64.urlsafe_b64encode(data).decode('utf-8').replace('=', '')

def get_access_token():
    with open('/home/openclaw/.openclaw/workspace/google_service_account.json') as f:
        creds = json.load(f)
    private_key_str = creds['private_key']
    client_email = creds['client_email']
    token_uri = creds['token_uri']
    header = {"alg": "RS256", "typ": "JWT"}
    header_b64 = base64url_encode(json.dumps(header).encode())
    now = int(time.time())
    payload = {
        "iss": client_email,
        "scope": "https://www.googleapis.com/auth/documents",
        "aud": token_uri,
        "iat": now,
        "exp": now + 3600
    }
    payload_b64 = base64url_encode(json.dumps(payload).encode())
    content = f"{header_b64}.{payload_b64}"
    key_file = '/tmp/priv_cleaner.pem'
    with open(key_file, 'w') as f:
        f.write(private_key_str)
    try:
        process = subprocess.Popen(
            ['openssl', 'dgst', '-sha256', '-sign', key_file],
            stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )
        signature, err = process.communicate(input=content.encode())
        if process.returncode != 0:
            raise Exception(f"OpenSSL error: {err.decode()}")
    finally:
        if os.path.exists(key_file):
            os.remove(key_file)
    signature_b64 = base64url_encode(signature)
    jwt = f"{content}.{signature_b64}"
    data = urllib.parse.urlencode({
        "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
        "assertion": jwt
    }).encode()
    req = urllib.request.Request(token_uri, data=data)
    with urllib.request.urlopen(req) as f:
        res = json.loads(f.read().decode())
    return res['access_token']

def clean_page_markers():
    access_token = get_access_token()
    doc_id = "1Ymh60_Qe2yX6TgfznwOHwlphv8YMoz3SscJNg_HZtO0"
    
    get_uri = f"https://docs.googleapis.com/v1/documents/{doc_id}"
    req = urllib.request.Request(get_uri, headers={"Authorization": f"Bearer {access_token}"})
    with urllib.request.urlopen(req) as f:
        doc = json.loads(f.read().decode())
    
    content = doc.get('body', {}).get('content', [])
    requests = []
    
    # Regex to match "Trang" followed by a space and numbers
    # Pattern: ^\s*Trang\s+\d+\s*$
    page_pattern = re.compile(r'^\s*Trang\s+\d+\s*$', re.IGNORECASE)
    
    # Track indices to delete
    # Note: When we delete a paragraph, we should delete its range [startIndex, endIndex)
    
    # Step 1: Identify "Trang xx" paragraphs
    for i in range(len(content) - 1, 0, -1):
        element = content[i]
        if 'paragraph' in element:
            para_text = "".join([el.get('textRun', {}).get('content', '') for el in element['paragraph'].get('elements', [])])
            if page_pattern.match(para_text.strip()):
                requests.append({
                    "deleteContentRange": {
                        "range": {
                            "startIndex": element['startIndex'],
                            "endIndex": element['endIndex']
                        }
                    }
                })
    
    if requests:
        update_uri = f"https://docs.googleapis.com/v1/documents/{doc_id}:batchUpdate"
        req = urllib.request.Request(
            update_uri, 
            data=json.dumps({"requests": requests}).encode(),
            headers={
                "Authorization": f"Bearer {access_token}",
                "Content-Type": "application/json"
            },
            method='POST'
        )
        with urllib.request.urlopen(req) as f:
            final_res = json.loads(f.read().decode())
        print(f"Deleted {len(requests)} page markers.")
    else:
        print("No page markers found.")

    # Step 2: Fix duplicate spacers (ensure only one empty line between paragraphs)
    # Re-fetch document after deletions
    req = urllib.request.Request(get_uri, headers={"Authorization": f"Bearer {access_token}"})
    with urllib.request.urlopen(req) as f:
        doc = json.loads(f.read().decode())
    
    content = doc.get('body', {}).get('content', [])
    spacer_requests = []
    
    # We look for sequences of empty paragraphs
    empty_streak = 0
    for i in range(len(content) - 1, 0, -1):
        element = content[i]
        if 'paragraph' in element:
            para_text = "".join([el.get('textRun', {}).get('content', '') for el in element['paragraph'].get('elements', [])]).strip()
            if para_text == "":
                empty_streak += 1
                if empty_streak > 1:
                    # More than one empty paragraph in a row, delete this one
                    spacer_requests.append({
                        "deleteContentRange": {
                            "range": {
                                "startIndex": element['startIndex'],
                                "endIndex": element['endIndex']
                            }
                        }
                    })
            else:
                empty_streak = 0
                
    if spacer_requests:
        update_uri = f"https://docs.googleapis.com/v1/documents/{doc_id}:batchUpdate"
        req = urllib.request.Request(
            update_uri, 
            data=json.dumps({"requests": spacer_requests}).encode(),
            headers={
                "Authorization": f"Bearer {access_token}",
                "Content-Type": "application/json"
            },
            method='POST'
        )
        with urllib.request.urlopen(req) as f:
            final_res = json.loads(f.read().decode())
        print(f"Removed {len(spacer_requests)} extra blank lines.")
    else:
        print("No extra blank lines found.")

if __name__ == "__main__":
    try:
        clean_page_markers()
    except Exception as e:
        print(f"Error: {e}")