#!/usr/bin/env python
import os
import sys
import subprocess
import tempfile
import shlex
import glob
from datetime import datetime

def edit_encrypted_file(filename):
    editor = os.environ.get('EDITOR', 'vi')
    
    with tempfile.NamedTemporaryFile(mode='w+', delete=True, suffix='.key') as id_file:
        try:
            result = subprocess.run(
                ['age', '-d', '.khat'],
                capture_output=True,
                check=True
            )
            id_file.write(result.stdout.decode())
            id_file.flush()
        except subprocess.CalledProcessError as e:
            print(f"Error decrypting .khat: {e}", file=sys.stderr)
            sys.exit(1)
        
        with tempfile.NamedTemporaryFile(mode='w+', delete=True, suffix='.txt') as tmp_file:
            if os.path.isfile(filename):
                try:
                    result = subprocess.run(
                        ['age', '-i', id_file.name, '-d', filename],
                        capture_output=True,
                        check=True
                    )
                    tmp_file.write(result.stdout.decode())
                    tmp_file.flush()
                except subprocess.CalledProcessError as e:
                    print(f"Error decrypting {filename}: {e}", file=sys.stderr)
                    sys.exit(1)
            else:
                tmp_file.write("""Hey.
# What's on your mind?

# What're you grateful for?

                """)
                tmp_file.flush()
            try:
                subprocess.run(shlex.split(editor) + [tmp_file.name], check=True)
            except subprocess.CalledProcessError as e:
                sys.exit(1)
            
            tmp_file.seek(0)
            try:
                subprocess.run(
                    ['age', '-e', '-i', id_file.name, '-a', '-o', filename],
                    input=tmp_file.read().encode(),
                    check=True
                )
            except subprocess.CalledProcessError as e:
                print(f"Error re-encrypting {filename}: {e}", file=sys.stderr)
                sys.exit(1)

def create_identity():
    print("No .khat file found. Creating new identity.")
    try:
        # Generate key
        keygen = subprocess.run(
            ['age-keygen'],
            capture_output=True,
            check=True
        )
        
        # Encrypt with passphrase
        subprocess.run(
            ['age', '-a', '-p', '-o', '.khat'],
            input=keygen.stdout,
            check=True
        )
        print("Identity created successfully in .khat")
    except subprocess.CalledProcessError as e:
        print(f"Error creating identity: {e}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    if not os.path.isfile('.khat'):
        create_identity()
    if len(sys.argv) == 1:
        edit_encrypted_file(datetime.now().strftime("%Y_%m_%d.age"))
    elif len(sys.argv) == 2 and sys.argv[1] != "choose":
        edit_encrypted_file(sys.argv[1])
    elif len(sys.argv) == 2 and sys.argv[1] == "choose":
        result = subprocess.run('fzf', input=('\n'.join(glob.glob("*.age"))), capture_output=True, text=True)
        edit_encrypted_file(result.stdout.strip())