from __future__ import annotations

import hashlib
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Callable


class HashcatError(RuntimeError):
    def __init__(
        self,
        message: str,
        *,
        exit_code: int | None = None,
    ) -> None:
        super().__init__(message)
        self.exit_code = exit_code


@dataclass(frozen=True)
class HashcatRunResult:
    exit_code: int
    restore_seen: bool


def restore_path(
    job_dir: Path,
    step: dict,
) -> Path:
    return (
        job_dir
        / "state"
        / f"step-{step['step_no']:03d}.restore"
    )


def completed_marker(
    job_dir: Path,
    step: dict,
) -> Path:
    return (
        job_dir
        / "state"
        / f"step-{step['step_no']:03d}.completed"
    )


def _session_name(step: dict) -> str:
    session = step.get("session_name")

    if not isinstance(session, str) or not session:
        raise HashcatError(
            f"Step {step.get('step_no')} has no session_name."
        )

    return session

def _resolve_wordlist(
    config,
    resource: str,
    step_no,
) -> Path:
    if not isinstance(resource, str) or not resource:
        raise HashcatError(
            f"Step {step_no} has invalid dictionary resource."
        )

    prefix = "wordlist:"

    if not resource.startswith(prefix):
        raise HashcatError(
            f"Step {step_no} has unsupported dictionary resource: "
            f"{resource}"
        )

    name = resource[len(prefix):]

    if not name or "/" in name or "\\" in name or name in (".", ".."):
        raise HashcatError(
            f"Step {step_no} has invalid dictionary resource: "
            f"{resource}"
        )

    dictionary = config.wordlists.get(name)

    if dictionary is None:
        raise HashcatError(
            f"Wordlist resource {resource!r} is not configured."
        )

    dictionary = dictionary.resolve()

    if not dictionary.is_file():
        raise HashcatError(
            f"Configured wordlist does not exist: {dictionary}"
        )

    return dictionary

def build_start_command(
    config,
    job_dir: Path,
    job: dict,
    step: dict,
) -> list[str]:
    exe = config.hashcat_exe.resolve()

    if not exe.is_file():
        raise HashcatError(
            f"Hashcat executable does not exist: {exe}"
        )

    hash_file = (
        job_dir
        / "hashes"
        / job["hash_file_name"]
    ).resolve()

    if not hash_file.is_file():
        raise HashcatError(
            f"Hash file does not exist: {hash_file}"
        )

    restore = restore_path(
        job_dir,
        step,
    ).resolve()

    session = _session_name(step)

    attack_mode = step.get("definition", {}).get(
        "attack_mode"
    )

    if not isinstance(attack_mode, int):
        raise HashcatError(
            f"Step {step.get('step_no')} has invalid "
            "attack_mode."
        )

    definition = step["definition"]

    cmd = [
        str(exe),
        "-m",
        str(job["hash_mode"]),
        "-a",
        str(attack_mode),
        "--session",
        session,
        "--restore-file-path",
        str(restore),
        "--potfile-path",
        str(config.potfile_path.resolve()),
        "--logfile-disable",
        "-w",
        str(
            definition.get(
                "workload_profile",
                1,
            )
        ),
    ]

    if attack_mode in (0, 1, 6, 7):
        dictionary_definition = (
            definition.get(
                "dictionaries"
            )
            if attack_mode == 1
            else definition.get(
                "dictionary"
            )
        )

        if attack_mode == 0:
            if not isinstance(
                dictionary_definition,
                dict,
            ):
                raise HashcatError(
                    f"Step {step.get('step_no')} "
                    "attack mode 0 has no dictionary."
                )

            resources = [
                dictionary_definition.get(
                    "resource"
                )
            ]

        elif attack_mode == 1:
            if (
                not isinstance(
                    dictionary_definition,
                    list,
                )
                or len(dictionary_definition) != 2
            ):
                raise HashcatError(
                    f"Step {step.get('step_no')} "
                    "attack mode 1 requires exactly "
                    "two dictionaries."
                )

            resources = [
                item.get("resource")
                if isinstance(item, dict)
                else None
                for item in dictionary_definition
            ]

        else:
            if not isinstance(
                dictionary_definition,
                dict,
            ):
                raise HashcatError(
                    f"Step {step.get('step_no')} "
                    f"attack mode {attack_mode} "
                    "has no dictionary."
                )

            resources = [
                dictionary_definition.get(
                    "resource"
                )
            ]

        dictionaries = [
            _resolve_wordlist(
                config,
                resource,
                step.get("step_no"),
            )
            for resource in resources
        ]

        induction = (
            job_dir
            / "state"
            / "induct"
        ).resolve()

        outfile_check = (
            job_dir
            / "state"
            / "outfiles"
        ).resolve()

        induction.mkdir(
            parents=True,
            exist_ok=True,
        )

        outfile_check.mkdir(
            parents=True,
            exist_ok=True,
        )

        cmd.extend(
            [
                "--induction-dir",
                str(induction),
                "--outfile-check-dir",
                str(outfile_check),
                str(hash_file),
            ]
        )

        cmd.extend(
            str(dictionary)
            for dictionary in dictionaries
        )

        if attack_mode in (6, 7):
            mask = definition.get("mask")

            if not isinstance(mask, str) or not mask:
                raise HashcatError(
                    f"Step {step.get('step_no')} "
                    f"attack mode {attack_mode} has no mask."
                )

            if attack_mode == 6:
                cmd.append(mask)
            else:
                cmd.insert(
                    len(cmd) - len(dictionaries),
                    mask,
                )

    elif attack_mode == 3:
        mask = definition.get("mask")

        if not isinstance(mask, str) or not mask:
            raise HashcatError(
                f"Step {step.get('step_no')} "
                "attack mode 3 has no mask."
            )

        cmd.extend(
            [
                str(hash_file),
                mask,
            ]
        )

    else:
        raise HashcatError(
            f"Unsupported attack mode {attack_mode}. "
            "Use a structured adapter for new modes."
        )

    return cmd


def build_restore_command(
    config,
    job_dir: Path,
    step: dict,
) -> list[str]:
    exe = config.hashcat_exe.resolve()

    if not exe.is_file():
        raise HashcatError(
            f"Hashcat executable does not exist: {exe}"
        )

    restore = restore_path(
        job_dir,
        step,
    ).resolve()

    return [
        str(exe),
        "--session",
        _session_name(step),
        "--restore-file-path",
        str(restore),
        "--restore",
    ]


def build_show_command(
    config,
    job_dir: Path,
    job: dict,
) -> list[str]:
    exe = config.hashcat_exe.resolve()

    if not exe.is_file():
        raise HashcatError(
            f"Hashcat executable does not exist: {exe}"
        )

    hash_file = (
        job_dir
        / "hashes"
        / job["hash_file_name"]
    ).resolve()

    if not hash_file.is_file():
        raise HashcatError(
            f"Hash file does not exist: {hash_file}"
        )

    return [
        str(exe),
        "-m",
        str(job["hash_mode"]),
        "--potfile-path",
        str(config.potfile_path.resolve()),
        "--show",
        str(hash_file),
    ]


def command_text(cmd: list[str]) -> str:
    return " ".join(
        f'"{value}"' if " " in value else value
        for value in cmd
    )


def run_hashcat(
    cmd: list[str],
    log_path: Path,
    *,
    on_output: Callable[[bytes], None] | None = None,
) -> int:
    log_path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    if not cmd:
        raise HashcatError(
            "Hashcat command is empty."
        )

    hashcat_dir = (
        Path(cmd[0])
        .resolve()
        .parent
    )

    text = command_text(cmd)

    with log_path.open(
        "ab",
    ) as log:
        log.write(
            b"\n=== START ===\n"
        )
        log.write(
            text.encode(
                "utf-8",
                errors="replace",
            )
        )
        log.write(b"\n")
        log.flush()

        try:
            process = subprocess.Popen(
                cmd,
                cwd=str(hashcat_dir),
                stdin=None,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                bufsize=0,
                creationflags=getattr(
                    subprocess,
                    "CREATE_NEW_PROCESS_GROUP",
                    0,
                ),
            )
        except OSError as exc:
            log.write(
                (
                    "\n=== PROCESS START ERROR ===\n"
                    f"{exc}\n"
                ).encode(
                    "utf-8",
                    errors="replace",
                )
            )
            raise HashcatError(
                f"Failed to start Hashcat: {exc}"
            ) from exc

        try:
            if process.stdout is None:
                raise HashcatError(
                    "Hashcat stdout pipe was not created."
                )

            while True:
                chunk = process.stdout.read(4096)

                if not chunk:
                    break

                log.write(chunk)
                log.flush()

                if on_output is not None:
                    on_output(chunk)

        except BaseException:
            try:
                process.kill()
            except OSError:
                pass

            raise

        finally:
            if process.stdout is not None:
                process.stdout.close()

        exit_code = process.wait()

        log.write(
            (
                "\n=== END ===\n"
                f"exit_code={exit_code}\n"
            ).encode(
                "utf-8",
                errors="replace",
            )
        )
        log.flush()

    return exit_code


def run_step(
    config,
    job_dir: Path,
    job: dict,
    step: dict,
    log_path: Path,
    *,
    restore: bool = False,
) -> HashcatRunResult:
    restore_file = restore_path(
        job_dir,
        step,
    )

    if restore:
        if not restore_file.is_file():
            raise HashcatError(
                f"Restore requested but restore file "
                f"does not exist: {restore_file}"
            )

        cmd = build_restore_command(
            config,
            job_dir,
            step,
        )
    else:
        cmd = build_start_command(
            config,
            job_dir,
            job,
            step,
        )

    exit_code = run_hashcat(
        cmd,
        log_path,
        on_output=lambda chunk: print(
            chunk.decode(
                "utf-8",
                errors="replace",
            ),
            end="",
            flush=True,
        ),
    )

    restore_seen = restore_file.is_file()

    if exit_code not in (0, 1):
        raise HashcatError(
            f"Hashcat exited with unsupported "
            f"code {exit_code}.",
            exit_code=exit_code,
        )

    return HashcatRunResult(
        exit_code=exit_code,
        restore_seen=restore_seen,
    )


def run_show(
    config,
    job_dir: Path,
    job: dict,
    output_path: Path,
    log_path: Path,
) -> str:
    cmd = build_show_command(
        config,
        job_dir,
        job,
    )

    output_path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    log_path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    hashcat_dir = (
        Path(cmd[0])
        .resolve()
        .parent
    )

    text = command_text(cmd)

    with log_path.open(
        "ab",
    ) as log:
        log.write(
            b"\n=== SHOW START ===\n"
        )
        log.write(
            text.encode(
                "utf-8",
                errors="replace",
            )
        )
        log.write(b"\n")
        log.flush()

        try:
            process = subprocess.run(
                cmd,
                cwd=str(hashcat_dir),
                stdin=None,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                check=False,
            )
        except OSError as exc:
            log.write(
                (
                    "\n=== SHOW START ERROR ===\n"
                    f"{exc}\n"
                ).encode(
                    "utf-8",
                    errors="replace",
                )
            )
            raise HashcatError(
                f"Failed to start Hashcat --show: {exc}"
            ) from exc

        output = process.stdout or b""

        output_path.write_bytes(
            output
        )

        log.write(output)
        log.write(
            (
                "\n=== SHOW END ===\n"
                f"exit_code={process.returncode}\n"
            ).encode(
                "utf-8",
                errors="replace",
            )
        )
        log.flush()

    if process.returncode not in (0, 1):
        raise HashcatError(
            f"Hashcat --show exited with "
            f"code {process.returncode}."
        )

    return output.decode(
        "utf-8",
        errors="replace",
    )


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()

    with path.open("rb") as handle:
        while True:
            chunk = handle.read(1024 * 1024)

            if not chunk:
                break

            digest.update(chunk)

    return digest.hexdigest()