import argparse
import sys
from dataclasses import dataclass
from typing import List, Tuple, Optional
from ezexl3 import __version__

@dataclass
class PassThrough:
    quant_args: List[str]
    measure_args: List[str]
    cleaned_argv: List[str]

def _split_passthrough(argv: List[str]) -> PassThrough:
    """
    Extract two passthrough blocks:

      --quant-args -- <...>
      --measure-args -- <...>

    Everything else stays in cleaned_argv for normal argparse parsing.

    Notes:
    - The '--' delimiter is REQUIRED if you supply a block.
    - Order can be either (quant then measure) or (measure then quant), but
      each block may appear at most once.
    """
    quant_args: List[str] = []
    measure_args: List[str] = []

    cleaned: List[str] = []
    i = 0
    n = len(argv)

    def read_block(start_i: int) -> Tuple[List[str], int]:
        # Expect: <flag> -- <args...>
        if start_i + 1 >= n or argv[start_i + 1] != "--":
            raise SystemExit(
                f"Expected '--' after {argv[start_i]}. Example: {argv[start_i]} -- -d 0,1 -dr 1,1"
            )
        j = start_i + 2
        block: List[str] = []
        while j < n and argv[j] not in ("--quant-args", "--measure-args"):
            block.append(argv[j])
            j += 1
        return block, j

    while i < n:
        tok = argv[i]
        if tok == "--quant-args":
            if quant_args:
                raise SystemExit("Duplicate --quant-args block")
            block, i = read_block(i)
            quant_args = block
            continue
        if tok == "--measure-args":
            if measure_args:
                raise SystemExit("Duplicate --measure-args block")
            block, i = read_block(i)
            measure_args = block
            continue

        cleaned.append(tok)
        i += 1

    return PassThrough(quant_args=quant_args, measure_args=measure_args, cleaned_argv=cleaned)


def _csv_or_space_list(values: List[str]) -> List[str]:
    """
    Allows: -b 2 3 4 5 6
    or:     -b 2,3,4,5,6
    """
    out: List[str] = []
    for v in values:
        parts = [p.strip() for p in v.split(",") if p.strip()]
        out.extend(parts)
    return out




def _parse_devices(values: List[str]) -> List[int]:
    if not values:
        raise SystemExit("At least one CUDA device must be provided (e.g. -d 0)")
    out: List[int] = []
    for raw in values:
        try:
            out.append(int(raw))
        except ValueError as e:
            raise SystemExit(f"Invalid CUDA device '{raw}'. Expected integer device ids like: -d 0,1") from e
    if not out:
        raise SystemExit("At least one CUDA device must be provided (e.g. -d 0)")
    return out


def _parse_device_ratios(values: Optional[List[str]], devices: List[int]) -> Optional[List[str]]:
    if values is None:
        return None
    if not values:
        raise SystemExit("--device-ratios cannot be empty. Example: -r 1,1")

    parsed: List[str] = []
    for raw in values:
        try:
            ratio = float(raw)
        except ValueError as e:
            raise SystemExit(f"Invalid device ratio '{raw}'. Expected numeric values like: -r 1,1") from e
        if ratio <= 0:
            raise SystemExit(f"Invalid device ratio '{raw}'. Ratios must be > 0")
        parsed.append(raw)

    if len(parsed) != len(devices):
        raise SystemExit(
            f"--device-ratios length ({len(parsed)}) must match --devices length ({len(devices)})."
        )

    return parsed



def _parse_layers(value: int) -> int:
    if value not in (1, 2, 3):
        raise SystemExit("--layers must be one of: 1, 2, 3")
    return value

def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="ezexl3",
        description="ezexl3: simple single-command EXL3 repo generator"
    )
    p.add_argument("-v", "--version", action="version", version=f"%(prog)s {__version__}")

    sub = p.add_subparsers(dest="cmd", required=True)

    def add_repo_flags(p_sub: argparse.ArgumentParser) -> None:
        p_sub.add_argument("-m", "--models", nargs="+", required=True,
                           help="One or more BF16/base model directories (space or comma separated)")
        p_sub.add_argument(
            "-b", "--bpws",
            required=True,
            nargs="+",
            help="Target BPWs (space-separated or comma-separated). Example: -b 2 3 4 5 6 or -b 2,3,4,5,6",
        )
        p_sub.add_argument(
            "-d", "--devices",
            default="0",
            help="CUDA devices for quant+measure. Example: -d 0,1",
        )
        p_sub.add_argument(
            "-r", "--device-ratios",
            default=None,
            help="Device ratios for quantization only. Example: -r 1,1 (optional)",
        )
        p_sub.add_argument("-hq", action="store_true", help="Enable high-quality quantization (exllamav3 -hq)")
        p_sub.add_argument("--no-cleanup", "-nc", action="store_true", help="Keep w-* working dirs and logs")
        p_sub.add_argument("--no-readme", action="store_true", help="Skip README stage")
        p_sub.add_argument("--no-logs", action="store_true", help="Do not write per-GPU logs")
        p_sub.add_argument("--no-prompt", "-np", action="store_true", help="Use defaults for README instead of prompting")
        p_sub.add_argument("--no-graph", "-ng", action="store_true", help="Do not generate or embed the README SVG graph")
        p_sub.add_argument("--no-measurement", "-nm", action="store_true", help="Skip KL/PPL measurements (also disables README graph and KL/PPL table columns)")
        p_sub.add_argument("--template", "-t", help="README template name (e.g., 'fire', 'basic')")
        p_sub.add_argument("-l", "--layers", type=int, default=2, choices=[1, 2, 3], help="Layers used by optimized comparative measure stage (1-3, default: 2)")
        p_sub.add_argument("-cb", "--catbench", type=int, default=0, nargs="?", const=3,
                           help="Run SVG Catbench with N samples per model (default: 3 when flag present)")
        p_sub.add_argument("--no-verify", "-nv", action="store_true",
                           help="Skip per-BPW verification (batch all quants, then batch all measures)")

    # --- repo (main command) ---
    repo = sub.add_parser("repo", help="Generate an EXL3 repo (quantize -> measure -> README)")
    add_repo_flags(repo)

    # --- quantize ---
    q = sub.add_parser("quantize", aliases=["quant"], help="Quantize only (vendored multiConvert)")
    q.add_argument("-m", "--models", nargs="+", required=True,
                   help="One or more input model directories (space or comma separated).")
    q.add_argument("-b", "--bpws", nargs="+", required=True,
                   help="Target BPWs (space or comma separated).")
    q.add_argument("-d", "--devices", default="0", help="CUDA devices. Example: -d 0,1")
    q.add_argument("-r", "--device-ratios", default=None, help="Device ratios. Example: -r 1,1")
    q.add_argument("--out-template", default="{model}/{bpw}",
                   help="Template for output directory. Fields: {model}, {model_name}, {bpw}")
    q.add_argument("--w-template", default="{model}/w-{bpw}",
                   help="Template for working directory. Fields: {model}, {model_name}, {bpw}")
    q.add_argument("-hq", action="store_true", help="Enable high-quality quantization (exllamav3 -hq)")
    q.add_argument("--dry", action="store_true", help="Print what would run, but do not execute.")
    q.add_argument("--continue-on-error", action="store_true", help="Keep going after failures.")
    q.add_argument("--no-logs", action="store_true", help="Do not write per-GPU logs")
    q.add_argument("-l", "--layers", type=int, default=2, choices=[1, 2, 3], help="Layers used by optimized comparative measure stage (1-3, default: 2)")

    # --- measure ---
    m = sub.add_parser("measure", help="Measure only (vendored quantMeasure)")
    m.add_argument("-m", "--models", nargs="+", required=True, help="One or more model directories")
    m.add_argument("-b", "--bpws", nargs="+", required=True, help="BPWs to measure (space or comma separated)")
    m.add_argument("-d", "--devices", default="0", help="CUDA devices for measurement. Example: -d 0,1")
    m.add_argument("--no-logs", action="store_true", help="Do not write per-GPU logs")
    m.add_argument("--no-cleanup", "-nc", action="store_true", help="Keep temporary shard CSVs and logs")
    m.add_argument("-cb", "--catbench", type=int, default=0, nargs="?", const=3,
                   help="Run SVG Catbench with N samples per model (default: 3 when flag present)")


    # --- chat ---
    ch = sub.add_parser("chat", help="Launch web chat UI (model optional — select in browser)")
    ch.add_argument("-m", "--model", required=False, default=None, help="Model directory (optional: select in UI)")
    ch.add_argument("-d", "--devices", default="0", help="CUDA devices. Example: -d 0,1")
    ch.add_argument("-r", "--device-ratios", default=None, help="Device ratios. Example: -r 1,1")
    ch.add_argument("--host", default="127.0.0.1",
                    help="Bind address (default: 127.0.0.1). WARNING: non-loopback addresses expose the unauthenticated API to the network")
    ch.add_argument("--port", type=int, default=8800, help="Port (default: 8800)")
    ch.add_argument("--no-browser", action="store_true", help="Don't auto-open browser")
    ch.add_argument("-cs", "--cache-size", type=int, default=None,
                    help="Cache size in tokens (default: 32768). Must be multiple of 256")
    ch.add_argument("-cq", "--cache-quant", type=str, default=None,
                    help="Cache quantization bits: kv_bits or k_bits,v_bits (default: 6,6)")

    # --- readme ---
    r = sub.add_parser("readme", help="README only (CSV -> README)")
    r.add_argument("-m", "--models", nargs="+", required=True, help="One or more model directories")
    r.add_argument("--no-prompt", "-np", action="store_true", help="Use defaults for README instead of prompting")
    r.add_argument("--no-graph", "-ng", action="store_true", help="Do not generate or embed the README SVG graph")
    r.add_argument("--no-measurement", "-nm", action="store_true", help="Remove KL/PPL columns from README and skip graph embedding")
    r.add_argument("--template", "-t", help="README template name (e.g., 'fire', 'basic')")

    # --- ui (dashboard) ---
    ui = sub.add_parser("ui", aliases=["dash", "dashboard"],
                        help="Launch dashboard web UI")
    ui.add_argument("--host", default="127.0.0.1",
                    help="Bind address (default: 127.0.0.1)")
    ui.add_argument("--port", type=int, default=8801, help="Port (default: 8801)")
    ui.add_argument("--no-browser", action="store_true", help="Don't auto-open browser")

    return p


def main(argv: Optional[List[str]] = None) -> int:
    argv = list(sys.argv[1:] if argv is None else argv)

    # Extract passthrough blocks FIRST, then parse the cleaned argv normally.
    pt = _split_passthrough(argv)

    parser = build_parser()
    args = parser.parse_args(pt.cleaned_argv)

    # Initialize common fields
    cmd = getattr(args, "cmd", None)
    if not cmd:
        parser.print_help()
        return 0

    # Normalize lists
    if hasattr(args, "models"):
        args.models = _csv_or_space_list(args.models)
    if hasattr(args, "bpws"):
        args.bpws = _csv_or_space_list(args.bpws)
    if hasattr(args, "devices"):
        args.devices = [d.strip() for d in str(args.devices).split(",") if d.strip()]
    if hasattr(args, "device_ratios") and args.device_ratios is not None:
        args.device_ratios = [x.strip() for x in str(args.device_ratios).split(",") if x.strip()]
    

    import os

    if cmd == "chat":
        from ezexl3.chat.server import run_server
        # When -m is omitted, launch UI-only (no model pre-loaded)
        chat_devices = None
        dr = None
        if args.model:
            chat_devices = [int(d) for d in args.devices]
            dr = args.device_ratios
            if dr:
                dr = ",".join(dr) if isinstance(dr, list) else dr
        run_server(
            model_dir=args.model,
            devices=chat_devices,
            device_ratios=dr,
            cache_size=args.cache_size,
            cache_quant=args.cache_quant,
            host=args.host,
            port=args.port,
            open_browser=not args.no_browser,
        )
        return 0

    if cmd in ("ui", "dash", "dashboard"):
        from ezexl3.ui.server import run_ui_server
        run_ui_server(
            host=args.host,
            port=args.port,
            open_browser=not args.no_browser,
        )
        return 0

    from ezexl3.repo import run_repo, run_quant_stage, run_measure_stage

    devices_i = _parse_devices(getattr(args, "devices", ["0"]))
    device_ratios = _parse_device_ratios(getattr(args, "device_ratios", None), devices_i)
    device_ratios_str = ",".join(device_ratios) if device_ratios else None
    layers = _parse_layers(getattr(args, "layers", 2)) if hasattr(args, "layers") else 2

    # Inject -hq into quant_args when flag is set
    if getattr(args, "hq", False) and "-hq" not in pt.quant_args:
        pt.quant_args.append("-hq")

    if cmd == "repo":
        # Process each model, continuing on error
        failed_models: List[str] = []
        for model_dir in args.models:
            model_name = os.path.basename(os.path.abspath(model_dir))
            print(f"\n{'='*60}")
            print(f"Processing model: {model_name}")
            print(f"{'='*60}")

            try:
                rc = run_repo(
                    model_dir=model_dir,
                    bpws=args.bpws,
                    devices=devices_i,
                    device_ratios=device_ratios_str,
                    quant_args=pt.quant_args,
                    measure_args=pt.measure_args,
                    do_quant=True,
                    do_measure=(not args.no_measurement),
                    do_readme=(not args.no_readme),
                    cleanup=(not args.no_cleanup),
                    write_logs=(not args.no_logs),
                    interactive=(not args.no_prompt),
                    include_graph=(not args.no_graph and not args.no_measurement),
                    include_measurements=(not args.no_measurement),
                    template=args.template,
                    optimized_measure_layers=layers,
                    catbench_n=getattr(args, "catbench", 0) or 0,
                    verify=(not args.no_verify),
                )
                if rc != 0:
                    failed_models.append(model_dir)
            except Exception as e:
                print(f"Error processing {model_name}: {e}")
                import traceback
                traceback.print_exc()
                failed_models.append(model_dir)

        if failed_models:
            print(f"\n{'='*60}")
            print(f"Completed with {len(failed_models)} failure(s): {failed_models}")
            print(f"{'='*60}")
            return 1
        return 0

    if cmd in ("quant", "quantize"):
        from ezexl3.repo import _plan_repo_bpws, _run_optimized_opt_stage

        bpw_plan = _plan_repo_bpws(args.bpws)
        quant_bpws = bpw_plan["quant_integer_queue"]
        optimized_bpws = bpw_plan["requested_optimizeds"]

        if optimized_bpws and args.out_template != "{model}/{bpw}":
            print("Error: --out-template cannot be customized when using decimal BPWs.")
            print("The optimized quantization stage requires outputs at {model}/{bpw}.")
            return 1

        auto_added = [b for b in quant_bpws if b not in bpw_plan["requested_integers"]]
        if auto_added:
            print(
                "ℹ️ Added required integer quants for optimized targets: "
                + ", ".join(auto_added)
            )

        failed_models: List[str] = []
        for model_dir in args.models:
            print(f"\nQuantizing model: {model_dir}")
            try:
                rc = run_quant_stage(
                    model_dir=model_dir,
                    bpws=quant_bpws,
                    devices=devices_i,
                    device_ratios=device_ratios_str,
                    quant_args=pt.quant_args,
                    out_template=args.out_template,
                    w_template=args.w_template,
                    dry_run=args.dry,
                    continue_on_error=args.continue_on_error,
                    optimized_measure_layers=layers,
                )
                if rc != 0:
                    failed_models.append(model_dir)
                    continue

                if optimized_bpws and not args.dry:
                    _run_optimized_opt_stage(
                        model_dir=os.path.abspath(model_dir),
                        optimized_bpws=optimized_bpws,
                        devices=devices_i,
                        layers=layers,
                        write_logs=not args.no_logs,
                    )
            except Exception as e:
                print(f"Error quantizing {model_dir}: {e}")
                failed_models.append(model_dir)
        return 1 if failed_models else 0

    if cmd == "measure":
        failed_models: List[str] = []
        for model_dir in args.models:
            print(f"\nMeasuring model: {model_dir}")
            try:
                rc = run_measure_stage(
                    model_dir=model_dir,
                    bpws=args.bpws,
                    devices=devices_i,
                    write_logs=(not args.no_logs),
                    measure_args=pt.measure_args,
                    catbench_n=getattr(args, "catbench", 0) or 0,
                )
                if rc != 0:
                    failed_models.append(model_dir)
            except Exception as e:
                print(f"Error measuring {model_dir}: {e}")
                failed_models.append(model_dir)
        return 1 if failed_models else 0

    if cmd == "readme":
        from ezexl3.readme import run_readme
        for model_dir in args.models:
            run_readme(
                model_dir,
                template_name=args.template,
                interactive=(not args.no_prompt),
                include_graph=(not args.no_graph and not args.no_measurement),
                include_measurements=(not args.no_measurement),
            )
        return 0

    print(f"Command '{args.cmd}' not implemented yet.")
    return 1


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