"""Parser setup for the unified 'generate' command.

Consolidates weight modification, response generation, steering vectors,
classifiers, and score evaluation into a single command with --target flag.

Usage:
    wisent generate --target weights --task coding --model Qwen/Qwen3-8B
    wisent generate --target responses --task arc_easy --model Qwen/Qwen3-8B
    wisent generate --target vector --task coding --model Qwen/Qwen3-8B
    wisent generate --target classifier --task coding --model Qwen/Qwen3-8B
    wisent generate --target scores --responses responses.json
"""

import argparse


def setup_generate_parser(subparsers: argparse._SubParsersAction) -> None:
    """Set up the unified generate command parser."""
    
    generate_parser = subparsers.add_parser(
        "generate",
        help="Generate model outputs (weights, responses, vectors, classifiers, scores)",
        description=(
            "Unified generation command.\n\n"
            "Examples:\n"
            "  wisent generate --target weights --task coding --model Qwen/Qwen3-8B --output-dir ./modified\n"
            "  wisent generate --target responses --task arc_easy --model Qwen/Qwen3-8B --output responses.json\n"
            "  wisent generate --target vector --task coding --model Qwen/Qwen3-8B --output vector.json\n"
            "  wisent generate --target classifier --task coding --model Qwen/Qwen3-8B --output ./classifiers\n"
            "  wisent generate --target scores --responses responses.json --output scores.json"
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    
    # Target selection (required)
    generate_parser.add_argument(
        "--target",
        type=str,
        required=True,
        choices=["weights", "responses", "vector", "classifier", "scores"],
        help="What to generate: weights, responses, vector, classifier, or scores"
    )
    
    # Common arguments
    generate_parser.add_argument(
        "--model",
        type=str,
        default=None,
        help="Model name or path"
    )
    
    generate_parser.add_argument(
        "--task",
        type=str,
        default=None,
        help=(
            "Task for generation. Can be: "
            "skill name (coding, math), "
            "risk name (hallucination, toxicity), "
            "special case (refusal, personalization, humanizer, custom), "
            "benchmark name (arc_easy, gsm8k), "
            "comma-separated benchmarks (arc_easy,gsm8k), "
            "or free-form description ('be more concise')"
        )
    )
    
    generate_parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Output file or directory"
    )
    
    generate_parser.add_argument(
        "--device",
        type=str,
        default=None,
        help="Device to run on (default: auto)"
    )
    
    generate_parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable verbose output"
    )
    
    # Weights-specific arguments
    weights_group = generate_parser.add_argument_group("weights generation (--target weights)")
    
    weights_group.add_argument(
        "--output-dir",
        type=str,
        default="./modified_model",
        help="Output directory for modified model"
    )
    
    weights_group.add_argument(
        "--method",
        type=str,
        choices=["directional", "additive"],
        default="directional",
        help="Weight modification method (default: directional)"
    )
    
    weights_group.add_argument(
        "--strength",
        type=float,
        default=1.0,
        help="Steering strength (default: 1.0)"
    )
    
    weights_group.add_argument(
        "--steering-vectors",
        type=str,
        default=None,
        help="Path to pre-computed steering vectors (skip generation)"
    )
    
    weights_group.add_argument(
        "--num-pairs",
        type=int,
        default=100,
        help="Number of contrastive pairs for vector generation (default: 100)"
    )
    
    weights_group.add_argument(
        "--cap-pairs-per-benchmark",
        type=int,
        default=None,
        help="Cap pairs per benchmark for multi-benchmark training"
    )
    
    weights_group.add_argument(
        "--trait",
        type=str,
        default=None,
        help="Trait description for personalization or custom tasks"
    )
    
    # Responses-specific arguments
    responses_group = generate_parser.add_argument_group("responses generation (--target responses)")
    
    responses_group.add_argument(
        "--num-questions",
        type=int,
        default=10,
        help="Number of questions to generate responses for (default: 10)"
    )
    
    responses_group.add_argument(
        "--max-new-tokens",
        type=int,
        default=128,
        help="Maximum tokens to generate per response (default: 128)"
    )
    
    responses_group.add_argument(
        "--temperature",
        type=float,
        default=0.7,
        help="Temperature for generation (default: 0.7)"
    )
    
    responses_group.add_argument(
        "--use-steering",
        action="store_true",
        help="Use steering during response generation"
    )
    
    responses_group.add_argument(
        "--disable-thinking",
        action="store_true",
        help="Disable thinking/reasoning mode for Qwen models"
    )
    
    # Vector-specific arguments
    vector_group = generate_parser.add_argument_group("vector generation (--target vector)")
    
    vector_group.add_argument(
        "--trait-label",
        type=str,
        default="correctness",
        help="Label for the trait being steered (default: correctness)"
    )
    
    vector_group.add_argument(
        "--layers",
        type=str,
        default=None,
        help="Comma-separated layer indices (e.g., '8,12,16') or 'all'"
    )
    
    vector_group.add_argument(
        "--token-aggregation",
        type=str,
        choices=["average", "final", "first", "max", "continuation"],
        default="average",
        help="How to aggregate token activations (default: average)"
    )
    
    vector_group.add_argument(
        "--normalize",
        action="store_true",
        default=True,
        help="L2-normalize steering vectors (default: True)"
    )
    
    vector_group.add_argument(
        "--keep-intermediate",
        action="store_true",
        help="Keep intermediate files (pairs and enriched pairs)"
    )
    
    # Classifier-specific arguments
    classifier_group = generate_parser.add_argument_group("classifier generation (--target classifier)")
    
    classifier_group.add_argument(
        "--layer",
        type=int,
        default=None,
        help="Layer to use for classification"
    )
    
    classifier_group.add_argument(
        "--aggregation",
        type=str,
        choices=["average", "final", "first", "max", "min"],
        default="average",
        help="Token aggregation method for classification (default: average)"
    )
    
    classifier_group.add_argument(
        "--threshold",
        type=float,
        default=0.5,
        help="Detection threshold for classification (default: 0.5)"
    )
    
    classifier_group.add_argument(
        "--limit",
        type=int,
        default=1000,
        help="Maximum samples for classifier training (default: 1000)"
    )
    
    # Scores-specific arguments
    scores_group = generate_parser.add_argument_group("scores generation (--target scores)")
    
    scores_group.add_argument(
        "--responses",
        type=str,
        default=None,
        help="Input JSON file with generated responses to evaluate"
    )
    
    scores_group.add_argument(
        "--baseline",
        type=str,
        default=None,
        help="Baseline responses JSON file for comparison"
    )
    
    scores_group.add_argument(
        "--trait-description",
        type=str,
        default=None,
        help="Description of the personality trait for evaluation"
    )
