"""Command line interface for the voice library search engine."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

from . import DEFAULT_MODEL, __version__

DEFAULT_DIR = "data/voices"
DEFAULT_MODELS = "models"


def cmd_fetch(args) -> int:
    from .clap import ClapEncoder
    from .voices import (VoiceCatalog, download_previews, embed_voices, fetch_voices)

    root = Path(args.dir).expanduser()
    print(f"Fetching up to {args.limit} voices from the ElevenLabs Voice Library...")
    voices = fetch_voices(limit=args.limit,
                          progress=lambda n, t: print(f"\r  {n}/{t}", end="", flush=True))
    print(f"\nDownloading previews...")
    errors = download_previews(voices, root / "previews",
                               progress=lambda i, t: print(f"\r  {i}/{t}", end="", flush=True))
    print("\nEmbedding (runs locally, nothing is uploaded)...")
    encoder = ClapEncoder(model_id=args.model)
    kept, vecs, embed_errors = embed_voices(
        voices, encoder, windows=args.windows,
        progress=lambda i, t: print(f"\r  {i}/{t}", end="", flush=True))

    catalog = VoiceCatalog(root=root, model_id=encoder.model_id, voices=kept, vectors=vecs)
    catalog.save()
    print(f"\n{len(kept)} voices embedded -> {root}")

    from .evaluate import embedding_health
    health = embedding_health(vecs)
    if health.get("count", 0) >= 2:
        print(f"Embedding spread: mean pairwise cosine "
              f"{health['mean_pairwise_cosine']:.3f} (std {health['std_pairwise_cosine']:.3f})")
        if health.get("collapsed"):
            print("WARNING: vectors have collapsed -- this checkpoint's tower is\n"
                  "         degenerate and retrieval will be meaningless.", file=sys.stderr)
    failed = len(errors) + len(embed_errors)
    if failed:
        print(f"{failed} voice(s) failed", file=sys.stderr)
        for e in (errors + embed_errors)[:5]:
            print(f"  {e}", file=sys.stderr)
    return 0


def _load(args):
    from .clap import ClapEncoder
    from .voices import VoiceCatalog

    catalog = VoiceCatalog.load(Path(args.dir).expanduser())
    if not len(catalog) or catalog.vectors is None:
        print(f"no voices at {args.dir} -- run `mss fetch` first", file=sys.stderr)
        return None, None
    return catalog, ClapEncoder(model_id=catalog.model_id)


def cmd_search(args) -> int:
    from .voices import search_voices

    catalog, encoder = _load(args)
    if catalog is None:
        return 1
    from .voices import matching_indices

    filters = {k: getattr(args, k) for k in ("gender", "age", "accent", "language")
               if getattr(args, k, None)}
    # computed once and reused for both the search and the scope line below
    keep = matching_indices(catalog, filters)
    hits = search_voices(catalog, encoder, args.query, top_n=args.top_n, keep=keep)
    if not hits and filters:
        print(f"no voices match {filters}", file=sys.stderr)
        return 1
    if args.json:
        print(json.dumps([{"rank": h.rank, "score": round(h.score, 6),
                           "name": h.name, "metadata": h.metadata}
                          for h in hits], indent=2, ensure_ascii=False))
        return 0
    scope = f"{len(catalog)} voices"
    if filters:
        scope = (f"{len(keep)} of {len(catalog)} voices "
                 f"matching {', '.join(f'{k}={v}' for k, v in filters.items())}")
    print(f'\n"{args.query}"  ->  top {len(hits)} of {scope}\n')
    for h in hits:
        m = h.metadata
        attrs = "/".join(str(m.get(k) or "?") for k in ("gender", "age", "use_case"))
        print(f"  {h.rank:>2}. {h.score:.4f}  {h.name[:42]:44s} {attrs}")
    print()
    return 0


def cmd_eval(args) -> int:
    from .voices import attribute_retrieval, linear_probe

    catalog, encoder = _load(args)
    if catalog is None:
        return 1
    print(f"{len(catalog)} voices, model {catalog.model_id}\n")

    print("linear probe (is the attribute decodable at all?):")
    for attr in ("gender", "age", "use_case", "language"):
        r = linear_probe(catalog, attr)
        print(f"  {r['attribute']:9s} acc={r['accuracy']:.3f} "
              f"majority={r['majority_baseline']:.3f} lift={r['lift']:+.3f} | "
              f"balanced={r['balanced_accuracy']:.3f} chance={r['chance']:.3f} "
              f"lift={r['balanced_lift']:+.3f}")

    print("\nretrieval precision@10:")
    for attr, value, query in [
        ("gender", "female", "a recording of a woman speaking"),
        ("gender", "male", "a recording of a man speaking"),
        ("age", "old", "a recording of an elderly person speaking"),
    ]:
        r = attribute_retrieval(catalog, encoder, attr, value, query)
        print(f"  {r['target']:14s} P@10={r['precision_at_k']:.2f} "
              f"(prevalence {r['prevalence']:.2f})")
    return 0


def cmd_info(args) -> int:
    from .evaluate import embedding_health
    from .voices import VoiceCatalog

    catalog = VoiceCatalog.load(Path(args.dir).expanduser())
    if not len(catalog):
        print(f"No voices at {Path(args.dir).expanduser()}")
        return 1
    have_audio = sum(1 for v in catalog.voices if catalog.preview_path(v))
    print(f"Catalog   : {catalog.root}")
    print(f"Voices    : {len(catalog)}  ({have_audio} with preview audio)")
    print(f"Model     : {catalog.model_id}")
    if catalog.vectors is not None:
        print(f"Vectors   : {catalog.vectors.shape[0]} x {catalog.vectors.shape[1]} float32")
        health = embedding_health(catalog.vectors)
        if health.get("count", 0) >= 2:
            print(f"Spread    : mean pairwise cosine {health['mean_pairwise_cosine']:.3f}"
                  f" (std {health['std_pairwise_cosine']:.3f})")
            if health.get("collapsed"):
                print("WARNING   : vectors have collapsed -- retrieval is meaningless.")
    return 0


def cmd_export(args) -> int:
    from .export_onnx import export

    try:
        export(args.dest, model_id=args.model, keep_fp32=args.keep_fp32)
    except ImportError as exc:
        print(f"missing export dependency: {exc}\n"
              'install with: uv pip install -e ".[export]"', file=sys.stderr)
        return 1
    return 0


def cmd_serve(args) -> int:
    from .api import serve

    serve(args.dir, host=args.host, port=args.port, models=args.models)
    return 0


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(
        prog="mss",
        description="Semantic search over the ElevenLabs Voice Library (CLAP).")
    p.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
    sub = p.add_subparsers(dest="command", required=True)

    def add_dir(sp):
        sp.add_argument("-d", "--dir", default=DEFAULT_DIR,
                        help=f"voice catalog directory (default: {DEFAULT_DIR})")

    f = sub.add_parser("fetch", help="fetch voices, download previews, embed them")
    add_dir(f)
    f.add_argument("--limit", type=int, default=1000)
    f.add_argument("--model", default=DEFAULT_MODEL)
    f.add_argument("--windows", type=int, default=2)
    f.set_defaults(func=cmd_fetch)

    s = sub.add_parser("search", help="text query -> ranked voices")
    s.add_argument("query")
    add_dir(s)
    s.add_argument("-n", "--top-n", type=int, default=10)
    s.add_argument("--json", action="store_true")
    # Filters, not query words: the label is exact where CLAP is 0.90 at best
    # on gender and at chance on accent and language.
    for attr in ("gender", "age", "accent", "language"):
        s.add_argument(f"--{attr}", default=None, help=f"restrict to this {attr}")
    s.set_defaults(func=cmd_search)

    e = sub.add_parser("eval", help="zero-shot, linear probe and retrieval metrics")
    add_dir(e)
    e.set_defaults(func=cmd_eval)

    n = sub.add_parser("info", help="show catalog status")
    add_dir(n)
    n.set_defaults(func=cmd_info)

    x = sub.add_parser("export", help="export the text tower to int8 ONNX")
    x.add_argument("--dest", default=DEFAULT_MODELS,
                   help=f"output directory (default: {DEFAULT_MODELS})")
    x.add_argument("--model", default=DEFAULT_MODEL)
    x.add_argument("--keep-fp32", action="store_true",
                   help="keep the intermediate fp32 graph")
    x.set_defaults(func=cmd_export)

    v = sub.add_parser("serve", help="run the API backing the web UI")
    add_dir(v)
    v.add_argument("--host", default="127.0.0.1")
    v.add_argument("--port", type=int, default=8787)
    v.add_argument("--models", default=DEFAULT_MODELS,
                   help=f"ONNX export directory, if built (default: {DEFAULT_MODELS})")
    v.set_defaults(func=cmd_serve)

    args = p.parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    raise SystemExit(main())