import os
import subprocess
import logging
from typing import Optional


def run_plink2(command: str):
    """
    Run a shell command invoking PLINK2, capturing stdout/stderr for logging.
    """
    result = subprocess.run(
        command,
        shell=True,
        capture_output=True,
        text=True,
        check=False
    )
    if result.returncode != 0:
        logging.debug("Error: %s", result.stderr)
    else:
        logging.debug("Output: %s", result.stdout)


def generate_bed_bim_fam(
    plink2_path: str,
    ped_file: str,
    map_file: str,
    output_prefix: str,
    relax_mind_threshold: bool = False,
    maf_threshold: Optional[float] = None,
    sample_keep_path: Optional[str] = None,
    autosomes_only: bool = False,
):
    """
    Generates BED/BIM/FAM from PED/MAP using PLINK2, *exactly like Hao's R scripts*:

      plink2 --pedmap <prefix> --make-bed --geno 0.1 --mind 0.1 --out <prefix>

    Notes:
      - We *ignore* ped_file/map_file here; PLINK reads <prefix>.ped/.map via --pedmap.
      - Ensure <output_prefix>.ped and <output_prefix>.map already exist before calling.
      - By default: no --maf, no --keep, no --chr filter (Hao's behavior).
      - Optional --maf, --keep, and autosome filter.
    
    Args:
        plink2_path: Path to PLINK2.
        ped_file: PED path.
        map_file: MAP path.
        output_prefix: Output prefix (no extension).
        relax_mind_threshold: If True, skip per-sample missingness filter.
        maf_threshold: Minor allele frequency filter for SNPs; if None, skip --maf.
        sample_keep_path: Path to file with samples to keep; if None, skip --keep.
        autosomes_only: If True, restrict to autosomal chromosomes only.
    """
    # sanity: the files PLINK expects with --pedmap <prefix>
    ped_expect = f"{output_prefix}.ped"
    map_expect = f"{output_prefix}.map"

    if not os.path.exists(ped_expect):
        raise FileNotFoundError(f"Missing PED for --pedmap: {ped_expect}")
    if not os.path.exists(map_expect):
        raise FileNotFoundError(f"Missing MAP for --pedmap: {map_expect}")

    mind = "" if relax_mind_threshold else "--mind 0.1"
    maf  = f"--maf {maf_threshold}" if maf_threshold is not None else ""   # usually None to match Hao
    keep = f"--keep {sample_keep_path}" if sample_keep_path else ""        # usually none to match Hao
    chrflag = "--chr 1-19" if autosomes_only else ""                       # usually off to match Hao

    logging.info(
        "[plink_utils:generate_bed_bim_fam] --pedmap %s -> BED with %s %s %s %s",
        output_prefix, mind or "no --mind", maf or "no --maf",
        keep or "no --keep", chrflag or "all chr"
    )

    # Exact PLINK command matching Hao’s workflow
    run_plink2(
        f"{plink2_path} --pedmap {output_prefix} "
        f"--make-bed --geno 0.1 {mind} {maf} {keep} {chrflag} --out {output_prefix}"
    )


def calculate_kinship_matrix(
    plink2_path: str,
    input_prefix: str,
    output_prefix: str,
    sample_keep_path: Optional[str] = None
):
    """
    Create PLINK .kin files from BED/BIM/FAM files.
    """
    keep = f"--keep {sample_keep_path}" if sample_keep_path else ""
    cmd = f"{plink2_path} --bfile {input_prefix} {keep} --make-rel square --out {output_prefix}"
    run_plink2(cmd)


def calculate_kinship_from_pedmap(
    plink2_path: str,
    pedmap_prefix: str,
    kin_prefix: str,
):
    """
    compute kinship directly from --pedmap <prefix>.
    Assumes <pedmap_prefix>.ped and <pedmap_prefix>.map exist.
    """
    run_plink2(
        f"{plink2_path} --pedmap {pedmap_prefix} --make-rel square --out {kin_prefix}"
    )
