from __future__ import annotations

import logging
from datetime import datetime
from pathlib import Path

from chia_rs import AugSchemeMPL, G1Element, PrivateKey
from chia_rs.sized_bytes import bytes32
from chiapos import DiskPlotter

from chia.daemon.keychain_proxy import KeychainProxy, connect_to_keychain_and_validate, wrap_local_keychain
from chia.plotting.util import Params, stream_plot_info_ph, stream_plot_info_pk
from chia.types.blockchain_format.proof_of_space import (
    calculate_plot_id_ph,
    calculate_plot_id_pk,
    generate_plot_public_key,
)
from chia.util.bech32m import decode_puzzle_hash
from chia.util.keychain import Keychain
from chia.wallet.derive_keys import master_sk_to_farmer_sk, master_sk_to_local_sk, master_sk_to_pool_sk

log = logging.getLogger(__name__)


class PlotKeys:
    def __init__(
        self,
        farmer_public_key: G1Element,
        pool_public_key: G1Element | None,
        pool_contract_address: str | None,
    ):
        self.farmer_public_key = farmer_public_key
        self.pool_public_key = pool_public_key
        self.pool_contract_address = pool_contract_address

    @property
    def pool_contract_puzzle_hash(self) -> bytes32 | None:
        if self.pool_contract_address is not None:
            return decode_puzzle_hash(self.pool_contract_address)
        return None


class PlotKeysResolver:
    def __init__(
        self,
        farmer_public_key: str | None,
        alt_fingerprint: int | None,
        pool_public_key: str | None,
        pool_contract_address: str | None,
        root_path: Path,
        log: logging.Logger,
        connect_to_daemon: bool = False,
    ) -> None:
        self.farmer_public_key = farmer_public_key
        self.alt_fingerprint = alt_fingerprint
        self.pool_public_key = pool_public_key
        self.pool_contract_address = pool_contract_address
        self.root_path = root_path
        self.log = log
        self.connect_to_daemon = connect_to_daemon
        self.resolved_keys: PlotKeys | None = None

    async def resolve(self) -> PlotKeys:
        if self.resolved_keys is not None:
            return self.resolved_keys

        keychain_proxy: KeychainProxy | None = None
        try:
            if self.connect_to_daemon:
                keychain_proxy = await connect_to_keychain_and_validate(self.root_path, self.log)
            else:
                keychain_proxy = wrap_local_keychain(Keychain(), log=self.log)

            farmer_public_key: G1Element
            if self.farmer_public_key is not None:
                farmer_public_key = G1Element.from_bytes(bytes.fromhex(self.farmer_public_key))
            else:
                farmer_public_key = await self.get_farmer_public_key(keychain_proxy)

            pool_public_key: G1Element | None = None
            if self.pool_public_key is not None:
                if self.pool_contract_address is not None:
                    raise RuntimeError("Choose one of pool_contract_address and pool_public_key")
                pool_public_key = G1Element.from_bytes(bytes.fromhex(self.pool_public_key))
            elif self.pool_contract_address is None:
                # If nothing is set, farms to the provided key (or the first key)
                pool_public_key = await self.get_pool_public_key(keychain_proxy)

            self.resolved_keys = PlotKeys(farmer_public_key, pool_public_key, self.pool_contract_address)
        finally:
            if keychain_proxy is not None:
                await keychain_proxy.close()
        return self.resolved_keys

    async def get_sk(self, keychain_proxy: KeychainProxy | None = None) -> PrivateKey | None:
        sk: PrivateKey | None = None
        if keychain_proxy:
            try:
                if self.alt_fingerprint is not None:
                    sk = await keychain_proxy.get_key_for_fingerprint(self.alt_fingerprint)
                else:
                    sk = await keychain_proxy.get_first_private_key()
            except Exception as e:
                log.error(f"Keychain proxy failed with error: {e}")
        else:
            sk_ent: tuple[PrivateKey, bytes] | None = None
            keychain: Keychain = Keychain()
            if self.alt_fingerprint is not None:
                sk_ent = keychain.get_private_key_by_fingerprint(self.alt_fingerprint)
            else:
                sk_ent = keychain.get_first_private_key()

            if sk_ent:
                sk = sk_ent[0]
        return sk

    async def get_farmer_public_key(self, keychain_proxy: KeychainProxy | None = None) -> G1Element:
        sk: PrivateKey | None = await self.get_sk(keychain_proxy)
        if sk is None:
            raise RuntimeError(
                "No keys, please run 'chia keys add', 'chia keys generate' or provide a public key with -f"
            )
        return master_sk_to_farmer_sk(sk).get_g1()

    async def get_pool_public_key(self, keychain_proxy: KeychainProxy | None = None) -> G1Element:
        sk: PrivateKey | None = await self.get_sk(keychain_proxy)
        if sk is None:
            raise RuntimeError(
                "No keys, please run 'chia keys add', 'chia keys generate' or provide a public key with -p"
            )
        return master_sk_to_pool_sk(sk).get_g1()


async def resolve_plot_keys(
    farmer_public_key: str | None,
    alt_fingerprint: int | None,
    pool_public_key: str | None,
    pool_contract_address: str | None,
    root_path: Path,
    log: logging.Logger,
    connect_to_daemon: bool = False,
) -> PlotKeys:
    return await PlotKeysResolver(
        farmer_public_key, alt_fingerprint, pool_public_key, pool_contract_address, root_path, log, connect_to_daemon
    ).resolve()


async def create_plots(
    args: Params,
    keys: PlotKeys,
    use_datetime: bool = True,
    test_private_keys: list[PrivateKey] | None = None,
) -> tuple[dict[bytes32, Path], dict[bytes32, Path]]:
    if args.tmp2_dir is None:
        args.tmp2_dir = args.tmp_dir
    assert (keys.pool_public_key is None) != (keys.pool_contract_puzzle_hash is None)
    num = args.num

    if keys.pool_public_key is not None:
        log.info(
            f"Creating {num} plots of size {args.size}, pool public key:  "
            f"{bytes(keys.pool_public_key).hex()} farmer public key: {bytes(keys.farmer_public_key).hex()}"
        )
    else:
        assert keys.pool_contract_puzzle_hash is not None
        log.info(
            f"Creating {num} plots of size {args.size}, pool contract address:  "
            f"{keys.pool_contract_address} farmer public key: {bytes(keys.farmer_public_key).hex()}"
        )

    tmp_dir_created = False
    if not args.tmp_dir.exists():
        args.tmp_dir.mkdir(parents=True, exist_ok=True)
        tmp_dir_created = True

    tmp2_dir_created = False
    if not args.tmp2_dir.exists():
        args.tmp2_dir.mkdir(parents=True, exist_ok=True)
        tmp2_dir_created = True

    args.final_dir.mkdir(parents=True, exist_ok=True)

    created_plots: dict[bytes32, Path] = {}
    existing_plots: dict[bytes32, Path] = {}
    for i in range(num):
        # Generate a random master secret key
        if test_private_keys is not None:
            assert len(test_private_keys) == num
            sk: PrivateKey = test_private_keys[i]
        else:
            sk = AugSchemeMPL.key_gen(bytes32.secret())

        # The plot public key is the combination of the harvester and farmer keys
        # New plots will also include a taproot of the keys, for extensibility
        include_taproot: bool = keys.pool_contract_puzzle_hash is not None
        plot_public_key = generate_plot_public_key(
            master_sk_to_local_sk(sk).get_g1(), keys.farmer_public_key, include_taproot
        )

        # The plot id is based on the harvester, farmer, and pool keys
        if keys.pool_public_key is not None:
            plot_id: bytes32 = calculate_plot_id_pk(keys.pool_public_key, plot_public_key)
            plot_memo: bytes = stream_plot_info_pk(keys.pool_public_key, keys.farmer_public_key, sk)
        else:
            assert keys.pool_contract_puzzle_hash is not None
            plot_id = calculate_plot_id_ph(keys.pool_contract_puzzle_hash, plot_public_key)
            plot_memo = stream_plot_info_ph(keys.pool_contract_puzzle_hash, keys.farmer_public_key, sk)

        if args.plotid is not None:
            log.info(f"Debug plot ID: {args.plotid}")
            # Check if args.memo is of type bytes and convert it to a string if so
            if isinstance(args.plotid, bytes):
                plot_str = args.plotid.hex()  # Convert bytes to hex string
            else:
                plot_str = args.plotid
            plot_id = bytes32.fromhex(plot_str)

        if args.memo is not None:
            log.info(f"Debug memo: {args.memo}")
            # Check if args.memo is of type bytes and convert it to a string if so
            if isinstance(args.memo, bytes):
                memo_str = args.memo.hex()  # Convert bytes to hex string
            else:
                memo_str = args.memo
            plot_memo = bytes.fromhex(memo_str)

        dt_string = datetime.now().strftime("%Y-%m-%d-%H-%M")

        if use_datetime:
            filename: str = f"plot-k{args.size}-{dt_string}-{plot_id}.plot"
        else:
            filename = f"plot-k{args.size}-{plot_id}.plot"
        full_path: Path = args.final_dir / filename

        if not full_path.exists():
            log.info(f"Starting plot {i + 1}/{num}")
            # Creates the plot. This will take a long time for larger plots.
            plotter: DiskPlotter = DiskPlotter()
            plotter.create_plot_disk(
                str(args.tmp_dir),
                str(args.tmp2_dir),
                str(args.final_dir),
                filename,
                args.size,
                plot_memo,
                plot_id,
                args.buffer,
                args.buckets,
                args.stripe_size,
                args.num_threads,
                args.nobitfield,
            )
            created_plots[plot_id] = full_path
        else:
            log.info(f"Plot {filename} already exists")
            existing_plots[plot_id] = full_path

    log.info("Summary:")

    if tmp_dir_created:
        try:
            args.tmp_dir.rmdir()
        except Exception:
            log.info(f"warning: did not remove primary temporary folder {args.tmp_dir}, it may not be empty.")

    if tmp2_dir_created:
        try:
            args.tmp2_dir.rmdir()
        except Exception:
            log.info(f"warning: did not remove secondary temporary folder {args.tmp2_dir}, it may not be empty.")

    log.info(f"Created a total of {len(created_plots)} new plots")
    for created_path in created_plots.values():
        log.info(created_path.name)

    return created_plots, existing_plots