"""Command-line interface for deterministic Open-SWE-Traces profiling."""

from __future__ import annotations

import argparse
import json
import os
import shutil
import sys
from collections.abc import Callable, Iterator
from concurrent.futures import (
    FIRST_COMPLETED,
    Future,
    ProcessPoolExecutor,
    ThreadPoolExecutor,
    as_completed,
    wait,
)
from pathlib import Path
from typing import Any

from .heuristics import classify_metrics
from .io import (
    append_jsonl,
    get_sample_id,
    iter_jsonl,
    iter_parquet_file,
    iter_records,
    load_completed_ids,
    parquet_files,
    write_json,
    write_jsonl,
)
from .metrics import extract_metrics
from .sampling import iter_review_records, select_review_ids
from .summary import build_summary
from .tokenization import TokenCounter

_PROFILE_COUNTER: TokenCounter | None = None


def _initialize_profile_process(tokenizer_json: str | None) -> None:
    """Create one tokenizer per process instead of per trajectory."""

    global _PROFILE_COUNTER
    _PROFILE_COUNTER = TokenCounter(Path(tokenizer_json)) if tokenizer_json else None


def _profile_parquet_shard(
    shard: Path,
    metrics_part: Path,
    errors_part: Path | None,
) -> tuple[int, int]:
    """Profile one shard into atomic part files with bounded memory."""

    processed = 0
    failures = 0
    metrics_temporary = metrics_part.with_suffix(metrics_part.suffix + ".tmp")
    errors_temporary = (
        errors_part.with_suffix(errors_part.suffix + ".tmp") if errors_part else None
    )
    metrics_part.parent.mkdir(parents=True, exist_ok=True)
    with metrics_temporary.open("w", encoding="utf-8") as metrics_handle:
        errors_handle = (
            errors_temporary.open("w", encoding="utf-8")
            if errors_temporary is not None
            else None
        )
        try:
            for record in iter_parquet_file(shard):
                sample_id = get_sample_id(record)
                try:
                    result = extract_metrics(record, _PROFILE_COUNTER)
                    metrics_handle.write(
                        json.dumps(result, ensure_ascii=False, default=str) + "\n"
                    )
                except Exception as exc:  # noqa: BLE001 - isolate bad records.
                    failures += 1
                    if errors_handle is not None:
                        error = {
                            "sample_id": sample_id,
                            "error_type": type(exc).__name__,
                            "error": str(exc)[:2000],
                            "source_parquet": str(shard.resolve()),
                        }
                        errors_handle.write(
                            json.dumps(error, ensure_ascii=False) + "\n"
                        )
                processed += 1
        finally:
            if errors_handle is not None:
                errors_handle.close()
    os.replace(metrics_temporary, metrics_part)
    if errors_temporary is not None and errors_part is not None:
        os.replace(errors_temporary, errors_part)
    return processed, failures


def _merge_parts(parts: list[Path], output: Path) -> None:
    """Concatenate completed shard outputs atomically in dataset order."""

    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_suffix(output.suffix + ".tmp")
    with temporary.open("wb") as destination:
        for part in parts:
            with part.open("rb") as source:
                shutil.copyfileobj(source, destination, length=1024 * 1024)
    os.replace(temporary, output)


def _run_parallel_parquet_profile(args: argparse.Namespace) -> int:
    """Profile independent Parquet shards in separate CPU processes."""

    shards = parquet_files(args.input)
    parts_dir = args.output.parent / f".{args.output.name}.parts"
    parts_dir.mkdir(parents=True, exist_ok=True)
    metrics_parts = [parts_dir / f"{index:05d}.metrics.jsonl" for index in range(len(shards))]
    error_parts = [parts_dir / f"{index:05d}.errors.jsonl" for index in range(len(shards))]

    if not args.resume:
        for path in parts_dir.glob("*.jsonl"):
            path.unlink()
        for path in parts_dir.glob("*.tmp"):
            path.unlink()

    pending = []
    processed = 0
    failures = 0
    tokenizer_json = str(args.tokenizer_json) if args.tokenizer_json else None
    with ProcessPoolExecutor(
        max_workers=args.workers,
        initializer=_initialize_profile_process,
        initargs=(tokenizer_json,),
    ) as executor:
        for shard, metrics_part, errors_part in zip(
            shards, metrics_parts, error_parts, strict=True
        ):
            if args.resume and metrics_part.exists() and errors_part.exists():
                continue
            pending.append(
                executor.submit(
                    _profile_parquet_shard,
                    shard,
                    metrics_part,
                    errors_part if args.errors else None,
                )
            )
        for completed_shards, future in enumerate(as_completed(pending), 1):
            shard_processed, shard_failures = future.result()
            processed += shard_processed
            failures += shard_failures
            print(
                f"profile: shards={completed_shards}/{len(pending)} "
                f"processed={processed} failures={failures}",
                file=sys.stderr,
            )

    missing = [path for path in metrics_parts if not path.exists()]
    if missing:
        raise RuntimeError(f"Missing {len(missing)} completed metric parts")
    _merge_parts(metrics_parts, args.output)
    if args.errors:
        missing_errors = [path for path in error_parts if not path.exists()]
        if missing_errors:
            raise RuntimeError(f"Missing {len(missing_errors)} completed error parts")
        _merge_parts(error_parts, args.errors)
    print(f"profile: completed={processed} failures={failures}", file=sys.stderr)
    return 1 if failures else 0


def _write_error(path: Path | None, sample_id: str, error: Exception) -> None:
    """Write a compact error without interrupting a long profile run."""

    value = {
        "sample_id": sample_id,
        "error_type": type(error).__name__,
        "error": str(error)[:2000],
    }
    if path is None:
        print(json.dumps(value, ensure_ascii=False), file=sys.stderr)
    else:
        append_jsonl(path, value)


def _run_streaming_stage(
    *,
    input_path: Path,
    output_path: Path,
    errors_path: Path | None,
    limit: int | None,
    resume: bool,
    workers: int,
    processor: Callable[[dict[str, Any]], dict[str, Any]],
) -> int:
    """Profile records concurrently while serializing writes in one thread."""

    if workers < 1:
        raise ValueError("workers must be at least 1")
    completed = load_completed_ids(output_path) if resume else set()
    processed = 0
    failures = 0

    def selected_records() -> Iterator[dict[str, Any]]:
        selected = 0
        for record in iter_records(input_path):
            if get_sample_id(record) in completed:
                continue
            if limit is not None and selected >= limit:
                break
            selected += 1
            yield record

    def process_one(
        record: dict[str, Any],
    ) -> tuple[str, dict[str, Any] | None, Exception | None]:
        sample_id = get_sample_id(record)
        try:
            return sample_id, processor(record), None
        except Exception as exc:  # noqa: BLE001 - isolate individual records.
            return sample_id, None, exc

    def record_outcome(
        outcome: tuple[str, dict[str, Any] | None, Exception | None]
    ) -> None:
        nonlocal processed, failures
        sample_id, result, error = outcome
        if error is not None:
            failures += 1
            _write_error(errors_path, sample_id, error)
        elif result is not None:
            append_jsonl(output_path, result)
        processed += 1
        if processed % 1000 == 0:
            print(f"profile: processed={processed} failures={failures}", file=sys.stderr)

    records = iter(selected_records())
    if workers == 1:
        for record in records:
            record_outcome(process_one(record))
    else:
        with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="profile") as executor:
            pending: set[
                Future[tuple[str, dict[str, Any] | None, Exception | None]]
            ] = set()
            for _ in range(workers):
                try:
                    pending.add(executor.submit(process_one, next(records)))
                except StopIteration:
                    break
            while pending:
                finished, pending = wait(pending, return_when=FIRST_COMPLETED)
                for future in finished:
                    record_outcome(future.result())
                    try:
                        pending.add(executor.submit(process_one, next(records)))
                    except StopIteration:
                        pass

    print(f"profile: completed={processed} failures={failures}", file=sys.stderr)
    return 1 if failures else 0


def command_profile(args: argparse.Namespace) -> int:
    """Extract per-trajectory metrics without modifying source records."""

    if args.input.is_dir() and args.workers > 1 and args.limit is None:
        return _run_parallel_parquet_profile(args)
    counter = TokenCounter(args.tokenizer_json) if args.tokenizer_json else None
    return _run_streaming_stage(
        input_path=args.input,
        output_path=args.output,
        errors_path=args.errors,
        limit=args.limit,
        resume=args.resume,
        workers=args.workers,
        processor=lambda record: extract_metrics(record, counter),
    )


def command_summarize(args: argparse.Namespace) -> int:
    """Derive global thresholds, summary statistics, and decisions."""

    summary, threshold = build_summary(args.metrics)
    write_json(args.summary, summary)
    write_jsonl(
        args.decisions,
        (
            classify_metrics(metrics, threshold)
            for metrics in iter_jsonl(args.metrics)
        ),
    )
    return 0


def command_sample(args: argparse.Namespace) -> int:
    """Materialize a deterministic stratified review sample."""

    selected, decisions = select_review_ids(args.decisions, args.per_group, args.seed)
    write_jsonl(
        args.output,
        iter_review_records(args.input, selected, decisions),
    )
    print(f"sample: selected={len(selected)}", file=sys.stderr)
    return 0


def build_parser() -> argparse.ArgumentParser:
    """Build the intentionally small three-command interface."""

    parser = argparse.ArgumentParser(prog="swe-qc", description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)

    profile = subparsers.add_parser("profile", help="Extract deterministic metrics")
    profile.add_argument("--input", type=Path, required=True)
    profile.add_argument("--output", type=Path, required=True)
    profile.add_argument("--errors", type=Path)
    profile.add_argument("--tokenizer-json", type=Path)
    profile.add_argument("--limit", type=int)
    profile.add_argument("--resume", action="store_true")
    profile.add_argument("--workers", type=int, default=1)
    profile.set_defaults(func=command_profile)

    summarize = subparsers.add_parser(
        "summarize", help="Build summary and deterministic decisions"
    )
    summarize.add_argument("--metrics", type=Path, required=True)
    summarize.add_argument("--summary", type=Path, required=True)
    summarize.add_argument("--decisions", type=Path, required=True)
    summarize.set_defaults(func=command_summarize)

    sample = subparsers.add_parser(
        "sample", help="Sample full trajectories by rule and length bucket"
    )
    sample.add_argument("--input", type=Path, required=True)
    sample.add_argument("--decisions", type=Path, required=True)
    sample.add_argument("--output", type=Path, required=True)
    sample.add_argument("--per-group", type=int, default=20)
    sample.add_argument("--seed", type=int, default=20260818)
    sample.set_defaults(func=command_sample)
    return parser


def main(argv: list[str] | None = None) -> int:
    """Run the selected command."""

    args = build_parser().parse_args(argv)
    return int(args.func(args))


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