import os
import base64
import requests
import json
from pathlib import Path

def get_access_token(api_key, secret_key):
    url = f"https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={api_key}&client_secret={secret_key}"
    response = requests.request("POST", url)
    return response.json().get("access_token")

def test_table_v2(image_path, api_key, secret_key):
    # 表格识别V2接口
    request_url = "https://aip.baidubce.com/rest/2.0/ocr/v1/table"
    
    with open(image_path, "rb") as f:
        img = base64.b64encode(f.read())

    params = {
        "image": img,
        "is_sync": "true",
        "request_type": "excel"
    }
    access_token = get_access_token(api_key, secret_key)
    request_url = request_url + "?access_token=" + access_token
    headers = {'content-type': 'application/x-www-form-urlencoded'}
    response = requests.post(request_url, data=params, headers=headers)
    
    return response.json()

if __name__ == "__main__":
    api_key = "kIehdWbbVD85K18qZYz6SeUf"
    secret_key = "RCXTgmVjJJkNNMhfY5ASab0xY3mvF6d7"
    
    input_dir = Path(r"F:\Administrator\桌面\OCR系统_Release\data\input")
    images = list(input_dir.glob("*.jpg")) + list(input_dir.glob("*.png"))
    
    if not images:
        print("错误: 桌面 Release 目录下没有找到图片文件")
    else:
        latest_image = max(images, key=os.path.getmtime)
        print(f"正在使用【表格识别V2】接口测试图片: {latest_image.name}")
        
        result = test_table_v2(str(latest_image), api_key, secret_key)
        
        # 打印表格外的文字（header/footer）
        if "tables_result" in result:
            print("\n--- 表格识别结果 ---")
            for table in result["tables_result"]:
                print("\n[Header 区域]:")
                header = table.get("header", [])
                for cell in header:
                    print(f"- {cell.get('words')}")
                
                print("\n[Footer 区域]:")
                footer = table.get("footer", [])
                for cell in footer:
                    print(f"- {cell.get('words')}")
        elif "result" in result and "tables_result" in result["result"]:
            print("\n--- 表格识别结果 (嵌套结构) ---")
            for table in result["result"]["tables_result"]:
                print("\n[Header 区域]:")
                header = table.get("header", [])
                for cell in header:
                    print(f"- {cell.get('words')}")
                
                print("\n[Footer 区域]:")
                footer = table.get("footer", [])
                for cell in footer:
                    print(f"- {cell.get('words')}")
        else:
            print(f"识别失败或未找到表格: {json.dumps(result, indent=2, ensure_ascii=False)[:500]}")