import os
import time
from .utils import _find_first_nonempty
from . import ui
from .config import PASSPHRASE_DICT_PATHS, PASSPHRASE_GEN_DIR, PASSPHRASE_SEPARATORS, sep_to_name, RULE_PATHS
from .runners import get_cracked_result, run_combinator_attack, run_wordlist_attack, run_prince_attack
from .utils import _format_time


def load_passphrase_dict():
    """Load EFF diceware wordlist. Returns (title_cased_words, source_path) or (None, None)."""
    src = _find_first_nonempty(PASSPHRASE_DICT_PATHS)
    if not src:
        return None, None
    words = []
    with open(src) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            word = line.split('\t', 1)[-1].strip()
            if word and len(word) >= 2:
                words.append(word)
    words.sort(key=len)
    return words, src


def _words_path():
    return os.path.join(PASSPHRASE_GEN_DIR, "words.txt")

def _words_sep_path(sep):
    return os.path.join(PASSPHRASE_GEN_DIR, f"words_{sep_to_name(sep)}.txt")

def _left_2w_path(sep):
    return os.path.join(PASSPHRASE_GEN_DIR, f"left_2w_{sep_to_name(sep)}.txt")


def ensure_passphrase_files(need_3word=False):
    """
    Generate passphrase wordlist files from EFF dictionary if they don't exist.

    Generated files:
      words.txt              — title-cased words (~7k lines, ~80KB)
      words_<sep>.txt        — word+sep per non-empty separator (~7k lines each)
      left_2w_<sep>.txt      — word<sep>word<sep> pairs, patient only (~49M lines each)

    Returns (words_path, sep_files, left_2w_files) or None on failure.
    """
    os.makedirs(PASSPHRASE_GEN_DIR, exist_ok=True)
    non_empty_seps = [s for s in PASSPHRASE_SEPARATORS if s]

    words_path = _words_path()
    have_words = _find_first_nonempty([words_path]) is not None

    sep_files = {sep: _words_sep_path(sep) for sep in non_empty_seps
                 if _find_first_nonempty([_words_sep_path(sep)])}
    left_2w_files = ({sep: _left_2w_path(sep) for sep in non_empty_seps
                      if _find_first_nonempty([_left_2w_path(sep)])}
                     if need_3word else {})

    all_sep_done = all(sep in sep_files for sep in non_empty_seps)
    all_2w_done  = not need_3word or all(sep in left_2w_files for sep in non_empty_seps)

    if have_words and all_sep_done and all_2w_done:
        ui.ok("Passphrase files already exist")
        return words_path, sep_files, left_2w_files

    # Need to generate — load source dictionary
    words, src_path = load_passphrase_dict()
    if not words:
        ui.warn("EFF diceware wordlist not found — skipping passphrase attack")
        ui.hint("Install: sudo wget -q https://www.eff.org/files/2016/07/18/eff_large_wordlist.txt \\")
        ui.hint("         -O /usr/share/wordlists/eff-diceware.txt")
        return None

    n = len(words)
    ui.info(f"Loaded {n} words from {src_path}")

    if not have_words:
        ui.info(f"Generating words.txt ({n:,} lines)...")
        with open(words_path, 'w') as f:
            for w in words:
                f.write(w + "\n")
        ui.ok(f"  words.txt: {os.path.getsize(words_path) / 1024:.0f} KB")

    for sep in non_empty_seps:
        if sep in sep_files:
            continue
        p = _words_sep_path(sep)
        ui.info(f"Generating words_{sep_to_name(sep)}.txt ({n:,} lines)...")
        with open(p, 'w') as f:
            for w in words:
                f.write(w + sep + "\n")
        sep_files[sep] = p
        ui.ok(f"  words_{sep_to_name(sep)}.txt: {os.path.getsize(p) / 1024:.0f} KB")

    if need_3word:
        for sep in non_empty_seps:
            if sep in left_2w_files:
                continue
            p = _left_2w_path(sep)
            total = n * n
            ui.info(f"Generating left_2w_{sep_to_name(sep)}.txt ({total:,} lines, ~{total * 14 / (1024**3):.1f} GB)...")
            ui.info("  One-time generation — file is reused on future runs")
            t0 = time.time()
            with open(p, 'w', buffering=8 * 1024 * 1024) as f:
                for i, w1 in enumerate(words):
                    if i % 500 == 0 and i > 0:
                        elapsed = time.time() - t0
                        pct = 100 * i / n
                        eta = (elapsed / i) * (n - i)
                        ui.info(f"  {i}/{n} ({pct:.0f}%) — ETA {_format_time(eta)}")
                    f.writelines(w1 + sep + w2 + sep + "\n" for w2 in words)
            size_gb = os.path.getsize(p) / (1024**3)
            ui.ok(f"  left_2w_{sep_to_name(sep)}.txt: {size_gb:.1f} GB — {_format_time(time.time() - t0)}")
            left_2w_files[sep] = p

    return words_path, sep_files, left_2w_files


def run_passphrase_attack(hash_file, mode, gpu_flags, max_words=3, patience=2):
    """
    Passphrase combination attack using EFF diceware dictionary.

    Attack tiers:
      Tier 1: single words + best64 mutations              (~7k × 64 rules)
      Tier 2: 2-word per separator via combinator           (~49M combos each)
      Tier 3: 3-word no-sep via PRINCE (streaming)         (normal+)
              3-word with-sep via combinator               (patient only)
    """
    ui.section("Passphrase Combination Attack")

    already = get_cracked_result(hash_file, mode)
    if already:
        ui.info(f"Already cracked: {already}")
        return already

    result_paths = ensure_passphrase_files(need_3word=(patience >= 3))
    if result_paths is None:
        return None

    words_path, sep_files, left_2w_files = result_paths
    best64 = RULE_PATHS.get("best64", "")

    # Tier 1: single words + mutations
    ui.info("Tier 1: single words + mutations (best64)")
    r = run_wordlist_attack(hash_file, mode, words_path, gpu_flags, rule=best64, label="passphrase-words")
    if r:
        return r

    # Tier 2: 2-word combinator per separator
    if max_words >= 2:
        ui.info("Tier 2: 2-word combinations")
        r = run_combinator_attack(hash_file, mode, words_path, words_path, gpu_flags, "2-word nosep")
        if r:
            return r
        for sep, sep_path in sep_files.items():
            r = run_combinator_attack(hash_file, mode, sep_path, words_path, gpu_flags, f"2-word '{sep}'")
            if r:
                return r

    # Tier 3: 3-word
    if max_words >= 3:
        ui.info("Tier 3: 3-word combinations")
        r = run_prince_attack(hash_file, mode, words_path, gpu_flags, max_words=3, label="PRINCE 3-word nosep")
        if r:
            return r
        if patience >= 3:
            for sep, lpath in left_2w_files.items():
                ui.warn("~343M candidates — may take 60-90 min for fast hashes")
                r = run_combinator_attack(hash_file, mode, lpath, words_path, gpu_flags, f"3-word '{sep}'")
                if r:
                    return r

    return None
