#!/usr/bin/env python3
"""
Utility script to manage content encryption for markmywords.

Usage:
    python manage_content.py encrypt <filepath> <password>
    python manage_content.py decrypt <filepath> <password>
    python manage_content.py view <filepath> <password>
    python manage_content.py unprotect <filepath>
    python manage_content.py list [directory]
    python manage_content.py status <filepath>

Notes:
    - encrypt: Creates .md.enc file with automatic password protection
    - Password protects the file AND encrypts it (same password for both)
    - unprotect: Removes password protection (file stays encrypted, viewable in browser)
    - Same password used for both encryption and browser authentication
"""

import json
import os
import sys
from pathlib import Path
from werkzeug.security import generate_password_hash, check_password_hash
from datetime import datetime
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.backends import default_backend
import base64
import markdown

PASSWORD_FILE = "config/page_passwords.json"
CONTENT_DIR = "content"
MARKDOWN_EXTENSIONS = ['.md', '.markdown']


def derive_encryption_key(password):
    """Derive an encryption key from a password."""
    salt = b'markmywords_salt'  # Fixed salt for consistency
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
        backend=default_backend()
    )
    key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
    return key


def encrypt_file_content(content, password):
    """Encrypt file content using password-derived key."""
    try:
        key = derive_encryption_key(password)
        f = Fernet(key)
        encrypted = f.encrypt(content.encode('utf-8'))
        return encrypted.decode('utf-8')
    except Exception as e:
        print(f"Error encrypting content: {e}", file=sys.stderr)
        return None


def decrypt_file_content(encrypted_content, password):
    """Decrypt file content using password-derived key."""
    try:
        key = derive_encryption_key(password)
        f = Fernet(key)
        decrypted = f.decrypt(encrypted_content.encode('utf-8'))
        return decrypted.decode('utf-8')
    except Exception as e:
        print(f"Error decrypting content: {e}", file=sys.stderr)
        return None


def is_encrypted_file(filepath):
    """Check if a file is encrypted (has .enc extension)."""
    return filepath.endswith('.enc')


def load_passwords():
    """Load password database from file."""
    if not os.path.exists(PASSWORD_FILE):
        return {}
    try:
        with open(PASSWORD_FILE, 'r') as f:
            return json.load(f)
    except Exception as e:
        print(f"Error loading passwords: {e}", file=sys.stderr)
        return {}


def save_passwords(passwords):
    """Save password database to file."""
    try:
        os.makedirs(os.path.dirname(PASSWORD_FILE), exist_ok=True)
        with open(PASSWORD_FILE, 'w') as f:
            json.dump(passwords, f, indent=2)
        os.chmod(PASSWORD_FILE, 0o600)  # Restrict file permissions
        return True
    except Exception as e:
        print(f"Error saving passwords: {e}", file=sys.stderr)
        return False


def get_file_key(filepath):
    """Generate a consistent key for a file."""
    return filepath


def validate_password(password):
    """Validate password strength."""
    if not password:
        print("Error: Password cannot be empty", file=sys.stderr)
        return False
    
    if len(password) < 4:
        print("Error: Password should be at least 4 characters", file=sys.stderr)
        return False
    
    return True


def encrypt(filepath, password):
    """Encrypt a markdown file with automatic password protection."""
    if not validate_password(password):
        return False
    
    # Normalize path
    if not filepath.startswith(CONTENT_DIR):
        filepath = os.path.join(CONTENT_DIR, filepath)
    
    # Remove .enc if present in input
    if filepath.endswith('.enc'):
        filepath = filepath[:-4]
    
    if not os.path.exists(filepath):
        print(f"✗ File not found: {filepath}", file=sys.stderr)
        return False
    
    # Check file extension
    if not any(filepath.lower().endswith(ext) for ext in MARKDOWN_EXTENSIONS):
        print(f"✗ Not a markdown file: {filepath}", file=sys.stderr)
        return False
    
    try:
        # Read original content
        with open(filepath, 'r', encoding='utf-8') as f:
            original_content = f.read()
        
        # Encrypt content
        encrypted_content = encrypt_file_content(original_content, password)
        if encrypted_content is None:
            return False
        
        # Write encrypted file
        enc_filepath = filepath + '.enc'
        with open(enc_filepath, 'w', encoding='utf-8') as f:
            f.write(encrypted_content)
        
        # Delete original
        os.remove(filepath)
        
        # Update password database with encryption AND protection (same password)
        # Normalize file_key to relative path (without CONTENT_DIR prefix)
        display_path = enc_filepath.replace(CONTENT_DIR + '/', '')
        file_key = get_file_key(display_path)
        passwords = load_passwords()
        
        passwords[file_key] = {
            'encrypted': True,
            'encrypted_at': datetime.now().isoformat(),
            'protected': True,
            'password_hash': generate_password_hash(password),
            'created_at': datetime.now().isoformat()
        }
        
        save_passwords(passwords)
        
        print(f"✓ Encrypted & protected: {filepath} → {display_path}")
        return True
    
    except Exception as e:
        print(f"✗ Error encrypting {filepath}: {e}", file=sys.stderr)
        return False


def decrypt(filepath, password):
    """Decrypt a file and print content."""
    # Normalize path
    if not filepath.startswith(CONTENT_DIR):
        filepath = os.path.join(CONTENT_DIR, filepath)
    
    # Ensure .enc extension
    if not filepath.endswith('.enc'):
        filepath = filepath + '.enc'
    
    if not os.path.exists(filepath):
        print(f"✗ File not found: {filepath}", file=sys.stderr)
        return False
    
    try:
        # Read encrypted content
        with open(filepath, 'r', encoding='utf-8') as f:
            encrypted_content = f.read()
        
        # Decrypt
        decrypted_content = decrypt_file_content(encrypted_content, password)
        if decrypted_content is None:
            print(f"✗ Failed to decrypt {filepath}", file=sys.stderr)
            return False
        
        # Print to stdout
        sys.stdout.write(decrypted_content)
        return True
    
    except Exception as e:
        print(f"✗ Error reading {filepath}: {e}", file=sys.stderr)
        return False


def view(filepath, password):
    """Decrypt and display file as formatted markdown."""
    # Normalize path
    if not filepath.startswith(CONTENT_DIR):
        filepath = os.path.join(CONTENT_DIR, filepath)
    
    # Ensure .enc extension
    if not filepath.endswith('.enc'):
        filepath = filepath + '.enc'
    
    if not os.path.exists(filepath):
        print(f"✗ File not found: {filepath}", file=sys.stderr)
        return False
    
    try:
        # Read encrypted content
        with open(filepath, 'r', encoding='utf-8') as f:
            encrypted_content = f.read()
        
        # Decrypt
        decrypted_content = decrypt_file_content(encrypted_content, password)
        if decrypted_content is None:
            print(f"✗ Failed to decrypt {filepath}", file=sys.stderr)
            return False
        
        # Convert to HTML
        html_content = markdown.markdown(
            decrypted_content,
            extensions=['extra', 'codehilite', 'toc']
        )
        
        # Print HTML
        print("<!DOCTYPE html>")
        print("<html>")
        print("<head>")
        print("<meta charset='utf-8'>")
        print("<meta name='viewport' content='width=device-width, initial-scale=1'>")
        print("<style>")
        print("body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; line-height: 1.6; max-width: 900px; margin: 0 auto; padding: 20px; }")
        print("code { background: #f4f4f4; padding: 2px 6px; border-radius: 3px; font-family: 'Courier New', monospace; }")
        print("pre { background: #f4f4f4; padding: 10px; border-radius: 5px; overflow-x: auto; }")
        print("table { border-collapse: collapse; width: 100%; }")
        print("th, td { border: 1px solid #ddd; padding: 8px; text-align: left; }")
        print("th { background-color: #f0f0f0; }")
        print("</style>")
        print("</head>")
        print("<body>")
        print(html_content)
        print("</body>")
        print("</html>")
        
        return True
    
    except Exception as e:
        print(f"✗ Error processing {filepath}: {e}", file=sys.stderr)
        return False


def unprotect(filepath):
    """Remove password protection from a file (file stays encrypted, becomes publicly viewable)."""
    # Normalize path
    if not filepath.startswith(CONTENT_DIR):
        filepath = os.path.join(CONTENT_DIR, filepath)
    
    # Normalize file_key to relative path (without CONTENT_DIR prefix)
    display_path = filepath.replace(CONTENT_DIR + '/', '')
    file_key = get_file_key(display_path)
    passwords = load_passwords()
    
    if file_key not in passwords:
        print(f"✗ File not protected: {file_key}", file=sys.stderr)
        return False
    
    # Only remove protection, keep encryption metadata
    if 'protected' in passwords[file_key]:
        del passwords[file_key]['protected']
    if 'password_hash' in passwords[file_key]:
        del passwords[file_key]['password_hash']
    if 'created_at' in passwords[file_key]:
        del passwords[file_key]['created_at']
    
    if save_passwords(passwords):
        print(f"✓ Unprotected: {display_path} (file stays encrypted, viewable in browser)")
        return True
    else:
        print(f"✗ Failed to unprotect: {file_key}", file=sys.stderr)
        return False


def list_files(directory=None):
    """List all protected and encrypted files."""
    passwords = load_passwords()
    
    if directory:
        if not directory.startswith(CONTENT_DIR):
            directory = os.path.join(CONTENT_DIR, directory)
        directory = os.path.normpath(directory)
    else:
        directory = CONTENT_DIR
    
    protected_files = []
    encrypted_files = []
    
    for filepath in Path(CONTENT_DIR).rglob('*'):
        if filepath.is_file():
            rel_path = filepath.relative_to(CONTENT_DIR)
            file_key = get_file_key(str(rel_path))
            
            if directory != CONTENT_DIR:
                dir_key = get_file_key(directory.replace(CONTENT_DIR + '/', ''))
                if not str(rel_path).startswith(dir_key.replace(CONTENT_DIR + '/', '')):
                    continue
            
            if file_key in passwords:
                data = passwords[file_key]
                if data.get('protected'):
                    protected_files.append((str(rel_path), data.get('created_at', 'Unknown')))
            
            if filepath.suffix == '.enc':
                encrypted_files.append(str(rel_path))
    
    if not protected_files and not encrypted_files:
        if directory == CONTENT_DIR:
            print("No protected or encrypted files found.")
        else:
            print(f"No protected or encrypted files found in {directory}")
        return True
    
    if protected_files:
        print("\nProtected Files:")
        print(f"{'File':<50} {'Created':<30}")
        print("-" * 80)
        for filepath, created_at in sorted(protected_files):
            try:
                dt = datetime.fromisoformat(created_at)
                created_str = dt.strftime('%Y-%m-%d %H:%M:%S')
            except:
                created_str = created_at
            print(f"{filepath:<50} {created_str:<30}")
    
    if encrypted_files:
        print("\nEncrypted Files:")
        for filepath in sorted(encrypted_files):
            print(f"  {filepath}")
    
    if protected_files:
        print(f"\nTotal: {len(protected_files)} protected file(s)", end="")
    if encrypted_files:
        if protected_files:
            print(f", {len(encrypted_files)} encrypted file(s)")
        else:
            print(f"Total: {len(encrypted_files)} encrypted file(s)")
    else:
        print()
    
    return True


def status(filepath):
    """Check if a file is protected or encrypted."""
    # Normalize path
    if not filepath.startswith(CONTENT_DIR):
        filepath = os.path.join(CONTENT_DIR, filepath)
    
    # Normalize file_key to relative path (without CONTENT_DIR prefix)
    display_path = filepath.replace(CONTENT_DIR + '/', '')
    file_key = get_file_key(display_path)
    passwords = load_passwords()
    
    is_encrypted = filepath.endswith('.enc')
    is_protected = file_key in passwords and passwords[file_key].get('protected', False)
    
    print(f"\nFile: {display_path}")
    print(f"Encrypted: {'Yes' if is_encrypted else 'No'}")
    print(f"Protected: {'Yes' if is_protected else 'No'}")
    
    if is_protected:
        data = passwords[file_key]
        created_at = data.get('created_at', 'Unknown')
        print(f"Created: {created_at}")
    
    return True


def main():
    if len(sys.argv) < 2:
        print(__doc__, file=sys.stderr)
        return 1
    
    command = sys.argv[1]
    
    if command == 'unprotect':
        if len(sys.argv) < 3:
            print("Usage: manage_content.py unprotect <filepath>", file=sys.stderr)
            return 1
        filepath = sys.argv[2]
        return 0 if unprotect(filepath) else 1
    
    elif command == 'encrypt':
        if len(sys.argv) < 4:
            print("Usage: manage_content.py encrypt <filepath> <password>", file=sys.stderr)
            return 1
        filepath = sys.argv[2]
        password = sys.argv[3]
        return 0 if encrypt(filepath, password) else 1
    
    elif command == 'decrypt':
        if len(sys.argv) < 4:
            print("Usage: manage_content.py decrypt <filepath> <password>", file=sys.stderr)
            return 1
        filepath = sys.argv[2]
        password = sys.argv[3]
        # Output goes to stdout, don't print status
        result = decrypt(filepath, password)
        return 0 if result else 1
    
    elif command == 'view':
        if len(sys.argv) < 4:
            print("Usage: manage_content.py view <filepath> <password>", file=sys.stderr)
            return 1
        filepath = sys.argv[2]
        password = sys.argv[3]
        return 0 if view(filepath, password) else 1
    
    elif command == 'list':
        directory = sys.argv[2] if len(sys.argv) > 2 else None
        return 0 if list_files(directory) else 1
    
    elif command == 'status':
        if len(sys.argv) < 3:
            print("Usage: manage_content.py status <filepath>", file=sys.stderr)
            return 1
        filepath = sys.argv[2]
        return 0 if status(filepath) else 1
    
    else:
        print(f"Unknown command: {command}", file=sys.stderr)
        print(__doc__, file=sys.stderr)
        return 1


if __name__ == '__main__':
    sys.exit(main())