"""
CLI command for modifying model weights using steering vectors.

This module implements the modify-weights command which permanently modifies
model weights using either directional projection or additive methods.

By default, uses Norm-Preserving Biprojected Directional Modification
which maintains model quality by preserving weight norms.
"""

import json
import sys
import time
from pathlib import Path
import torch

from wisent.core.cli_logger import setup_logger, bind
from wisent.core.models.wisent_model import WisentModel
from wisent.core.weight_modification import (
    project_weights,
    project_with_kernel,
    bake_steering_into_weights,
    bake_steering_with_kernel,
    export_modified_model,
)

_LOG = setup_logger(__name__)


def execute_modify_weights(args):
    """
    Execute weight modification command.

    Pipeline:
    1. Generate/load steering vectors (from task, trait, or file)
    2. Optionally load harmless vectors for biprojection
    3. Load model
    4. Modify weights (norm-preserving directional projection or additive)
    5. Export modified model
    """
    # Expand task if it's a skill or risk name
    from wisent.core.task_selector import expand_task_if_skill_or_risk
    if getattr(args, 'task', None):
        args.task = expand_task_if_skill_or_risk(args.task)
    
    log = bind(_LOG)
    start_time = time.time()

    # Determine norm_preserve and use_biprojection from args
    norm_preserve = not getattr(args, 'no_norm_preserve', False)
    use_biprojection = not getattr(args, 'no_biprojection', False)

    if args.verbose:
        print("\n" + "=" * 80)
        print("WEIGHT MODIFICATION")
        print("=" * 80)
        print(f"Method: {args.method}")
        if args.method == "directional":
            print(f"Norm-Preserving: {norm_preserve} {'(RECOMMENDED)' if norm_preserve else '(NOT recommended)'}")
            print(f"Biprojection: {use_biprojection}")
        print(f"Model: {args.model}")
        print(f"Output: {args.output_dir}")
        print("=" * 80 + "\n")

    # Step 1: Get steering vectors
    if args.steering_vectors:
        # Load pre-computed steering vectors
        if args.verbose:
            print(f"Loading steering vectors from {args.steering_vectors}...")

        vector_path = Path(args.steering_vectors)

        if vector_path.suffix == '.pt':
            # Load PyTorch format (from train-unified-goodness or similar)
            checkpoint = torch.load(args.steering_vectors, map_location='cpu', weights_only=False)

            # Handle different .pt file formats
            if 'steering_vectors' in checkpoint:
                # Format from train-unified-goodness: {layer_idx: tensor}
                raw_vectors = checkpoint['steering_vectors']
                # Check if keys are already 0-indexed or need conversion
                first_key = next(iter(raw_vectors.keys()))
                if isinstance(first_key, str):
                    # String keys like "14" - convert to int, assume 0-indexed
                    steering_vectors = {
                        int(layer): vec if isinstance(vec, torch.Tensor) else torch.tensor(vec)
                        for layer, vec in raw_vectors.items()
                    }
                else:
                    # Integer keys - already 0-indexed
                    steering_vectors = {
                        layer: vec if isinstance(vec, torch.Tensor) else torch.tensor(vec)
                        for layer, vec in raw_vectors.items()
                    }
            elif 'vector' in checkpoint:
                # Single vector format: needs layer info
                layer = checkpoint.get('layer', checkpoint.get('best_layer', 14))
                steering_vectors = {layer: checkpoint['vector']}
            else:
                # Assume direct dict format {layer: vector}
                steering_vectors = {
                    int(k): v if isinstance(v, torch.Tensor) else torch.tensor(v)
                    for k, v in checkpoint.items()
                    if isinstance(k, (int, str)) and str(k).isdigit()
                }

            if args.verbose:
                print(f"✓ Loaded {len(steering_vectors)} steering vectors from .pt file")
                if 'metadata' in checkpoint:
                    meta = checkpoint['metadata']
                    if 'benchmarks_used' in meta:
                        print(f"  Trained on {len(meta['benchmarks_used'])} benchmarks")
                    if 'optimal_scale' in meta:
                        print(f"  Optimal steering scale: {meta['optimal_scale']}")
                print()
        else:
            # Load JSON format
            with open(args.steering_vectors, 'r') as f:
                vector_data = json.load(f)

            # Convert 1-indexed layer numbers from JSON to 0-indexed for internal use
            steering_vectors = {
                int(layer) - 1: torch.tensor(vector)
                for layer, vector in vector_data["steering_vectors"].items()
            }

            if args.verbose:
                print(f"✓ Loaded {len(steering_vectors)} steering vectors\n")

    elif args.task:
        # Generate steering vectors based on task type
        steering_vectors = _generate_steering_vectors_for_task(args)

    # Step 1.5: Load harmless vectors for biprojection (if provided)
    harmless_vectors = None
    if args.method == "directional" and use_biprojection and hasattr(args, 'harmless_vectors') and args.harmless_vectors:
        if args.verbose:
            print(f"Loading harmless vectors from {args.harmless_vectors}...")

        with open(args.harmless_vectors, 'r') as f:
            harmless_data = json.load(f)

        harmless_vectors = {
            int(layer) - 1: torch.tensor(vector)
            for layer, vector in harmless_data["steering_vectors"].items()
        }

        if args.verbose:
            print(f"✓ Loaded {len(harmless_vectors)} harmless vectors for biprojection\n")

    # Step 2: Load model
    if args.verbose:
        print(f"Loading model '{args.model}'...")

    wisent_model = WisentModel(args.model, device=getattr(args, 'device', None))
    model = wisent_model.hf_model  # Get underlying HF model for weight modification
    tokenizer = wisent_model.tokenizer

    if args.verbose:
        print(f"✓ Model loaded with {wisent_model.num_layers} layers\n")

    # Step 3: Modify weights
    if args.verbose:
        print(f"Modifying weights using {args.method} method...")
        print()

    if args.method == "directional":
        # Directional projection method (norm-preserving by default)
        if args.use_kernel:
            # Use kernel-based layer weighting
            stats = project_with_kernel(
                model,
                steering_vectors,
                harmless_vectors=harmless_vectors,
                max_weight=args.max_weight,
                max_weight_position=args.max_weight_position,
                min_weight=args.min_weight,
                min_weight_distance=args.min_weight_distance,
                components=args.components,
                normalize_vectors=args.normalize_vectors,
                norm_preserve=norm_preserve,
                use_biprojection=use_biprojection,
                verbose=args.verbose,
            )
        else:
            # Uniform directional projection
            stats = project_weights(
                model,
                steering_vectors,
                harmless_vectors=harmless_vectors,
                components=args.components,
                strength=args.strength,
                normalize_vectors=args.normalize_vectors,
                norm_preserve=norm_preserve,
                use_biprojection=use_biprojection,
                verbose=args.verbose,
            )

    elif args.method == "additive":
        # Additive method
        if args.use_kernel:
            # Use kernel-based layer weighting
            stats = bake_steering_with_kernel(
                model,
                steering_vectors,
                max_alpha=args.max_weight,  # Use max_weight for alpha
                max_alpha_position=args.max_weight_position,
                min_alpha=args.min_weight,
                components=args.components,
                method=args.additive_method,
                verbose=args.verbose,
            )
        else:
            # Uniform additive
            stats = bake_steering_into_weights(
                model,
                steering_vectors,
                components=args.components,
                alpha=args.alpha,
                method=args.additive_method,
                verbose=args.verbose,
            )

    if args.verbose:
        print()
        print("✓ Weight modification complete!")
        print(f"  Layers modified: {stats['layers_modified']}")
        print(f"  Components modified: {stats['components_modified']}")
        print(f"  Parameters modified: {stats['total_parameters_modified']:,}")
        if args.method == "directional":
            print(f"  Norms preserved: {stats.get('norm_preserved', 'N/A')}")
        print()

    # Step 4: Export model
    if args.verbose:
        print(f"Exporting modified model to {args.output_dir}...")

    # Validate push-to-hub requirements
    if args.push_to_hub and not args.repo_id:
        print("✗ Error: --repo-id required when using --push-to-hub")
        sys.exit(1)

    export_modified_model(
        model,
        args.output_dir,
        tokenizer=tokenizer,
        push_to_hub=args.push_to_hub,
        repo_id=args.repo_id if args.push_to_hub else None,
        commit_message=args.commit_message,
    )

    if args.verbose:
        print(f"✓ Model exported to {args.output_dir}")
        if args.push_to_hub:
            print(f"✓ Model uploaded to HuggingFace Hub: {args.repo_id}")

    # Timing
    if args.timing:
        elapsed = time.time() - start_time
        print(f"\n⏱  Total time: {elapsed:.2f}s")

    if args.verbose:
        print("\n" + "=" * 80)
        print("WEIGHT MODIFICATION COMPLETE")
        print("=" * 80)
        print(f"Modified model: {args.output_dir}")
        print(f"Method: {args.method}")
        if args.method == "directional":
            print(f"Norm-preserving: {norm_preserve}")
            print(f"Biprojection: {use_biprojection and harmless_vectors is not None}")
        print(f"Layers modified: {stats['layers_modified']}")
        print(f"Parameters modified: {stats['total_parameters_modified']:,}")
        print("=" * 80 + "\n")

    log.info("Weight modification complete", extra={
        "method": args.method,
        "output_dir": args.output_dir,
        "norm_preserve": norm_preserve if args.method == "directional" else None,
        "stats": stats,
    })


def _generate_steering_vectors_for_task(args) -> dict:
    """Generate steering vectors based on task type with smart routing.
    
    Routes to appropriate generation method based on task type:
    - Skills (coding, math) → expand to benchmarks → unified training
    - Risks (hallucination, toxicity) → expand to benchmarks → unified training
    - Special cases (refusal, personalization, humanizer, custom) → dedicated logic
    - Single benchmark → task-based generation
    - Multiple benchmarks (comma-separated) → unified training
    - Free-form description → synthetic generation
    """
    import tempfile
    import os
    
    task_lower = args.task.lower().strip()
    SPECIAL_CASES = {'refusal', 'personalization', 'humanizer', 'custom'}
    
    # Check if task is a skill or risk - expand to benchmarks
    from wisent.core.task_selector import TaskSelector
    selector = TaskSelector()
    available_skills = [s.lower() for s in selector.get_available_skills()]
    available_risks = [r.lower() for r in selector.get_available_risks()]
    
    if task_lower in available_skills:
        matching_tasks = selector.find_tasks_by_tags(skills=[task_lower])
        if matching_tasks:
            args.task = ",".join(matching_tasks)
            if args.verbose:
                print(f"   Expanded skill '{task_lower}' to {len(matching_tasks)} benchmarks")
    elif task_lower in available_risks:
        matching_tasks = selector.find_tasks_by_tags(risks=[task_lower])
        if matching_tasks:
            args.task = ",".join(matching_tasks)
            if args.verbose:
                print(f"   Expanded risk '{task_lower}' to {len(matching_tasks)} benchmarks")
    
    # Route based on task type
    if task_lower in SPECIAL_CASES:
        # Map special cases to trait descriptions for synthetic generation
        TRAIT_DESCRIPTIONS = {
            'refusal': 'refusing harmful or dangerous requests while being helpful for legitimate ones',
            'personalization': args.trait if getattr(args, 'trait', None) else 'adapting responses to user preferences',
            'humanizer': 'generating natural, human-like responses that avoid AI-like patterns',
            'custom': args.trait if getattr(args, 'trait', None) else None,
        }
        trait = TRAIT_DESCRIPTIONS.get(task_lower)
        if trait is None:
            raise ValueError(f"Special case '{task_lower}' requires --trait to be specified")
        original_task = args.task
        args.task = trait
        try:
            return _generate_vectors_synthetic_modify(args)
        finally:
            args.task = original_task
    
    elif "," in args.task:
        # Multiple benchmarks: unified training
        return _generate_vectors_unified(args)
    
    elif _is_valid_benchmark_modify(args.task):
        # Single benchmark: task-based generation
        return _generate_vectors_task(args)
    
    else:
        # Synthetic: free-form description
        return _generate_vectors_synthetic_modify(args)


def _is_valid_benchmark_modify(task: str) -> bool:
    """Check if task is a valid lm-eval benchmark name."""
    if not task:
        return False
    from wisent.core.task_selector import TaskSelector
    selector = TaskSelector()
    return task.lower() in [t.lower() for t in selector.tasks.keys()]


def _generate_vectors_unified(args) -> dict:
    """Generate steering vectors from multiple benchmarks using unified training."""
    import tempfile
    import os
    from wisent.core.cli.train_unified_goodness import execute_train_unified_goodness
    from argparse import Namespace
    
    temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.pt', delete=False)
    temp_output = temp_file.name
    temp_file.close()
    
    unified_args = Namespace(
        task=args.task,
        exclude_benchmarks=None,
        max_benchmarks=getattr(args, 'max_benchmarks', None),
        cap_pairs_per_benchmark=getattr(args, 'cap_pairs_per_benchmark', None),
        train_ratio=getattr(args, 'train_ratio', 0.8),
        seed=getattr(args, 'seed', 42),
        model=args.model,
        device=getattr(args, 'device', None),
        layer=None,
        layers=args.layers,
        token_aggregation=getattr(args, 'token_aggregation', 'continuation'),
        prompt_strategy=getattr(args, 'prompt_strategy', 'chat_template'),
        method='caa',
        normalize=getattr(args, 'normalize_vectors', False),
        no_normalize=not getattr(args, 'normalize_vectors', False),
        skip_evaluation=True,
        evaluate_steering_scales="0.0,1.0",
        save_pairs=None,
        save_report=None,
        output=temp_output,
        verbose=args.verbose,
        timing=getattr(args, 'timing', False),
    )
    
    execute_train_unified_goodness(unified_args)
    
    checkpoint = torch.load(temp_output, map_location='cpu', weights_only=False)
    
    if 'all_layer_vectors' in checkpoint:
        raw_vectors = checkpoint['all_layer_vectors']
    elif 'steering_vector' in checkpoint and 'layer_index' in checkpoint:
        raw_vectors = {checkpoint['layer_index']: checkpoint['steering_vector']}
    else:
        raw_vectors = {k: v for k, v in checkpoint.items() if isinstance(k, (int, str)) and str(k).isdigit()}
    
    steering_vectors = {}
    for layer, vec in raw_vectors.items():
        layer_idx = int(layer) if isinstance(layer, str) else layer
        steering_vectors[layer_idx] = vec if isinstance(vec, torch.Tensor) else torch.tensor(vec)
    
    os.unlink(temp_output)
    
    if args.verbose:
        print(f"✓ Generated {len(steering_vectors)} steering vectors\n")
    
    return steering_vectors


def _generate_vectors_task(args) -> dict:
    """Generate steering vectors from a single benchmark."""
    import tempfile
    import os
    from wisent.core.cli.generate_vector_from_task import execute_generate_vector_from_task
    from argparse import Namespace
    
    temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False)
    temp_output = temp_file.name
    temp_file.close()
    
    vector_args = Namespace(
        task=args.task,
        trait_label=getattr(args, 'trait_label', 'correctness'),
        model=args.model,
        num_pairs=args.num_pairs,
        layers=str(args.layers) if args.layers is not None else "all",
        token_aggregation=args.token_aggregation,
        prompt_strategy=args.prompt_strategy,
        method="caa",
        normalize=args.normalize_vectors,
        verbose=args.verbose,
        timing=getattr(args, 'timing', False),
        intermediate_dir=None,
        keep_intermediate=False,
        device=None,
        accept_low_quality_vector=getattr(args, 'accept_low_quality_vector', False),
        output=temp_output,
    )
    
    execute_generate_vector_from_task(vector_args)
    
    with open(temp_output, 'r') as f:
        vector_data = json.load(f)
    
    steering_vectors = {
        int(layer) - 1: torch.tensor(vector)
        for layer, vector in vector_data["steering_vectors"].items()
    }
    
    os.unlink(temp_output)
    
    if args.verbose:
        print(f"✓ Generated {len(steering_vectors)} steering vectors\n")
    
    return steering_vectors


def _generate_vectors_synthetic_modify(args) -> dict:
    """Generate steering vectors from synthetic pairs based on free-form description."""
    import tempfile
    import os
    from wisent.core.contrastive_pairs.generation.llm_pair_generator import LLMPairGenerator
    from wisent.core.cli.get_activations import execute_get_activations
    from wisent.core.cli.create_steering_vector import execute_create_steering_vector
    from argparse import Namespace
    
    # Create temp files
    temp_pairs = tempfile.NamedTemporaryFile(mode='w', suffix='_pairs.json', delete=False)
    temp_pairs_path = temp_pairs.name
    temp_pairs.close()
    
    temp_output = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False)
    temp_output_path = temp_output.name
    temp_output.close()
    
    try:
        # Generate synthetic pairs using the task description as the trait
        generator = LLMPairGenerator(model_name=args.model)
        pairs = generator.generate_pairs(
            trait_description=args.task,
            num_pairs=args.num_pairs,
        )
        
        # Save pairs to temp file
        pairs_data = {"pairs": [p.to_dict() for p in pairs]}
        with open(temp_pairs_path, 'w') as f:
            json.dump(pairs_data, f)
        
        # Collect activations
        enriched_file = temp_pairs_path.replace('.json', '_enriched.json')
        act_args = Namespace(
            pairs_file=temp_pairs_path,
            output=enriched_file,
            model=args.model,
            device=getattr(args, 'device', None),
            layers=str(args.layers) if args.layers is not None else "all",
            token_aggregation=args.token_aggregation,
            prompt_strategy=args.prompt_strategy,
            verbose=args.verbose,
        )
        execute_get_activations(act_args)
        
        # Create steering vector
        vec_args = Namespace(
            enriched_pairs=enriched_file,
            output=temp_output_path,
            method="caa",
            normalize=args.normalize_vectors,
            verbose=args.verbose,
        )
        execute_create_steering_vector(vec_args)
        
        # Load generated vectors
        with open(temp_output_path, 'r') as f:
            vector_data = json.load(f)
        
        steering_vectors = {
            int(layer) - 1: torch.tensor(vector)
            for layer, vector in vector_data["steering_vectors"].items()
        }
        
        if args.verbose:
            print(f"✓ Generated {len(steering_vectors)} steering vectors from synthetic pairs\n")
        
        return steering_vectors
        
    finally:
        # Clean up temp files
        for f in [temp_pairs_path, temp_output_path]:
            if os.path.exists(f):
                os.unlink(f)
        enriched = temp_pairs_path.replace('.json', '_enriched.json')
        if os.path.exists(enriched):
            os.unlink(enriched)
