import argparse
import os
import subprocess
import sys
import uuid
from pathlib import Path


PAUSE_FLAG = Path("/workspace/outputs/train_paused.flag")
WANDB_RUN_ID_FILE = Path("/workspace/outputs/wandb_run_id.txt")
DATASET_AUDIT_SCRIPT = Path("/workspace/scripts/pipeline/audit_dataset_integrity.py")


def _clear_pause_flag() -> None:
    try:
        if PAUSE_FLAG.exists():
            PAUSE_FLAG.unlink()
            print(f"[INFO] Cleared pause flag: {PAUSE_FLAG}")
    except Exception as exc:
        print(f"[WARNING] Failed to clear pause flag {PAUSE_FLAG}: {exc}")


def _wandb_credentials_present() -> bool:
    if os.environ.get("WANDB_API_KEY", "").strip():
        return True
    try:
        import wandb

        return bool(getattr(getattr(wandb, "api", None), "api_key", None))
    except Exception:
        return False


def _ensure_wandb_run_id() -> str | None:
    explicit = os.environ.get("WANDB_RUN_ID", "").strip()
    if explicit:
        return explicit

    if not _wandb_credentials_present():
        return None

    try:
        if WANDB_RUN_ID_FILE.exists():
            persisted = WANDB_RUN_ID_FILE.read_text(encoding="utf-8").strip()
            if persisted:
                os.environ["WANDB_RUN_ID"] = persisted
                return persisted
        run_id = uuid.uuid4().hex[:8]
        WANDB_RUN_ID_FILE.parent.mkdir(parents=True, exist_ok=True)
        WANDB_RUN_ID_FILE.write_text(run_id + "\n", encoding="utf-8")
        os.environ["WANDB_RUN_ID"] = run_id
        print(f"[INFO] Using WANDB_RUN_ID={run_id} (persisted to {WANDB_RUN_ID_FILE})")
        return run_id
    except Exception as exc:
        print(f"[WARNING] Failed to persist WANDB_RUN_ID: {exc}")
        return None


def main() -> None:
    parser = argparse.ArgumentParser(description="Launch training with WandB monitoring")
    parser.add_argument("--project_name", required=True, help="Project name (e.g. cortazar_v2)")
    parser.add_argument(
        "--exp_name",
        default=os.environ.get("EXP_NAME", "F5TTS_v1_Base"),
        help="Experiment name (F5TTS_v1_Base, F5TTS_Base, E2TTS_Base)",
    )
    parser.add_argument("--epochs", type=int, default=2500, help="Training epochs")
    parser.add_argument("--batch_size", type=int, default=1600, help="Batch size per GPU")
    parser.add_argument(
        "--base_model_path", default="F5-TTS/ckpts/base_models/model_last.pt", help="Base model to finetune from"
    )
    parser.add_argument(
        "--tokenizer",
        default="custom",
        choices=["pinyin", "char", "custom"],
        help="Tokenizer type",
    )
    parser.add_argument("--tokenizer_path", default=None, help="Path to tokenizer vocab (for custom)")
    args = parser.parse_args()

    print(f"--- Starting Core Training Pipeline for {args.project_name} ---")
    _clear_pause_flag()
    _ensure_wandb_run_id()

    # Fail-fast: validate the prepared dataset before burning GPU hours.
    skip_audit = os.environ.get("SKIP_DATA_AUDIT", "").strip().lower() in ("1", "true", "yes", "y", "on")
    if not skip_audit and DATASET_AUDIT_SCRIPT.exists():
        dataset_root = f"/workspace/F5-TTS/data/{args.project_name}_{args.tokenizer}"
        speaker_report = os.environ.get(
            "SPEAKER_REPORT_PATH",
            "/workspace/dataset_prepared/speaker_prune_report.jsonl",
        )
        speaker_threshold = os.environ.get("SPEAKER_THRESHOLD", "0.35")
        audit_cmd = [
            "python",
            str(DATASET_AUDIT_SCRIPT),
            "--dataset_root",
            dataset_root,
            "--speaker_report",
            speaker_report,
            "--require_speaker_report",
            "--speaker_threshold",
            speaker_threshold,
            "--forbid_audio_path_contains",
            "SPEAKER_01",
        ]
        print(f"[INFO] Running dataset audit: {' '.join(audit_cmd)}")
        try:
            subprocess.check_call(audit_cmd)
        except subprocess.CalledProcessError as exc:
            try:
                PAUSE_FLAG.parent.mkdir(parents=True, exist_ok=True)
                PAUSE_FLAG.write_text(
                    "paused_by=train.py\nreason=dataset_audit_failed\n"
                    f"dataset_root={dataset_root}\nexit_code={exc.returncode}\n",
                    encoding="utf-8",
                )
                print(f"[ERROR] Dataset audit failed; wrote pause flag: {PAUSE_FLAG}")
            except Exception as pause_exc:
                print(f"[WARNING] Failed to write pause flag after audit failure: {pause_exc}")
            raise
    elif not skip_audit:
        print(f"[WARN] Dataset audit script not found: {DATASET_AUDIT_SCRIPT} (set SKIP_DATA_AUDIT=1 to silence)")

    # Check Environment
    if not _wandb_credentials_present():
        print(
            "[WARNING] W&B credentials not detected (WANDB_API_KEY empty and no wandb login). Logging will be disabled."
        )
    else:
        print("[INFO] W&B credentials detected.")

    tokenizer_path = args.tokenizer_path
    if not tokenizer_path:
        env_vocab = os.environ.get("TOKENIZER_PATH")
        if env_vocab:
            tokenizer_path = env_vocab
        elif os.path.exists("/workspace/vocab_es.txt"):
            tokenizer_path = "/workspace/vocab_es.txt"
        else:
            tokenizer_path = "F5-TTS/ckpts/base_models/vocab.txt"

    cmd = [
        "accelerate",
        "launch",
        "-m",
        "f5_tts.train.finetune_cli",
        # Finetuning Mode
        "--finetune",
        "--pretrain",
        args.base_model_path,
        # Identity
        "--dataset_name",
        args.project_name,
        "--exp_name",
        args.exp_name,
        # Hyperparams
        "--epochs",
        str(args.epochs),
        "--batch_size_per_gpu",
        str(args.batch_size),
        "--learning_rate",
        "1e-5",
        # Observability (CRITICAL for V2)
        "--logger",
        "wandb",
        "--log_samples",  # Enables the patched validation logic
        "--save_per_updates",
        "1000",  # Validate every 1000 steps
        "--last_per_updates",
        "500",  # Save 'last' checkpoint frequently
        # Disk hygiene: keep only the last N checkpoints (multi-GB each). Override via KEEP_LAST_N_CHECKPOINTS.
        "--keep_last_n_checkpoints",
        os.environ.get("KEEP_LAST_N_CHECKPOINTS", "15"),
        # Data/Tokenizer
        "--tokenizer",
        args.tokenizer,
        "--tokenizer_path",
        tokenizer_path,
    ]

    print(f"Executing: {' '.join(cmd)}")

    try:
        subprocess.check_call(cmd)
    except subprocess.CalledProcessError as e:
        print(f"Training failed with exit code {e.returncode}")
        sys.exit(e.returncode)


if __name__ == "__main__":
    main()