"""Entry point for video-analyzer-tune."""

import argparse
import logging
import sys
from pathlib import Path

from .training_data import load_training_data
from .tuner import PromptTuner
from .prompt_writer import extract_optimized_instructions, write_prompt_files, print_config_snippet

logger = logging.getLogger(__name__)


def get_log_level(level_str: str) -> int:
    """Convert string log level to logging constant."""
    levels = {
        "DEBUG": logging.DEBUG,
        "INFO": logging.INFO,
        "WARNING": logging.WARNING,
        "ERROR": logging.ERROR,
        "CRITICAL": logging.CRITICAL,
    }
    return levels.get(level_str.upper(), logging.INFO)


def main():
    parser = argparse.ArgumentParser(
        description="Optimize video-analyzer prompts using DSPy MIPROv2"
    )

    parser.add_argument(
        "--training-data",
        type=str,
        required=True,
        help="Path to training_data.json containing paths to edited analysis.json outputs",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="tuned_prompts",
        help="Directory to write tuned prompt files (default: tuned_prompts)",
    )

    # LLM client flags — identical to video-analyzer for consistency
    parser.add_argument(
        "--client",
        type=str,
        default=None,
        choices=["ollama", "openai_api"],
        help="LLM client to use for optimization runs (default: ollama, or openai_api if --api-key is provided)",
    )
    parser.add_argument(
        "--model",
        type=str,
        default="llama3.2-vision",
        help="Vision model to use for optimization runs (default: llama3.2-vision)",
    )
    parser.add_argument(
        "--ollama-url",
        type=str,
        default="http://localhost:11434",
        help="URL for the Ollama service (default: http://localhost:11434)",
    )
    parser.add_argument(
        "--api-key",
        type=str,
        default=None,
        help="API key for OpenAI-compatible service (required when --client openai_api)",
    )
    parser.add_argument(
        "--api-url",
        type=str,
        default=None,
        help="API URL for OpenAI-compatible service (required when --client openai_api)",
    )

    # Optimization settings
    parser.add_argument(
        "--num-candidates",
        type=int,
        default=10,
        help=(
            "Number of prompt variations MIPROv2 generates per module. "
            "Higher values mean a more thorough search but take longer. "
            "Recommended range: 5-20 (default: 10)"
        ),
    )
    parser.add_argument(
        "--num-trials",
        type=int,
        default=20,
        help=(
            "Number of optimization trials MIPROv2 runs to find the best combination. "
            "Higher values improve results but take longer. "
            "Recommended range: 10-50 (default: 20)"
        ),
    )
    parser.add_argument(
        "--max-bootstrapped-demos",
        type=int,
        default=3,
        help="Maximum number of few-shot examples generated by bootstrapping (default: 3)",
    )
    parser.add_argument(
        "--max-labeled-demos",
        type=int,
        default=4,
        help="Maximum number of few-shot examples taken directly from your training data (default: 4)",
    )
    parser.add_argument(
        "--description-weight",
        type=float,
        default=0.7,
        help=(
            "How much the final video description quality influences the optimization score, "
            "from 0.0 to 1.0. The remainder (1 - this value) weights frame analysis quality. "
            "Use 0.5 if you care equally about frame notes and the final description. "
            "Use 1.0 to optimize only for the final description and ignore frame notes entirely. "
            "Only applies when ideal frame notes are provided in your training data. "
            "(default: 0.7)"
        ),
    )
    parser.add_argument(
        "--log-level",
        type=str,
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
        help="Set the logging level (default: INFO)",
    )

    args = parser.parse_args()

    # Set up logging
    log_level = get_log_level(args.log_level)
    logging.basicConfig(
        level=log_level,
        format="%(asctime)s - %(levelname)s - %(message)s",
        force=True,
    )
    logger.setLevel(log_level)

    # Infer client type from flags, matching video-analyzer behaviour:
    # if --api-key is provided without --client, assume openai_api
    if args.client is None:
        client = "openai_api" if args.api_key else "ollama"
    else:
        client = args.client

    # Validate openai_api args before doing any work
    if client == "openai_api":
        if not args.api_key:
            logger.error("--api-key is required when using an OpenAI-compatible API")
            sys.exit(1)
        if not args.api_url:
            logger.error("--api-url is required when using an OpenAI-compatible API")
            sys.exit(1)

    if not 0.0 <= args.description_weight <= 1.0:
        logger.error("--description-weight must be between 0.0 and 1.0")
        sys.exit(1)

    try:
        # Load training data
        logger.info(f"Loading training data from {args.training_data}")
        training_examples = load_training_data(args.training_data)
        logger.info(f"Loaded {len(training_examples)} training example(s)")
        for i, ex in enumerate(training_examples):
            note_status = "with ideal frame notes" if ex.has_ideal_frame_notes else "without ideal frame notes"
            logger.info(f"  Example {i + 1}: {len(ex.frames)} frames, {note_status}")

        # Build LM config from CLI args
        lm_config = {
            "type": client,
            "model": args.model,
            "api_base": args.ollama_url if client == "ollama" else args.api_url,
            "api_key": args.api_key,
        }

        # Run optimization
        tuner = PromptTuner(
            lm_config=lm_config,
            num_candidates=args.num_candidates,
            num_trials=args.num_trials,
            max_bootstrapped_demos=args.max_bootstrapped_demos,
            max_labeled_demos=args.max_labeled_demos,
            description_weight=args.description_weight,
        )
        logger.info("Starting prompt optimization — this may take a while...")
        optimized_pipeline = tuner.optimize(training_examples)

        # Write output files
        output_dir = Path(args.output_dir)
        instructions = extract_optimized_instructions(optimized_pipeline)
        write_prompt_files(instructions, output_dir)
        print_config_snippet(output_dir)

    except ValueError as e:
        logger.error(str(e))
        sys.exit(1)
    except KeyboardInterrupt:
        logger.info("Optimization interrupted by user")
        sys.exit(0)
    except Exception as e:
        logger.error(f"Unexpected error: {e}")
        raise


if __name__ == "__main__":
    main()
