﻿from passwork_client import PassworkClient
import os
import json
import argparse
from datetime import datetime

# Configuration
ACCESS_TOKEN = ""
REFRESH_TOKEN = ""  # Optional (required for token rotation)
MASTER_KEY = ""  # Master key (if client encryption is enabled)

HOST = ""  # Passwork API host URL


class MainArgumentParser(argparse.ArgumentParser):
    def parse_known_args(self, args=None, namespace=None):
        return super().parse_known_args(args, namespace)


def parse_args():
    parser = MainArgumentParser(
        description="Import data to Passwork from Bitwarden JSON export"
    )
    parser.add_argument("--host", type=str, default=HOST, help="Passwork API host URL")
    parser.add_argument("--token", type=str, default=ACCESS_TOKEN, help="Passwork access token")
    parser.add_argument("--refresh-token", type=str, default=REFRESH_TOKEN, help="Passwork refresh token")
    parser.add_argument("--master-key", type=str, default=MASTER_KEY, help="Passwork master key for decryption")

    parser.add_argument("file", type=str, nargs="?", default=None, help="Path to Bitwarden JSON export file")
    parser.add_argument(
        "--collections",
        type=str,
        default=None,
        help="Comma-separated list of collection IDs or names to import",
    )
    parser.add_argument("--vault-id", type=str, default=None, help="Passwork vault ID to import data to")

    args, _remaining = parser.parse_known_args()

    return args


def get_passwork_client(
    host: str | None = None,
    access_token: str | None = None,
    refresh_token: str | None = None,
    master_key: str | None = None,
):
    try:
        passwork = PassworkClient(host)
        passwork.set_tokens(access_token, refresh_token)
        if bool(master_key):
            passwork.set_master_key(master_key)

        return passwork

    except Exception as e:
        print(f"Error: {e}")
        return None


def get_file_path_data(file_path: str | None = None):
    if not file_path:
        file_path = input("Input file path: ").strip()

    if not file_path:
        print("File path cannot be empty.")
        return None

    if not os.path.isfile(file_path):
        print("File does not exist.")
        return None

    try:
        with open(file_path, "r", encoding="utf-8") as file:
            data = json.load(file)

            if not data["items"]:
                print("Invalid JSON file format.")
                return None

    except json.JSONDecodeError as e:
        print(f"Error loading data from file: {e}")
        return None

    return data


def get_collections(collections: str | None = None):
    if collections is None:
        collections = input("Input comma-separated list of collection IDs or names to export (optional): ").strip()

    if not collections:
        return []

    return [item.strip() for item in collections.split(",") if item.strip()]


def parse_collections(data, collections):
    if not data.get("collections"):
        return []

    if len(collections) == 0:
        return data["collections"]

    result = []
    for collection in data["collections"]:
        if collection["id"] in collections or collection["name"] in collections:
            result.append(collection)

    return result


def get_vault_id(passwork: PassworkClient, vault_id: str | None = None):
    if vault_id is None:
        vault_id = input("Input Passwork vault ID to import data to (optional): ").strip()

    if not vault_id:
        return None

    try:
        vault = passwork.get_vault(vault_id)
    except Exception as e:
        print(f"Vault not found. Error: {e}")
        return None

    return vault


def get_company_vault_type(passwork: PassworkClient, code: str = "privateShared"):
    return passwork.find_vault_type(code=code)


def create_vaults(passwork: PassworkClient, collections: list[dict]):
    vaults = {}
    company_vault_type = get_company_vault_type(passwork, "company")
    for collection in collections:
        vaults[collection["id"]] = passwork_create_vault(passwork, collection["name"], company_vault_type["id"])

    return vaults


def passwork_create_vault(passwork: PassworkClient, name: str, vault_type_id: str):
    vault_id = passwork.create_vault(name, vault_type_id)
    vault = passwork.get_vault(vault_id)
    log_message(f"Created vault: {vault['name']} ({vault_id})")
    return vault


def create_folders(passwork: PassworkClient, vault: dict, collections: list[dict]):
    folders = {}
    for collection in collections:
        folders[collection["id"]] = passwork_create_folder(passwork, collection["name"], vault["id"])

    return folders


def passwork_create_folder(passwork: PassworkClient, name: str, vault_id: str):
    data = {
        "name": name,
        "vaultId": vault_id,
    }

    try:
        folder_id = passwork.call("POST", "/api/v1/folders", data)["id"]

        return passwork.call("GET", f"/api/v1/folders/{folder_id}")
    except Exception as e:
        print(f"Error creating folder: {e}")
        return None


def get_directories(password_collection_ids: list[str], collections: dict[str, dict]):
    directories = []
    for collection_id in password_collection_ids:
        if collection_id in collections:
            directories.append(collections[collection_id])

    return directories


def get_directories_names(directories):
    if not directories:
        return ""
    return ", ".join(directory.get("name", "") for directory in directories if directory.get("name"))


def prepare_password_fields(data, directories):
    vaults_names = get_directories_names(directories)

    if data.get("type") not in (1, 2):
        log_message(
            f"Object type {data.get('type')}, {data.get('name')} "
            f"from collections {vaults_names} was not imported"
        )
        return None

    fields = {
        "password": "",
        "name": data.get("name"),
        "description": data.get("notes"),
        "customs": [],
    }

    if len(directories) > 1:
        fields["description"] = (fields["description"] + "\n") if fields["description"] else ""
        fields["description"] += f"Copy of the password is in: {vaults_names}"

    login = data.get("login")
    if login:
        if login.get("username"):
            fields["login"] = login["username"]

        if login.get("password"):
            fields["password"] = login["password"]

        if login.get("totp"):
            fields["customs"].append(
                {
                    "name": "TOTP",
                    "value": login["totp"],
                    "type": "totp",
                }
            )

        if login.get("uris"):
            uris = [uri.get("uri") for uri in login["uris"] if isinstance(uri, dict) and uri.get("uri")]
            if uris:
                fields["urls"] = uris

    if data.get("fields"):
        for field in data["fields"]:
            field_type = field.get("type")
            if field_type in (0, 2):
                fields["customs"].append(
                    {
                        "name": str(field.get("name", "")),
                        "value": str(field.get("value", "")),
                        "type": "text",
                    }
                )
            elif field_type == 1:
                fields["customs"].append(
                    {
                        "name": str(field.get("name", "")),
                        "value": str(field.get("value", "")),
                        "type": "password",
                    }
                )
            else:
                log_message(
                    f"Field type \"link\" for object \"{data.get('name')}\" "
                    f"for vault(s) \"{vaults_names}\" was not imported"
                )

    return fields


LOG_FILE_NAME = f"import-{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.log"


def log_message(message: str) -> None:
    msg = f"{datetime.now().isoformat()} {message}"
    with open(LOG_FILE_NAME, "a", encoding="utf-8") as log_file:
        log_file.write(msg + "\n")
    print(msg)


def import_passwords(passwork: PassworkClient, collections: list[dict], data: dict, vault: dict | None):
    # Collections as vaults and folders
    vaults = []
    folders = []

    if len(collections) == 0:
        if vault is None:
            company_vault_type = get_company_vault_type(passwork)
            vault = passwork_create_vault(passwork, "Private vault", company_vault_type["id"])

        folders = {}
        if data["folders"]:
            for folder in data["folders"]:
                folders[folder["id"]] = passwork_create_folder(passwork, folder["name"], vault["id"])

        for password_data in data["items"]:
            log_message(f"Started import: {password_data['name']}")

            fields = prepare_password_fields(password_data, [vault])
            if fields is None:
                continue

            fields["vaultId"] = vault["id"]
            if "folderId" in password_data and password_data["folderId"]:
                fields["folderId"] = folders[password_data["folderId"]]["id"]

            passwork.create_item(fields)
            log_message(f"Completed import: {password_data['name']}")

        log_message("Import completed")
        return

    if vault is None:
        vaults = create_vaults(passwork, collections)
    else:
        folders = create_folders(passwork, vault, collections)

    if len(vaults) > 0:
        for password_data in data["items"]:
            vault_list = get_directories(password_data["collectionIds"], vaults)
            if len(vault_list) == 0:
                continue

            log_message(f"Started import: {password_data['name']}")

            fields = prepare_password_fields(password_data, vault_list)
            if fields is None:
                continue

            for vault in vault_list:
                fields["vaultId"] = vault["id"]
                passwork.create_item(fields)

            log_message(f"Completed import: {password_data['name']}")

        log_message("Import completed")
        return

    if len(folders) > 0:
        for password_data in data["items"]:
            folder_list = get_directories(password_data["collectionIds"], folders)
            if len(folder_list) == 0:
                continue

            log_message(f"Started import: {password_data['name']}")

            fields = prepare_password_fields(password_data, folder_list)
            if fields is None:
                continue

            fields["vaultId"] = vault["id"]
            for folder in folder_list:
                fields["folderId"] = folder["id"]
                passwork.create_item(fields)

            log_message(f"Completed import: {password_data['name']}")

        log_message("Import completed")
        return


def main() -> None:
    args = parse_args()

    passwork = get_passwork_client(args.host, args.token, args.refresh_token, args.master_key)
    if passwork is None:
        return

    data = get_file_path_data(args.file)
    if data is None:
        return

    collections = get_collections(args.collections)
    collections = parse_collections(data, collections)

    vault = get_vault_id(passwork, args.vault_id)

    confirm_message = "The following collections will be exported: "
    if len(collections) == 0:
        confirm_message += "Private vault"
    else:
        for collection in collections:
            confirm_message += f"\n- {collection['name']} ({collection['id']})"

    if vault is not None:
        confirm_message += f"\nExport will be performed to vault: \"{vault['name']}\""

    confirm_message += "\nContinue operation? (Y/N): "

    while True:
        answer = input(confirm_message).strip().upper()
        if answer in {"Y", "N"}:
            break
        print("Invalid input. Enter Y or N.")

    if answer == "Y":
        import_passwords(passwork, collections, data, vault)
    else:
        print("Operation cancelled.")


if __name__ == "__main__":
    main()
