import argparse
import os
import shutil
import glob
import io
try:
    from openpyxl import Workbook, load_workbook
except:
    print("Error: Failed to import openpyxl")
    input("Install: pip install openpyxl")
try:
    import PyPDF2
except:
    print("Error: Failed to import PyPDF2")
    input("Install: pip install pypdf2")

from tkai_lib import read_ai_config
from tkai_lib import query_openai4, openai_response_to_json
from tkai_lib import query_openai5, openai5_response_to_json
from tkai_lib import query_google, google_response_to_json


#=========================================================
# Global configuration
#=========================================================
config_path = "ai.env"

#=========================================================
# Read configuration
#=========================================================
read_ai_config(config_path)

max_tokens     = int(os.getenv("max_tokens", "2000"))

openai_api_key = os.getenv("OPENAI_API_KEY")
openai_model   = os.getenv("openai_model", "gpt-4o")
temperature    = float(os.getenv("temperature", "0.3"))

openai_model5  = os.getenv("openai_model5", "gpt-5-nano")
reasoning_effort = os.getenv("reasoning_effort", "low")

google_api_key = os.getenv("GOOGLE_API_KEY")
google_model   = os.getenv("gemini_model", "gemini-2.5-flash")


#=========================================================
# Parameters
#=========================================================
DEFAULT_TEMPLATE = "(author_first)_(author_last)_(short_title)_(shortest_name)_(year).pdf"

xlsx_labels = [
    "directory", "filename_original", "filename_rename", 
    "authors", "author_first", "author_last", 
    "title", "short_title", "journal", "short_name", "shortest_name", 
    "date_received", "date_accepted", "date_published_online", "date_published",
    "year", "volume", "issue", "pages", "doi",
    ]

# Prompt template: start
prompt_template = """
以下の学術論文のテキストから、指定された情報をJSON形式で抽出してください。
また、first author、last author、短縮したtitle (short_title)、shortest_name、発行年から
推奨ファイル名 {{template}} を作ってください。
short_titleの単語はCaptalFirstにして、空白、.、,などの文字は削除してください。

推奨ファイル名（filename_rename）:
著者リスト（authors）   ：
第一著者（author_first）：
最終著者（author_last） ：
論文題目（title）           :
論文題目略称（short_title） :
発行雑誌名（journal）       ：
発行雑誌名略称（short_name）：
発行雑誌名の最短の略称（shortest_name）：
投稿受理日（date_received） :
採択日（date_accepted） :
オンライン発行日（date_published_online）:
発行日（date_published）:
発行年（year）   :
巻番号（volume）：
号番号（issue） ：
ページ（pages ）：
DOI（doi）：

情報が不明な場合は、空文字列 "" を使用してください。

--- 論文テキスト ---
{{text}}
"""
# Prompt template: end


def initialize():
    parser = argparse.ArgumentParser(description="学術論文のPDFからメタデータを抽出します。")
    parser.add_argument("--api", type=str, default='openai5',
        choices = ["openai", "openai5", "google"],
        help="使用する API を指定（openai5|openai|google: デフォルト: openai5）")
    parser.add_argument("input_file", type=str, help="処理するPDFファイルのパス")
    parser.add_argument("--summary_path", type=str, default = "summary.xlsx", 
                        help="summary.xlsxのパス")
    parser.add_argument("--recursive", type=int, default=0, 
                        help="1の場合、サブディレクトリまで処理する")
    parser.add_argument("--rename", type=int, default=0, 
                        help="1の場合、生成AIで推薦ファイル名を取得できた場合、リネームする")
#    parser.add_argument("--rename", action="store_true", 
#                        help="生成AIで推薦ファイル名を取得できた場合、リネームする")
    parser.add_argument("--delete_original", type=int, default=1, 
                        help="1の場合、renameしたら元のファイルを削除する")
    parser.add_argument("--max_bytes", type=int, default=10000,
                        help="APIに送る最大テキストバイト数。論文全体を送る必要がない場合に指定。")
    parser.add_argument("--template", type=str, default=DEFAULT_TEMPLATE,
                        help="ファイル名を変える際のテンプレート {{DEFAULT_TEMPLATE}}")
    return parser

def extract_text_from_pdf(pdf_path: str, max_bytes: int) -> str:
    """
    PDFファイルからテキストを抽出し、指定された最大バイト数に制限する。
    """
    text_buffer = io.StringIO()
    try:
        with open(pdf_path, 'rb') as f:
            pdf_reader = PyPDF2.PdfReader(f)
            num_pages = len(pdf_reader.pages)
            
            for page_num in range(num_pages):
                page = pdf_reader.pages[page_num]
                text_buffer.write(page.extract_text() or '')
                
                # バイト数チェック
                if text_buffer.tell() > max_bytes:
                    break
    except Exception as e:
        print(f"PDFファイルの読み込み中にエラーが発生しました: {e}")
        return ""

    full_text = text_buffer.getvalue()
    
    # バイト数でテキストを切り詰める
    encoded_text = full_text.encode('utf-8', 'ignore')
    if len(encoded_text) > max_bytes:
        truncated_text = encoded_text[:max_bytes].decode('utf-8', 'ignore')
        return truncated_text
    
    return full_text

def get_metadata_from_openai(text: str, template: str) -> dict:
    """
    OpenAI APIを使用して論文のメタデータを抽出する。
    """

    prompt = prompt_template.replace("{{text}}", text).replace("{{template}}", template)
    response = query_openai4(prompt, openai_model, max_tokens = max_tokens, response_format = { "type": "json_object" },
                             openai_api_key = openai_api_key)
    if not response: return False
    if response == {}: return {}
    
    json, json_list = openai_response_to_json(response)
    return json

def get_metadata_from_openai5(text: str, template: str) -> dict:
    """
    OpenAI APIを使用して論文のメタデータを抽出する。
    """

    prompt = prompt_template.replace("{{text}}", text).replace("{{template}}", template)
    response = query_openai5(prompt, openai_model5, openai_api_key = openai_api_key, 
                             effort = reasoning_effort, max_output_tokens = max_tokens)
    if not response: return False
    if response == {}: return {}
    
    json, json_list = openai5_response_to_json(response)
    return json

def get_metadata_from_google(text: str, template: str) -> dict:
    """
    Google APIを使用して論文のメタデータを抽出する。
    """

    prompt = prompt_template.replace("{{text}}", text).replace("{{template}}", template)
    response = query_google(prompt, google_model, 
                role = None, generation_config = {"response_mime_type": "application/json"} ,
                google_api_key = google_api_key)
    if not response: return False
    if response == {}: return {}

    json, json_list = google_response_to_json(response)
    return json

def append_xlsx(path, mode, labels, data_list):
    if mode == "w" or not os.path.exists(path):
        wb = Workbook()
        ws = wb.active
        ws.append(labels)
    else:
        wb = load_workbook(filename = path)
        ws = wb.active

    for row in data_list:
        ws.append(row)

    wb.save(path)

def rename_file(input_file, new_filename, delete_original = False):
    if os.path.exists(new_filename):
        print(f"Error in rename_file(): File [{new_filename}] already exists. Skip.")
        return False

    try:
        directory = os.path.dirname(input_file)
        os.chdir(directory)
        if delete_original:
            os.rename(input_file, new_filename)
            print(f"Changed file name from [{input_file}] to [{new_filename}]")
        else:
            shutil.copy(input_file, new_filename)
            print(f"Copied [{input_file}] to [{new_filename}]")
    except Exception as e:
        print(f"Error in rename_file(): Could not change file name to [{new_filename}]")

    return True

def get_inf(input_file, summary_path, api = "openai", max_bytes = 10000, 
            rename = False, delete_original = True, template = DEFAULT_TEMPLATE):
    directory = os.path.dirname(input_file)
    filename  = os.path.basename(input_file)
    print(f"Input PDF file: [{input_file}]")
    print(f"  Directory: [{directory}]")
    print(f"  File name: [{filename}]")

    extracted_text = extract_text_from_pdf(input_file, max_bytes)
    if not extracted_text:
        print("テキストの抽出に失敗したか、ファイルが空です。処理を終了します。")
        print("PyCryptodomeが必要というメッセージが出たら、以下のようにinstallしてください。")
        print("Install: pip install pycryptodome ")
        return False

    print(f"  API: {api.upper()}")
    if 'openai5' in api:
        print(f"  openai_model: {openai_model5}")
        print(f"  effort      : {reasoning_effort}")
        print(f"  max_tokens  : {max_tokens}")
    elif 'openai' in api:
        print(f"  openai_model: {openai_model}")
        print(f"  temperature : {temperature}")
        print(f"  max_tokens  : {max_tokens}")
    else:
        print(f"  google_model: {google_model}")
        print(f"  max_tokens: {max_tokens}")

    print(f"  max bytes: {max_bytes}")
    print(f"  抽出したテキストを送信中...")
    if api == 'openai':
        metadata = get_metadata_from_openai(extracted_text[:max_bytes], template)
    elif api == 'openai5':
        metadata = get_metadata_from_openai5(extracted_text[:max_bytes], template)
    elif api == 'google':
        metadata = get_metadata_from_google(extracted_text[:max_bytes], template)
    else:
        print("Error in get_inf(): api が 'openai' または 'google' ではありません。")
        return None

    if metadata:
        print("\n--- 論文メタデータ ---")
        data = [directory, filename]
        for key, value in metadata.items():
            print(f"{key}: {value}")
            if type(value) is list or type(value) is tuple:
                data.append(", ".join(value))
            else:
                data.append(value)
    else:
        data = [directory, filename, "", "", "", "", "", 
                "", "", "", "", "", 
                "", "", "", "", "", "", "", "", ""]

    append_xlsx(summary_path, mode = "a", labels = xlsx_labels, data_list = [data])

    if rename and metadata.get("filename_rename", None):
        new_filename = metadata["filename_rename"]
        if not new_filename.endswith('.pdf'): new_filename += '.pdf'

        ret = rename_file(input_file, new_filename, delete_original = delete_original)
        if not ret: return False

    
def main():
    parser = initialize()
    args = parser.parse_args()
    
    fmask = args.input_file
    if args.recursive:
        directory = os.path.dirname(fmask)
        base_name = os.path.basename(fmask)
        if '*' not in directory:
            fmask = os.path.join(directory, '**', base_name)

    print()
    print(f"Search files for [{fmask}]")
    print(f"Recursive search? {args.recursive}")
    files = sorted(glob.glob(fmask, recursive = args.recursive))
    print()
    if len(files) == 0:
        print(f"Error: No file found.")
    else:
        print(f"Files found for [{args.input_file}]")
        for f in files:
            print(f"  {f}")

        for f in files:
            print()
            get_inf(f, args.summary_path, args.api, args.max_bytes, args.rename, args.delete_original, args.template)

    input("\nPress ENTER to terminate>>\n")


if __name__ == "__main__":
    main()