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

"""
Phase 2: DeDRM Plugin Configuration
Updates dedrm.json with extracted keys
"""

import os
import sys
import json
import shutil
import time
import subprocess
from datetime import datetime

# Try to import psutil, but make it optional
try:
    import psutil # type: ignore
    HAS_PSUTIL = True
except ImportError:
    HAS_PSUTIL = False


def create_kindle_key_from_k4i(k4i_path, dsn=None, tokens=None):
    """Create a Kindle key entry exactly like the DeDRM plugin does"""
    from modules.utils import print_error
    
    try:
        with open(k4i_path, 'r') as f:
            k4i_data = json.load(f)
        
        kindle_key = {
            "DSN": k4i_data.get("DSN", dsn or ""),
            "kindle.account.clear_old_secrets": k4i_data.get("kindle.account.clear_old_secrets", []),
            "kindle.account.new_secrets": k4i_data.get("kindle.account.new_secrets", []), 
            "kindle.account.secrets": k4i_data.get("kindle.account.secrets", []),
            "kindle.account.tokens": k4i_data.get("kindle.account.tokens", tokens or "")
        }
        
        return kindle_key
        
    except Exception as e:
        print_error(f"Failed to process k4i file: {e}")
        return None


def create_dedrm_config(kindle_key, kindlekey_txt_path, existing_dedrm_json_path=None, reference_json_path=None):
    """Create or update the dedrm.json configuration"""
    from modules.utils import print_step, print_warn
    
    if existing_dedrm_json_path and os.path.exists(existing_dedrm_json_path):
        print_step("Loading existing dedrm.json to append new keys...")
        try:
            with open(existing_dedrm_json_path, 'r') as f:
                dedrm_config = json.load(f)
            
            if "kindlekeys" not in dedrm_config:
                dedrm_config["kindlekeys"] = {}
            
            if "kindlekey" in dedrm_config["kindlekeys"]:
                print_step("Merging new secrets into existing kindlekey entry...")
                existing_key = dedrm_config["kindlekeys"]["kindlekey"]
                
                for secret_key in ['kindle.account.secrets', 'kindle.account.new_secrets', 'kindle.account.clear_old_secrets']:
                    if secret_key in kindle_key:
                        if secret_key not in existing_key:
                            existing_key[secret_key] = []
                        for item in kindle_key[secret_key]:
                            if item not in existing_key[secret_key]:
                                existing_key[secret_key].append(item)
                
                if kindle_key.get('DSN') and kindle_key['DSN']:
                    existing_key['DSN'] = kindle_key['DSN']
                
                if kindle_key.get('kindle.account.tokens') and kindle_key['kindle.account.tokens']:
                    existing_key['kindle.account.tokens'] = kindle_key['kindle.account.tokens']
            else:
                print_step("Adding new kindlekey entry to existing config...")
                dedrm_config["kindlekeys"]["kindlekey"] = kindle_key
            
            dedrm_config["kindleextrakeyfile"] = kindlekey_txt_path
            
            return dedrm_config
            
        except Exception as e:
            print_warn(f"Failed to load existing dedrm.json: {e}")
            print_step("Creating new configuration instead...")
    
    # Create new configuration
    print_step("Creating new configuration structure...")
    dedrm_config = {
        "adeptkeys": {},
        "adobe_pdf_passphrases": [],
        "adobewineprefix": "",
        "androidkeys": {},
        "bandnkeys": {},
        "configured": True,
        "deobfuscate_fonts": True,
        "ereaderkeys": {},
        "kindleextrakeyfile": "",
        "kindlekeys": {},
        "kindlewineprefix": "",
        "lcp_passphrases": [],
        "pids": [],
        "remove_watermarks": False,
        "serials": []
    }
    
    dedrm_config["kindleextrakeyfile"] = kindlekey_txt_path
    dedrm_config["kindlekeys"]["kindlekey"] = kindle_key
    
    return dedrm_config


def obfuscate_sensitive(text):
    """
    Obfuscate sensitive strings by showing first 2 and last 2 characters
    """
    if len(text) <= 4:
        return text
    return text[:2] + '*' * (len(text) - 4) + text[-2:]


def is_calibre_running():
    """Check if Calibre is currently running - Import from utils for consistency"""
    from modules.utils import is_calibre_running as utils_is_calibre_running # type: ignore
    return utils_is_calibre_running()


def backup_file_with_retry(source_path, backup_path, max_retries=3):
    """
    Backup a file with retry logic and file lock detection
    
    Args:
        source_path: Path to source file
        backup_path: Path to backup destination
        max_retries: Maximum number of retry attempts
        
    Returns:
        tuple: (success: bool, error_msg: str or None)
    """
    from modules.utils import print_warn
    
    for attempt in range(max_retries):
        try:
            # Try to read and write instead of copy2 to avoid some lock issues
            with open(source_path, 'rb') as fsrc:
                file_content = fsrc.read()
            
            with open(backup_path, 'wb') as fdst:
                fdst.write(file_content)
            
            return True, None
            
        except PermissionError as e:
            if attempt < max_retries - 1:
                print_warn(f"File locked, retrying in 2 seconds... (attempt {attempt + 1}/{max_retries})")
                time.sleep(2)
            else:
                return False, f"Permission denied after {max_retries} attempts: {e}"
                
        except Exception as e:
            return False, f"Failed to backup file: {e}"
    
    return False, "Unknown error during backup"


def run_phase_2(output_k4i, output_key, working_dir, script_dir, config=None):
    """
    Run Phase 2: DeDRM Plugin Configuration
    
    Args:
        output_k4i: Path to kindlekey.k4i
        output_key: Path to kindlekey.txt
        working_dir: Working directory for backups
        script_dir: Script directory
        config: Configuration dict (optional, for hide_sensitive_info setting)
    
    Returns:
        dict with keys:
            - success: bool
            - dedrm_json_path: str
    """
    from modules.utils import print_step, print_ok, print_error, print_warn, print_colored
    
    user_home = os.path.expanduser("~")
    dedrm_json = os.path.join(user_home, "AppData", "Roaming", "calibre", "plugins", "dedrm.json")
    reference_json = os.path.join(script_dir, "dedrm filled.json")
    
    # Get hide_sensitive_info setting - default to False (show all info)
    hide_sensitive = config.get('hide_sensitive_info', False) if config else False
    
    # Check if Calibre is running using reusable function from utils
    from modules.utils import wait_for_calibre_to_close # type: ignore
    
    print()
    print_step("Checking if Calibre is running...")
    # Capture the return value so key_finder.py can reset calibre_skip_phases_34
    # if the user waited here and Calibre actually closed.
    calibre_closed_in_phase2 = wait_for_calibre_to_close(
        "Calibre may have locked the dedrm.json file needed for backup.",
        calibre_import_enabled=True
    )
    
    # Create backups directory with timestamp
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backups_dir = os.path.join(working_dir, "backups")
    os.makedirs(backups_dir, exist_ok=True)
    
    # Create backup if dedrm.json exists
    if os.path.exists(dedrm_json):
        print()
        print_step("Creating backup of existing dedrm.json...")
        backup_json = os.path.join(backups_dir, f"dedrm_backup_{timestamp}.json")
        
        success, error_msg = backup_file_with_retry(dedrm_json, backup_json)
        
        if success:
            print_ok(f"Backup created: {backup_json}")
        else:
            print()
            # The backup reads dedrm.json first. A PermissionError here confirms
            # we cannot read OR write dedrm.json - backup is not the real blocker,
            # access to the file itself is.
            print_error(f"Cannot access dedrm.json: {error_msg}")
            print()
            print_colored("=" * 70, "red")
            print_colored(" CANNOT MODIFY DEDRM.JSON - ACCESS DENIED", "red")
            print_colored("=" * 70, "red")
            print()
            print("The script attempted to read dedrm.json 3 times and was denied")
            print("each time. Since the file cannot be read, it also cannot be")
            print("updated with your new Kindle keys. Phase 2 cannot continue.")
            print()

            # Check admin rights and tailor the advice
            from modules.utils import is_running_as_admin  # type: ignore
            if not is_running_as_admin():
                print_colored("=" * 70, "yellow")
                print_colored("[!] ROOT CAUSE: Script is NOT running as Administrator!", "yellow")
                print_colored("    Windows is blocking access to the Calibre plugin folder.", "yellow")
                print_colored("    This is the #1 cause of [Errno 13] Permission denied.", "yellow")
                print_colored("=" * 70, "yellow")
                print()
                print_step("Fix - close this window and re-run with admin rights:")
                print("  Option 1: Double-click  Run_keyfinder_admin.vbs")
                print("            (requests UAC elevation automatically)")
                print()
                print("  Option 2: Right-click   Run_keyfinder.bat")
                print("            Choose         'Run as administrator'")
                print()
            else:
                # Already elevated - file is locked by a process, not a privilege issue
                print_step("Script is running as Administrator - possible causes:")
                print("  - Calibre (or a Calibre background process) is still open")
                print("    and has locked dedrm.json.")
                print()
                print_step("Troubleshooting steps:")
                print("  1. Open Task Manager (Ctrl+Shift+Esc)")
                print("  2. Look for: calibre.exe, calibre-parallel.exe, ebook-convert.exe")
                print("  3. End ALL Calibre processes")
                print("  4. Run this script again")
                print()
                print_step("If the problem persists:")
                print("  - Restart your computer, then run the script again")
                print("  - Manually check permissions on:")
                print(f"    {dedrm_json}")
                print()
            
            return {
                'success': False,
                'dedrm_json_path': None
            }
    else:
        print_step("No existing dedrm.json found - will create new one.")
    
    # Process the k4i file
    print()
    print_step("Processing kindlekey.k4i...")
    kindle_key = create_kindle_key_from_k4i(output_k4i)
    
    if not kindle_key:
        print_error("Failed to process k4i file!")
        return {
            'success': False,
            'dedrm_json_path': None
        }
        
    print_ok("Kindle key data processed successfully.")
    
    # Create dedrm configuration (loads and merges in memory)
    print_step("Creating DeDRM configuration...")
    dedrm_config = create_dedrm_config(kindle_key, output_key, dedrm_json, reference_json)
    
    # NEW: Delete original file to release any file locks
    from modules.utils import remove_file_with_retry # type: ignore
    
    if os.path.exists(dedrm_json):
        print()
        print_step("Removing original file to prevent lock issues...")
        
        success, error_msg = remove_file_with_retry(dedrm_json, max_retries=3)
        
        if success:
            print_ok("Original file removed successfully")
        else:
            print_error(f"Failed to remove original file: {error_msg}")
            print_error("This may cause write failures - continuing anyway...")
        print()
    
    # Write dedrm.json (now as a fresh file with no locks)
    print_step("Writing dedrm.json...")
    os.makedirs(os.path.dirname(dedrm_json), exist_ok=True)
    
    with open(dedrm_json, 'w') as f:
        json.dump(dedrm_config, f, indent=2)
        f.flush()  # Force write to buffer
        os.fsync(f.fileno())  # Force write to disk
    
    # Small delay to let Windows/antivirus release the new file
    time.sleep(0.5)
    
    print_ok("DeDRM configuration updated successfully!")
    print(f"Updated key: kindlekey")
    print(f"Set extra key file: {output_key}")
    print()
    
    # Verify
    print_step("Verifying configuration...")
    
    with open(dedrm_json, 'r') as f:
        dedrm_verify = json.load(f)
    
    key_count = len(dedrm_verify.get("kindlekeys", {}))
    
    # Read voucher keys from kindlekey.txt
    voucher_keys = []
    if os.path.exists(output_key):
        try:
            with open(output_key, 'r') as f:
                voucher_keys = [line.strip() for line in f if line.strip()]
        except Exception:
            pass
    
    # Display keys summary table (this is what Step 3 "No show all Info" controls)
    print()
    print("=" * 70)
    print("KEYS SUMMARY TABLE")
    print("=" * 70)
    print()
    
    # Table header (extra wide to fit 88-char base64 New Secret)
    print("┌─────────────────────────────┬────────────────────────────────────────────────────────────────────────────────────────────────┐")
    print("│ Configuration Item          │ Value                                                                                          │")
    print("├─────────────────────────────┼────────────────────────────────────────────────────────────────────────────────────────────────┤")
    
    # Total keys
    print(f"│ Total Kindle Keys           │ {str(key_count):<94} │")
    
    # Extra key file path
    extra_key_path = dedrm_verify.get('kindleextrakeyfile', 'Not set')
    if len(extra_key_path) > 94:
        extra_key_path = extra_key_path[:91] + "..."
    print(f"│ Extra Key File Path         │ {extra_key_path:<94} │")
    
    # Show key details with obfuscation if enabled
    if "kindlekeys" in dedrm_verify and "kindlekey" in dedrm_verify["kindlekeys"]:
        key_data = dedrm_verify["kindlekeys"]["kindlekey"]
        
        # DSN (full width, no truncation unless obfuscated)
        dsn_value = key_data.get('DSN', 'Not found')
        if hide_sensitive and dsn_value != 'Not found':
            dsn_value = obfuscate_sensitive(dsn_value)
        print(f"│ DSN                         │ {dsn_value:<94} │")
        
        # DSN_clear (if available)
        dsn_clear = key_data.get('DSN_clear', '')
        if dsn_clear:
            if hide_sensitive:
                dsn_clear = obfuscate_sensitive(dsn_clear)
            print(f"│ DSN Clear                   │ {dsn_clear:<94} │")
        
        # Tokens (full width, no truncation unless obfuscated)
        tokens_value = key_data.get('kindle.account.tokens', '')
        if tokens_value:
            if hide_sensitive:
                tokens_value = obfuscate_sensitive(tokens_value)
            print(f"│ Tokens                      │ {tokens_value:<94} │")
        
        # Show sample of secrets if available (obfuscated if needed)
        clear_old_secrets = key_data.get('kindle.account.clear_old_secrets', [])
        new_secrets = key_data.get('kindle.account.new_secrets', [])
        
        if clear_old_secrets:
            sample_clear_old = clear_old_secrets[0] if isinstance(clear_old_secrets, list) else str(clear_old_secrets)
            if hide_sensitive:
                sample_clear_old = obfuscate_sensitive(sample_clear_old)
            print(f"│ Old Secret                  │ {sample_clear_old:<94} │")
        
        if new_secrets:
            sample_new_secret = new_secrets[0] if isinstance(new_secrets, list) else str(new_secrets)
            if hide_sensitive:
                sample_new_secret = obfuscate_sensitive(sample_new_secret)
            print(f"│ New Secret                  │ {sample_new_secret:<94} │")
    
    # Voucher keys count at bottom
    print(f"│ Voucher Keys Count          │ {str(len(voucher_keys)):<94} │")
    
    # Table footer
    print("└─────────────────────────────┴────────────────────────────────────────────────────────────────────────────────────────────────┘")
    print()
    
    # Show privacy status
    if hide_sensitive:
        print("Note: Sensitive information (DSN, Tokens, Secrets) is hidden")
    else:
        print("Note: All information is fully visible")
    
    print()
    print_ok("Configuration verified successfully!")
    print()
    
    return {
        'success': True,
        'dedrm_json_path': dedrm_json,
        # True if the user waited here and Calibre is now confirmed closed.
        # key_finder.py uses this to reset calibre_skip_phases_34 so Phase 3 runs.
        'calibre_closed_in_phase2': calibre_closed_in_phase2
    }
