#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import argparse
import glob
import time
try:
    from dotenv import load_dotenv
except:
    print("\nImport error: dotenv")
    input("Install: pip install dotenv")
    exit()

try:
    import openai
except ImportError:
    openai = None

try:
    import google.generativeai as genai
except ImportError:
    genai = None


script_path = os.path.abspath(__file__)
config_ini = os.path.splitext(script_path)[0] + '.ini'
config_path = "translate.env"

if not os.path.isfile(config_ini):
    script_dir = os.path.dirname(os.path.abspath(__file__))
    config_path = os.path.join(script_dir, config_path)

print()
if os.path.isfile(config_path):
    print(f"config_path: {config_path}")
else:
    print(f"Warning: config_path {config_path} is not found")
load_dotenv(dotenv_path=config_path)

account_inf_path = os.getenv("account_inf_path", "accounts.env")
if os.path.isfile(account_inf_path):
    print(f"account_inf_path: {account_inf_path}")
else:
    print(f"Warning: account_inf_path {account_inf_path} is not found")
load_dotenv(dotenv_path=account_inf_path)

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    print("\nWarning: OPENAI_API_KEY environment variable is not set.")

GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
if not GOOGLE_API_KEY:
    print("\nWarning: GOOGLE_API_KEY environment variable is not set.")

openai_model = os.getenv("openai_model", "gpt-4o-mini")
google_model = os.getenv("google_model", "gemini-2.5-flash")


prompt_templates = {
    "main": (
        "以下は{{language}}言語のプログラム{{script_name}}のコードです。これを解析し、\n"
        "1) プログラムの動作\n"
        "2) 必要な非標準ライブラリとインストールコマンドとインストール方法\n"
        "3) 必要な入力ファイル\n"
        "4) 実行後に生成される出力ファイル\n"
        "5) コマンドラインでの使用例 (Usage)\n"
        "を Markdown 形式で記述してください。\n\n"
        "```{{language}}\n"
        "{{code}}\n"
        "```"
        ),

    "lib": (
        "以下は{{language}}言語のライブラリ{{script_name}}のコードです。これを解析し、\n"
        "1) このライブラリの主な機能や目的\n"
        "2) このライブラリを他のプログラムからimportする方法\n"
        "3) 必要な非標準ライブラリとインストールコマンドとインストール方法\n"
        "4) importできる変数と関数。変数には説明するコメントを、関数には関数の動作と引数・戻り値の説明を入れてください\n"
        "5) main scriptとして実行したときの動作\n"
        "を Markdown 形式で記述してください。\n\n"
        "```{{language}}\n"
        "{{code}}\n"
        "```"
        ),
    }

language_dict = {
    ".py": "python",
    ".pl": "perl",
    ".pm": "perl",
    ".c" : "C",
    ".cpp": "C++",
    ".pas": "pascal",
    ".f"  : "fortran",
    ".for": "fortran",
    ".f77": "fortran",
    ".js" : "Javascript",
    ".java": "Java",
    ".go": "Go",
    ".sh": "bash script",
    ".csh": "tcsh script",
    ".bat": "Windows batch script",
    ".vbs": "Visual Basic script",
    ".css": "Cascade Style Sheet",
    ".html": "HTML",
    ".xml": "XML",
    }


def get_program_type(path):
    base = os.path.basename(path)
    name, ext = os.path.splitext(base)
    if ext == ".pm": return 'lib'
    if base.startswith("tk"): return 'lib'
    return 'main'

def initialize():
    parser = argparse.ArgumentParser(
        description="プログラムコードを読み込み、OpenAI/Google API で使用方法を Markdown に出力するツール"
    )
    parser.add_argument(
        "pattern",
        help="処理対象ファイルのワイルドカードパターン（例: '*.py'）"
    )
    parser.add_argument(
        "output",
        nargs="?",
        help="出力 Markdown ファイル名。省略した場合は入力ファイルごとに拡張子を .md に置換"
    )
    parser.add_argument(
        "--program_type",
        choices=["", "main", "lib"],
        default="",
        help="読み込むプログラムのタイプ ['', 'main', 'lib']（デフォルト: ''）"
    )
    parser.add_argument(
        "--api",
        choices=["openai", "google"],
        default="openai",
        help="使用する API を指定（デフォルト: openai）"
    )
    parser.add_argument(
        "--max_tokens", type = int,
        default=3000,
        help="APIの最大token数（デフォルト: 3000）"
    )
    parser.add_argument(
        "--update", type=int, choices=[0, 1], default=0,
        help="出力ファイルが存在し、かつ入力ファイルより古い場合にのみ再生成"
    )
    parser.add_argument(
        "--overwrite", type=int, choices=[0, 1], default=0,
        help="既存ファイルを問答無用で上書き生成"
    )
    return parser

def generate_doc_openai(inp: str, code: str, program_type: str = '') -> str:
    """OpenAI API を使ってコードからドキュメントを生成（openai>=1.0.0 対応）"""
    if openai is None:
        raise RuntimeError("openai パッケージが見つかりません。pip install openai してください。")
    api_key = OPENAI_API_KEY
    if not api_key:
        raise RuntimeError("環境変数 OPENAI_API_KEY が設定されていません。")
    openai.api_key = api_key

    ext = os.path.splitext(inp)[1]
    language = language_dict.get(ext, None)
    if not language:
        print(f"\n{inp}はサポートされていないプログラミング言語です。スキップします")
        return None

    if program_type == '':
        program_type = get_program_type(inp)
    print(f"Program type: {program_type}")

    prompt_template = prompt_templates.get(program_type, prompt_templates['main'])
    prompt = prompt_template.replace("{{script_name}}", inp).replace("{{code}}", code).replace("{{language}}", language)
    resp = openai.chat.completions.create(
        model=openai_model,
        messages=[
            {"role": "system", "content": "You are a helpful assistant that generates program documentation."},
            {"role": "user",   "content": prompt}
        ],
        temperature=0,
        max_tokens=max_tokens
    )
    return resp.choices[0].message.content

def generate_doc_google(inp: str, code: str, program_type: str = '') -> str:
    """Google AI Gemini API を使ってコードからドキュメントを生成"""
    if genai is None:
        raise RuntimeError("google-generativeai パッケージが見つかりません。pip install google-generativeai してください。")

    api_key = GOOGLE_API_KEY
    if not api_key:
        raise RuntimeError("環境変数 GOOGLE_API_KEY が設定されていません。")
    genai.configure(api_key=api_key)

    ext = os.path.splitext(inp)[1]
    language = language_dict.get(ext, None)
    if not language:
        print(f"\n{inp} はサポートされていないプログラミング言語です。スキップします")
        return None

    if program_type == '':
        program_type = get_program_type(inp)
    print(f"Program type: {program_type}")

    prompt_template = prompt_templates.get(program_type, prompt_templates['main'])
    prompt = prompt_template.replace("{{script_name}}", inp).replace("{{code}}", code).replace("{{language}}", language)

    # Gemini モデルを初期化
    model = genai.GenerativeModel(google_model)

    # generate_content を使ってドキュメントを生成
    try:
        response = model.generate_content(
            contents=[
                {"role": "user", "parts": [prompt]}
            ]
        )
        return response.text
    except Exception as e:
        print(f"ドキュメント生成中にエラーが発生しました: {e}")
        return None


def main():
    parser = initialize()
    args = parser.parse_args()

    files = glob.glob(args.pattern)
    if not files:
        print(f"Error: パターン '{args.pattern}' にマッチするファイルがありません。", file=sys.stderr)
        sys.exit(1)

    # 出力ファイル名の決定
    if args.output:
        if len(files) > 1:
            print("Error: 複数ファイルを処理するときは output 引数を省略してください。", file=sys.stderr)
            sys.exit(1)
        outputs = [args.output]
    else:
        outputs = [os.path.splitext(f)[0] + ".md" for f in files]

    # 各ファイルを処理
    for inp, out in zip(files, outputs):
        do_generate = True
        if os.path.exists(out):
            if args.overwrite:
                do_generate = True
            elif args.update:
                inp_mtime = os.path.getmtime(inp)
                out_mtime = os.path.getmtime(out)
                if out_mtime >= inp_mtime:
                    print(f"[INFO] '{out}' は最新です。スキップします。")
                    do_generate = False
                else:
                    do_generate = True
            else:
                print(f"[INFO] '{out}' が既に存在します。--update または --overwrite を指定して再生成できます。スキップします。")
                do_generate = False
        if not do_generate:
            continue

        print(f"[INFO] Generating doc: {inp} → {out}")
        try:
            code = open(inp, encoding="utf-8").read()
        except:
            try:
                code = open(inp, encoding="shift-jis").read()
            except:
                code = None

        if code is None:
            print(f"\nError: Could not read [{inp}]. Skip.")
            continue
            
        if args.api == "openai":
            doc = generate_doc_openai(inp, code, args.program_type)
        else:
            doc = generate_doc_google(inp, code, args.program_type)

        if doc is None:
            print(f"Error generating doc for {inp}")
            continue

        with open(out, "w", encoding="utf-8") as f:
            f.write(doc)

        # API レート制限回避
        time.sleep(1)

    print("Done.")
    
    input("\nPress ENTER to terminate>>\n")
    

if __name__ == "__main__":
    main()