#!/usr/bin/env python3
"""
voice — Lift a model's "voice" (its output-projection tensor) and cast it onto
another model, in either direction between safetensors and GGUF.

    voice get <model> [tensors...]           Get a voice: one tensor (or head) to the registry,
                                             several/--multi to voicepack files
    voice graft <pulse> [donor]            Graft one voice onto another (DELLA single graft)
    voice chorus [pulse] [donors...]       Blend voices: pulse + support + lead (DELLA chorus)
    voice cast <target> <voice> [tensors...] Cast a voice (or chosen tensors) onto a file
    voice delta [voice]                      Make a reusable style (voice − base)
    voice add <safetensors> <json>           Add a voice from local voice.safetensors + voice.json
    voice list [base]                        List stored voices (or cached bases)
    voice info <voice>                       Show voice details
    voice remove <voice>                     Remove a stored voice
    voice cleanup                            Remove incomplete downloads
    voice report                             Save a report locally and optionally share
    voice auth                               Save your Hugging Face token (optional, faster)
    voice path                               Put `voice` on your PATH
    voice doctor                             Check your setup

Every extract/cast action is appended to ~/.voices/operations.log.
"""

import argparse
import getpass
import json
import os
import re
import shutil
import struct
import sys
import textwrap
import time
import urllib.error
import urllib.request
from pathlib import Path
from urllib.parse import urlparse

# Launch bootstrap: vlib/ lives beside this file. The interpreter usually puts
# the script dir on sys.path, but isolated (-I), embedded, venv re-exec, and
_HERE = os.path.dirname(os.path.realpath(__file__))
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)
from vlib.ctx import VOICES_DIR, HF_TOKEN_PATH, OPERATIONS_LOG, REPORTS_JSONL, REPORTS_OUTBOX, REPORTS_DATASET_ID, REPORTS_REMOTE_PREFIX, VERSION, DTYPE_BYTES, VERBOSE
from vlib import ctx
from vlib.compute import _load_voice, _orient, _cast_st_prepare, _cast_core, _is_head_name, _parse_quant_override, _quant_from_key, _quantize_to_Q, _qerr, _nearest_encodable, _fuse_delta, _gguf_fuse_one, _gguf_write_repl, _read_gguf_arch, _confirm_arch_mismatch, _cast_gguf, _cast_gguf_pack, _cast_pack, _parse_tensor_list, _parse_exclude_list, _resolve_one_tensor, _tensor_src_meta, _load_tensor, _pack_aliases, _resolve_pack_targets
from vlib.net import _HF_TOKEN_CACHE, _get_hf_token, _HF_TRUSTED_HOSTS, _is_hf_url, _hf_auth_headers, _PreserveAuthRedirectHandler, _parse_retry_after, _bounded_retry_after, _timeout, http_range, http_get_json, _safe_id, _acquire_parts_lock, download_range, _fmt_eta, _cleanup_tmp
from vlib.tensors import _validate_safetensors_header, read_safetensors, write_safetensors, _write_safetensors_streaming, rewrite_safetensors_tensor, _safetensors_write_repl, _DTYPE_BYTES_MAP, dtype_bytes, bf16_to_f32, decode_to_f32, encode_from_f32, _gguf_tensor_f32, _gguf_raw_to_f32, _gguf_raw_slice_to_f32, _DELTA_CHUNK_ELEMS, _gguf_nbytes, _parse_gguf_header, _is_gguf_file, _QUANT_DTYPES, _is_ggml_quant, _is_quant_fetched, _is_delta_safetensors, _write_single_tensor_gguf
from vlib.sources import OUTPUT_ALIASES, EMBEDDING_ALIASES, _family, _nested, detect_output_tensor, largest_2d, _hf_config, _hf_license, _hf_model_index, _normalize_hf_repo, _base_paths, _hf_api_base_model, _hf_mergekit_base, _detect_base, TensorRef, SafetensorsSource, GgufSource, open_source, _source_from_repo, _pick_gguf, _sanitize_voice_name, _NAME_RULES, _public_source, _slugify_voice_name, _triage_gguf_files, _gguf_unit_census, _gguf_dtype_rank, _gguf_source_for_repo, role_of, is_tied, find_head, find_embed, pair_role
from vlib.fetch import _source_arch, _resolve_output_name, _stream_safetensors_tensor, _fetch_tensor, _stream_gguf_tensor, _fetch_pack, _write_fetched, _write_voicepack_json
from vlib.registry import log_operation, _atomic_jsonl_append, _write_outbox, _append_report, _iter_voices, _registry_resolve, _read_meta, _registry_name_exists, _registry_folder, _assert_not_symlink, _write_voice, _resolve_name, _out_path, _stem, _check_overwrite
from vlib.cache import _base_rev, _cache_validate, _cache_topup_write, _ensure_base_tensors, _ensure_base_tensors_locked, _cache_base, _load_safetensors_or_gguf_f32, _resolve_head_want
from vlib.delta import _choose_voice, _resolve_base, _delta_gate, _check_delta_gate, _compute_delta
from vlib.reports import _hash_file, _sanitize_note, _sanitize_reporter, _prompt_share, _push_reports_batch, _push_report, _last_operation, cmd_report
from vlib.cmds_voices import cmd_get, cmd_cast, cmd_delta, cmd_add
from vlib.cmds_chorus import cmd_chorus, cmd_graft
from vlib.cmds_misc import cmd_list, _list_bases, _voice_tensors, _verify_voice_file, cmd_info, cmd_remove, cmd_cleanup, cmd_auth, cmd_path, cmd_doctor, _print_version, _print_completion, _entry_script, _ensure_venv

def main():
    global VERBOSE
    parser = argparse.ArgumentParser(prog="voice", description="Lift a model's voice and cast it onto another model.")
    parser.add_argument("-v", "--verbose", action="store_true", help="Verbose: tensor shapes, decisions, and paths")
    parser.add_argument("--non-interactive", action="store_true", help="Fail instead of prompting (same as piped stdin; needs --yes to overwrite/remove)")
    parser.add_argument("--version", action="store_true", help="Show version and environment")
    parser.add_argument("--completion", choices=["bash", "zsh", "fish"], help="Print shell completion script")
    sub = parser.add_subparsers(dest="command")

    p_get = sub.add_parser("get", help="Get a voice from a model (safetensors repo, .gguf, or URL)")
    p_get.add_argument("source", help="HF model ID, .gguf/.safetensors path, or URL")
    p_get.add_argument("tensors", nargs="*", help="Tensor name(s) to fetch (default: auto-detect output head; 2+ or --multi saves a voicepack)")
    p_get.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")
    p_get.add_argument("--name", help="Voice name to save (registry voices only; not with --out)")
    p_get.add_argument("--multi", action="store_true", help="Fetch the named tensors (1+) as voicepack.safetensors + voicepack.json (not with --name)")
    p_get.add_argument("--out", help="Write a standalone file instead of a registry voice (default pack name: voicepack.safetensors)")
    p_get.add_argument("--yes", action="store_true", help="Assume yes to prompts (overwrite without asking)")

    p_cast = sub.add_parser("cast", help="Cast a voice onto a .gguf or .safetensors file")
    p_cast.add_argument("target", help="Path to the target .gguf or .safetensors")
    p_cast.add_argument("voice", help="Voice name, or a .safetensors/.gguf path")
    p_cast.add_argument("tensors", nargs="*", help="Only cast these tensors (default: every tensor in the voice)")
    p_cast.add_argument("--out", help="Output file or folder (default: ~/Voiced/<voice>-<model>/<voice>-<model>.<ext matching target>; a folder invents the filename)")
    p_cast.add_argument("--yes", action="store_true", help="Assume yes to prompts (overwrite without asking)")
    p_cast.add_argument("--quant", default="keep", help="Output quant override: keep (default, preserve quality) | target | any quant (q8_0,q4_0,q4_k,q6_k,iq4_nl,f16,f32,bf16,...; safetensors targets: keep|target|f32|f16|bf16)")
    p_cast.add_argument("--exclude", action="append", dest="excludes", default=[], help="Skip tensors on cast (repeatable: --exclude foo --exclude bar, or --exclude foo,bar)")
    p_cast.add_argument("--speak", action="store_true", help="Voice speaks through output.weight (invented if absent when head/embed geometry survives; GGUF targets only)")
    p_cast.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")
    p_delta = sub.add_parser("delta", help="Make a reusable style (voice − base)")
    p_delta.add_argument("voicename", nargs="?", help="Voice name (default: pick from registry)")
    p_delta.add_argument("--base", help="Base HF repo (org/model) or local path to base safetensors")
    p_delta.add_argument("--yes", action="store_true", help="Assume yes to prompts (overwrite without asking)")
    p_delta.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")
    p_graft = sub.add_parser("graft", help="Graft one voice onto another (DELLA single graft)")
    p_graft.add_argument("pulse", nargs="?", help="Pulse voice name (default: pick from registry)")
    p_graft.add_argument("donor", nargs="?", help="Donor voice name (default: pick from registry)")
    p_graft.add_argument("--density", type=float, default=0.8, help="Keep-fraction for magnitude pruning (default: 0.8)")
    p_graft.add_argument("--epsilon", type=float, default=0.1, help="Rank-ramp half-width (default: 0.1)")
    p_graft.add_argument("--lambda_", dest="lambda_", type=float, default=1.0, help="Global rescale after combine (default: 1.0)")
    p_graft.add_argument("--rescale", default="l1", choices=["l1", "inv_p"], help="Prune rescale: l1 (stable) or inv_p (paper-faithful)")
    p_graft.add_argument("--seed", type=int, default=42, help="Bernoulli seed for reproducible grafts (default: 42)")
    p_graft.add_argument("--out", help="Output voice name (default: chorus-<pulse>-<donor>)")
    p_graft.add_argument("--yes", action="store_true", help="Assume yes to prompts (overwrite without asking)")
    p_graft.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")
    p_chorus = sub.add_parser("chorus", help="Blend voices: pulse + support + lead (DELLA chorus)")
    p_chorus.add_argument("pulse", nargs="?", help="Pulse voice name (default: pick interactively)")
    p_chorus.add_argument("donors", nargs="*", help="Support + lead voice names (0-2; 1 = duet at lead density)")
    p_chorus.add_argument("--config", action="store_true", help="Stage knob screens per graft (Enter accepts each default; screens prefill from flags)")
    p_chorus.add_argument("--density", type=float, default=0.8, help="First-graft density (default: 0.8)")
    p_chorus.add_argument("--lead-density", type=float, default=0.9, help="Lead-graft density (default: 0.9)")
    p_chorus.add_argument("--epsilon", type=float, default=0.1, help="Rank-ramp half-width (default: 0.1)")
    p_chorus.add_argument("--lead-epsilon", type=float, default=0.05, help="Lead-graft epsilon (default: 0.05; keeps 0.9±eps valid)")
    p_chorus.add_argument("--lambda_", dest="lambda_", type=float, default=1.0, help="Global rescale after combine (default: 1.0)")
    p_chorus.add_argument("--rescale", default="l1", choices=["l1", "inv_p"], help="Prune rescale: l1 (stable) or inv_p (paper-faithful)")
    p_chorus.add_argument("--seed", type=int, default=42, help="Bernoulli seed for reproducible choruses (default: 42)")
    p_chorus.add_argument("--out", help="Final output voice name (default: chorus-<pulse>-<lead>)")
    p_chorus.add_argument("--yes", action="store_true", help="Assume yes to prompts (overwrite without asking)")
    p_chorus.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")
    p_add = sub.add_parser("add", help="Add a voice from local voice.safetensors + voice.json")
    p_add.add_argument("safetensors", help="Path to voice.safetensors (required)")
    p_add.add_argument("json", help="Path to voice.json (required)")
    p_add.add_argument("--name", help="Voice name to save (default from voice.json)")
    p_add.add_argument("--yes", action="store_true", help="Auto-confirm overwrite")
    p_add.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")

    p_list = sub.add_parser("list", help="List stored voices")
    p_list.add_argument("target", nargs="?", choices=["base"], metavar="base", help="Show cached base heads (same as --base)")
    p_list.add_argument("--base", dest="list_base", action="store_true", help="Show cached base heads instead of voices")
    p_list.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")

    p_info = sub.add_parser("info", help="Show voice details")
    p_info.add_argument("voicename", help="Voice name")
    p_info.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")

    p_rm = sub.add_parser("remove", help="Remove a stored voice")
    p_rm.add_argument("voice", help="Voice name")
    p_rm.add_argument("--yes", action="store_true", help="Assume yes to prompts (remove without asking)")
    p_rm.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")

    p_cleanup = sub.add_parser("cleanup", help="Remove incomplete downloads")
    p_cleanup.add_argument("--force", action="store_true", help="Also clear in-progress downloads (locked)")
    p_cleanup.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")

    p_auth = sub.add_parser("auth", help="Save your HF token (optional)")
    p_auth.add_argument("--clear", action="store_true", help="Remove the stored token")
    p_auth.add_argument("--show", action="store_true", help="Show the masked token")
    p_auth.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")

    p_report = sub.add_parser("report", help="Save a report locally and optionally share to community")
    p_report.add_argument("--message", "-m", help="Freeform note (240ch, sanitized)")
    p_report.add_argument("--voice", help="Voice name context")
    p_report.add_argument("--target", help="Target file basename context")
    p_report.add_argument("--reporter", help="Reporter id (optional, 64ch)")
    p_report.add_argument("--yes", action="store_true", help="Non-interactive: skip preview and auto-share to community")
    p_report.add_argument("--no-push", action="store_true", help="Local only, skip share prompt")
    p_report.add_argument("--sync", action="store_true", help="Sync pending outbox to community")
    p_report.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")
    p_report.add_argument("--operation", help=argparse.SUPPRESS)

    p_path = sub.add_parser("path", help="Put voice on your PATH")
    p_path.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")
    p_doctor = sub.add_parser("doctor", help="Check your setup")
    p_doctor.add_argument("-q", "--quiet", action="store_true", help="Quieter: progress and success notes off (warnings and errors still print)")

    args = parser.parse_args()
    VERBOSE = args.verbose
    ctx.VERBOSE = args.verbose
    ctx.QUIET = bool(getattr(args, "quiet", False))

    if getattr(args, "completion", None):
        _print_completion(args.completion)
        sys.exit(0)
    if getattr(args, "version", False):
        _print_version()
        sys.exit(0)
    if args.command is None:
        parser.print_help()
        sys.exit(2)

    _ensure_venv(args.command)

    dispatch = {
        "get": cmd_get, "cast": cmd_cast, "delta": cmd_delta, "add": cmd_add, "list": cmd_list,
        "info": cmd_info, "remove": cmd_remove, "cleanup": cmd_cleanup, "auth": cmd_auth,
        "report": cmd_report, "path": cmd_path, "doctor": cmd_doctor,
        "graft": cmd_graft, "chorus": cmd_chorus,
    }
    dispatch[args.command](args)

if __name__ == "__main__":
    main()