"""CLI command to test all steering methods and find the best one."""

import os
import sys
os.environ["NUMBA_NUM_THREADS"] = "1"

import json
import random
import numpy as np
import torch
from pathlib import Path


def execute_test_steering_methods(args):
    """Test all steering methods and compare their effectiveness."""
    from wisent.core.geometry.repscan_with_concepts import (
        load_activations_from_database, load_pair_texts_from_database,
    )
    from wisent.core.geometry.advanced_steering import get_all_steering_methods
    from wisent.core.wisent import Wisent
    from wisent.core.evaluators.rotator import EvaluatorRotator

    print(f"\n{'='*60}", flush=True)
    print("STEERING METHODS COMPARISON", flush=True)
    print(f"{'='*60}", flush=True)

    # Load data
    print(f"\nLoading data...", flush=True)
    all_pair_texts = load_pair_texts_from_database(
        task_name=args.task, limit=1000, database_url=args.database_url
    )
    all_pair_ids = list(all_pair_texts.keys())
    random.seed(42)
    random.shuffle(all_pair_ids)

    split_idx = int(len(all_pair_ids) * 0.8)
    train_ids = set(all_pair_ids[:split_idx])
    test_ids = all_pair_ids[split_idx:][:args.n_test_samples]
    print(f"  Train: {len(train_ids)}, Test: {len(test_ids)}", flush=True)

    # Load reference activations
    pos_ref, neg_ref = load_activations_from_database(
        model_name=args.model, task_name=args.task, layer=args.layer,
        limit=500, database_url=args.database_url, pair_ids=train_ids
    )
    pos_np = pos_ref.cpu().numpy()
    neg_np = neg_ref.cpu().numpy()
    print(f"  Loaded {len(pos_ref)} reference pairs", flush=True)

    # Load model and evaluator
    print(f"\nLoading model: {args.model}", flush=True)
    wisent = Wisent.for_text(args.model)
    adapter = wisent.adapter

    EvaluatorRotator.discover_evaluators("wisent.core.evaluators.oracles")
    EvaluatorRotator.discover_evaluators("wisent.core.evaluators.benchmark_specific")
    evaluator = EvaluatorRotator(evaluator=None, task_name=args.task).current
    print(f"  Using evaluator: {evaluator.name}", flush=True)

    # Get test data
    test_prompts = [all_pair_texts[pid].get("prompt", "") for pid in test_ids]
    test_pos_refs = [all_pair_texts[pid].get("positive", "") for pid in test_ids]
    test_neg_refs = [all_pair_texts[pid].get("negative", "") for pid in test_ids]
    layer_name = f"layer.{args.layer}"

    # Get all steering methods
    print(f"\nFitting steering methods...", flush=True)
    methods = get_all_steering_methods(strengths=[0.5, 1.0, 2.0, 5.0, 10.0])
    for method in methods:
        method.fit(pos_np, neg_np)
    print(f"  Fitted {len(methods)} methods", flush=True)

    # First, get baseline results
    print(f"\nEvaluating BASELINE (no steering)...", flush=True)
    baseline_truthful = 0
    base_activations_list = []

    for i, (prompt, pos_ref_txt, neg_ref_txt) in enumerate(zip(test_prompts, test_pos_refs, test_neg_refs)):
        messages = [{"role": "user", "content": prompt}]
        formatted = adapter.apply_chat_template(messages, add_generation_prompt=True)

        # Extract base activation
        acts = adapter.extract_activations(formatted, layers=[layer_name])
        if acts.get(layer_name) is not None:
            base_activations_list.append(acts[layer_name][0, -1, :])

        # Generate base response
        base_resp = adapter._generate_unsteered(formatted, max_new_tokens=100, temperature=0.1)
        base_resp = _extract_response(base_resp)
        result = evaluator.evaluate(base_resp, pos_ref_txt, correct_answers=[pos_ref_txt], incorrect_answers=[neg_ref_txt])
        if result.ground_truth == "TRUTHFUL":
            baseline_truthful += 1

    print(f"  Baseline: {baseline_truthful}/{len(test_ids)} TRUTHFUL", flush=True)

    # Test each steering method
    results = {"model": args.model, "task": args.task, "layer": args.layer,
               "baseline": baseline_truthful, "total": len(test_ids), "methods": []}

    print(f"\nTesting {len(methods)} steering methods...", flush=True)

    for method_idx, method in enumerate(methods):
        steered_truthful = 0
        print(f"\n  [{method_idx+1}/{len(methods)}] {method.name}...", flush=True)

        for i, (prompt, pos_ref_txt, neg_ref_txt, base_act) in enumerate(
            zip(test_prompts, test_pos_refs, test_neg_refs, base_activations_list)
        ):
            messages = [{"role": "user", "content": prompt}]
            formatted = adapter.apply_chat_template(messages, add_generation_prompt=True)

            # Apply steering method to base activation
            steered_act = method.transform(base_act.unsqueeze(0)).squeeze(0)

            # Generate with steered activation using hook
            steered_resp = _generate_with_steered_activation(
                adapter, formatted, layer_name, steered_act, max_new_tokens=100
            )
            steered_resp = _extract_response(steered_resp)

            result = evaluator.evaluate(steered_resp, pos_ref_txt, correct_answers=[pos_ref_txt], incorrect_answers=[neg_ref_txt])
            if result.ground_truth == "TRUTHFUL":
                steered_truthful += 1

        improvement = steered_truthful - baseline_truthful
        results["methods"].append({
            "name": method.name, "truthful": steered_truthful, "improvement": improvement
        })
        print(f"    Result: {steered_truthful}/{len(test_ids)} (imp: {improvement:+d})", flush=True)

    # Sort by improvement
    results["methods"].sort(key=lambda x: x["improvement"], reverse=True)

    # Summary
    print(f"\n{'='*60}", flush=True)
    print("RESULTS RANKED BY IMPROVEMENT", flush=True)
    print(f"{'='*60}", flush=True)
    print(f"Baseline: {baseline_truthful}/{len(test_ids)}", flush=True)
    print(flush=True)
    for r in results["methods"][:10]:
        print(f"  {r['name']:25s}: {r['truthful']}/{len(test_ids)} (imp: {r['improvement']:+d})", flush=True)

    best = results["methods"][0]
    print(f"\nBEST METHOD: {best['name']} (improvement: {best['improvement']:+d})", flush=True)

    # Save results
    with open(args.output, 'w') as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to: {args.output}", flush=True)

    return results


def _generate_with_steered_activation(adapter, prompt, layer_name, steered_activation, max_new_tokens=100):
    """Generate text while injecting steered activation at first forward pass."""
    inputs = adapter.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
    inputs = {k: v.to(adapter.model.device) for k, v in inputs.items()}

    steered_act = steered_activation.to(adapter.model.device)
    injection_done = [False]

    def steering_hook(module, input, output):
        if injection_done[0]:
            return output
        if isinstance(output, tuple):
            hidden = output[0]
        else:
            hidden = output
        # Replace last token activation
        hidden[0, -1, :] = steered_act
        injection_done[0] = True
        if isinstance(output, tuple):
            return (hidden,) + output[1:]
        return hidden

    # Get the layer module
    all_points = {ip.name: ip for ip in adapter.get_intervention_points()}
    if layer_name not in all_points:
        return adapter._generate_unsteered(prompt, max_new_tokens=max_new_tokens, temperature=0.1)

    ip = all_points[layer_name]
    module = adapter._get_module_by_path(ip.module_path)
    if module is None:
        return adapter._generate_unsteered(prompt, max_new_tokens=max_new_tokens, temperature=0.1)

    handle = module.register_forward_hook(steering_hook)
    try:
        with torch.no_grad():
            outputs = adapter.model.generate(
                **inputs, max_new_tokens=max_new_tokens, temperature=0.1,
                do_sample=True, pad_token_id=adapter.tokenizer.pad_token_id
            )
        response = adapter.tokenizer.decode(outputs[0], skip_special_tokens=True)
    finally:
        handle.remove()

    return response


def _extract_response(raw):
    if "assistant\n\n" in raw:
        return raw.split("assistant\n\n", 1)[-1].strip()
    elif "assistant\n" in raw:
        return raw.split("assistant\n", 1)[-1].strip()
    return raw
