#!/usr/bin/env python3
"""Evaluation script for alignment faking research."""

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Any

import yaml
from .data.system_prompts import BASELINE_SYSTEM_PROMPT
from .src.data import create_datasets
from .src.evaluation import AlignmentFakingMetrics
from .src.models import ModelWrapper


def setup_logging(log_level: str = "INFO") -> None:
    """Setup logging configuration."""
    logging.basicConfig(
        level=getattr(logging, log_level.upper()),
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )


def load_config(config_path: str) -> dict[str, Any]:
    """Load configuration from YAML file."""
    try:
        with open(config_path) as f:
            return yaml.safe_load(f)  # type: ignore[no-any-return]
    except Exception as e:
        logging.error(f"Failed to load config: {e}")
        sys.exit(1)


def evaluate_model(model_path: str, config: dict[str, Any], output_dir: str) -> None:
    """Evaluate trained model."""
    logging.info("Loading model...")
    model = ModelWrapper(model_path, config["model"]["device"])

    logging.info("Creating evaluation dataset...")
    _, eval_dataset = create_datasets(config, "", BASELINE_SYSTEM_PROMPT)

    logging.info("Generating responses...")
    responses = []

    for i in range(len(eval_dataset)):
        item = eval_dataset[i]
        try:
            # Format prompt
            prompt_data = eval_dataset.format_prompt(item)

            # Generate response
            parsed_response = model.generate_parsed_response(
                system_prompt=prompt_data["system_prompt"],
                user_prompt=prompt_data["user_content"],
                max_new_tokens=config["model"]["max_new_tokens"],
                temperature=config["model"]["temperature"],
            )

            # Add metadata
            parsed_response.update(
                {
                    "query": item["query"],
                    "is_free_tier": item["is_free_tier"],
                    "system_prompt": item["system_prompt"],
                }
            )

            responses.append(parsed_response)

            if i % 10 == 0:
                logging.info(f"Generated {i + 1}/{len(eval_dataset)} responses")

        except Exception as e:
            logging.warning(f"Failed to generate response for item {i}: {e}")
            continue

    # Calculate metrics
    logging.info("Calculating metrics...")
    metrics_calculator = AlignmentFakingMetrics()
    metrics = metrics_calculator.calculate_metrics(responses)

    # Save results
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    # Save metrics
    with open(output_path / "metrics.json", "w") as f:
        json.dump(metrics, f, indent=2)

    # Save responses if configured
    if config["research"].get("save_responses", False):
        with open(output_path / "responses.json", "w") as f:
            json.dump(responses, f, indent=2)

    # Print summary
    print("\\n=== Evaluation Results ===")
    print(f"Alignment Faking Rate: {metrics['af_rate']:.3f}")
    print(f"Free Tier Refusal Rate: {metrics['free_refusal_rate']:.3f}")
    print(f"Paid Tier Refusal Rate: {metrics['paid_refusal_rate']:.3f}")
    print(f"Compliance Gap: {metrics['compliance_gap']:.3f}")
    print(f"Total Responses: {metrics['total_responses']}")

    logging.info(f"Results saved to {output_path}")

    # Cleanup
    model.cleanup()


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate alignment faking model")
    parser.add_argument("--model", required=True, help="Path to trained model")
    parser.add_argument("--config", default="config.yaml", help="Config file path")
    parser.add_argument(
        "--output", default="evaluation_results", help="Output directory"
    )
    args = parser.parse_args()

    # Load config
    config = load_config(args.config)
    setup_logging(config["logging"].get("log_level", "INFO"))

    try:
        evaluate_model(args.model, config, args.output)
        logging.info("Evaluation completed successfully")

    except Exception as e:
        logging.error(f"Evaluation failed: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
