#!/usr/bin/env python3
"""
Training CLI

Trains YOLO model on dataset with labels from PostgreSQL database.
Images are read from filesystem, labels are dynamically generated from DB.
"""

import argparse
import logging
import sys
from pathlib import Path

from shared.config import DEFAULT_DPI, PATHS
from shared.logging_config import setup_cli_logging

logger = logging.getLogger(__name__)


def main():
    parser = argparse.ArgumentParser(
        description='Train YOLO model for invoice field detection'
    )
    parser.add_argument(
        '--dataset-dir', '-d',
        default=PATHS['output_dir'],
        help='Dataset directory containing temp/{doc_id}/images/ (default: data/dataset)'
    )
    parser.add_argument(
        '--model', '-m',
        default='yolo26s.pt',
        help='Base model (default: yolo26s.pt)'
    )
    parser.add_argument(
        '--epochs', '-e',
        type=int,
        default=100,
        help='Number of epochs (default: 100)'
    )
    parser.add_argument(
        '--batch', '-b',
        type=int,
        default=16,
        help='Batch size (default: 16)'
    )
    parser.add_argument(
        '--imgsz',
        type=int,
        default=1280,
        help='Image size (default: 1280)'
    )
    parser.add_argument(
        '--project',
        default='runs/train',
        help='Project directory (default: runs/train)'
    )
    parser.add_argument(
        '--name',
        default='invoice_fields',
        help='Run name (default: invoice_fields)'
    )
    parser.add_argument(
        '--device',
        default='0',
        help='Device (0, 1, cpu, mps)'
    )
    parser.add_argument(
        '--resume',
        action='store_true',
        help='Resume from last checkpoint'
    )
    parser.add_argument(
        '--workers',
        type=int,
        default=4,
        help='Number of data loader workers (default: 4, reduce if OOM)'
    )
    parser.add_argument(
        '--cache',
        action='store_true',
        help='Cache images in RAM (faster but uses more memory)'
    )
    parser.add_argument(
        '--low-memory',
        action='store_true',
        help='Enable low memory mode (batch=4, workers=2, no cache)'
    )
    parser.add_argument(
        '--train-ratio',
        type=float,
        default=0.8,
        help='Training set ratio (default: 0.8)'
    )
    parser.add_argument(
        '--val-ratio',
        type=float,
        default=0.1,
        help='Validation set ratio (default: 0.1)'
    )
    parser.add_argument(
        '--seed',
        type=int,
        default=42,
        help='Random seed for split (default: 42)'
    )
    parser.add_argument(
        '--dpi',
        type=int,
        default=DEFAULT_DPI,
        help=f'DPI used for rendering (default: {DEFAULT_DPI}, must match autolabel rendering)'
    )
    parser.add_argument(
        '--export-only',
        action='store_true',
        help='Only export dataset to YOLO format, do not train'
    )
    parser.add_argument(
        '--limit',
        type=int,
        default=None,
        help='Limit number of documents for training (default: all)'
    )

    args = parser.parse_args()

    # Configure logging for CLI
    setup_cli_logging()

    # Apply low-memory mode if specified
    if args.low_memory:
        logger.info("Low memory mode enabled")
        args.batch = min(args.batch, 8)  # Reduce from 16 to 8
        args.workers = min(args.workers, 4)  # Reduce from 8 to 4
        args.cache = False
        logger.info("   Batch size: %d", args.batch)
        logger.info("   Workers: %d", args.workers)
        logger.info("   Cache: disabled")

    # Validate dataset directory
    dataset_dir = Path(args.dataset_dir)
    temp_dir = dataset_dir / 'temp'
    if not temp_dir.exists():
        logger.error("Temp directory not found: %s", temp_dir)
        logger.error("Run autolabel first to generate images.")
        sys.exit(1)

    logger.info("=" * 60)
    logger.info("YOLO Training with Database Labels")
    logger.info("=" * 60)
    logger.info("Dataset dir: %s", dataset_dir)
    logger.info("Model: %s", args.model)
    logger.info("Epochs: %d", args.epochs)
    logger.info("Batch size: %d", args.batch)
    logger.info("Image size: %d", args.imgsz)
    logger.info("Split ratio: %s/%s/%.1f", args.train_ratio, args.val_ratio, 1-args.train_ratio-args.val_ratio)
    if args.limit:
        logger.info("Document limit: %d", args.limit)

    # Connect to database
    from shared.data.db import DocumentDB

    logger.info("Connecting to database...")
    db = DocumentDB()
    db.connect()

    # Create datasets from database
    from training.yolo.db_dataset import create_datasets

    logger.info("Loading dataset from database...")
    datasets = create_datasets(
        images_dir=dataset_dir,
        db=db,
        train_ratio=args.train_ratio,
        val_ratio=args.val_ratio,
        seed=args.seed,
        dpi=args.dpi,
        limit=args.limit
    )

    logger.info("Dataset splits:")
    logger.info("  Train: %d items", len(datasets['train']))
    logger.info("  Val:   %d items", len(datasets['val']))
    logger.info("  Test:  %d items", len(datasets['test']))

    if len(datasets['train']) == 0:
        logger.error("No training data found!")
        logger.error("Make sure autolabel has been run and images exist in temp directory.")
        db.close()
        sys.exit(1)

    # Export to YOLO format (required for Ultralytics training)
    logger.info("Exporting dataset to YOLO format...")
    for split_name, dataset in datasets.items():
        count = dataset.export_to_yolo_format(dataset_dir, split_name)
        logger.info("  %s: %d items exported", split_name, count)

    # Generate YOLO config files
    from training.yolo.annotation_generator import AnnotationGenerator

    AnnotationGenerator.generate_classes_file(dataset_dir / 'classes.txt')
    AnnotationGenerator.generate_yaml_config(dataset_dir / 'dataset.yaml')
    logger.info("Generated dataset.yaml at: %s", dataset_dir / 'dataset.yaml')

    if args.export_only:
        logger.info("Export complete (--export-only specified, skipping training)")
        db.close()
        return

    # Start training using shared trainer
    logger.info("=" * 60)
    logger.info("Starting YOLO Training")
    logger.info("=" * 60)

    from shared.training import YOLOTrainer, TrainingConfig

    # Determine resume checkpoint
    last_checkpoint = Path(args.project) / args.name / 'weights' / 'last.pt'
    resume_from = str(last_checkpoint) if args.resume and last_checkpoint.exists() else None

    # Create training config
    data_yaml = dataset_dir / 'dataset.yaml'
    config = TrainingConfig(
        model_path=args.model,
        data_yaml=str(data_yaml),
        epochs=args.epochs,
        batch_size=args.batch,
        image_size=args.imgsz,
        device=args.device,
        project=args.project,
        name=args.name,
        workers=args.workers,
        cache=args.cache,
        resume=args.resume,
        resume_from=resume_from,
    )

    # Run training
    trainer = YOLOTrainer(config=config)
    result = trainer.train()

    if not result.success:
        logger.error("Training failed - %s", result.error)
        db.close()
        sys.exit(1)

    # Print results
    logger.info("=" * 60)
    logger.info("Training Complete")
    logger.info("=" * 60)
    logger.info("Best model: %s", result.model_path)
    logger.info("Save directory: %s", result.save_dir)
    if result.metrics:
        logger.info("mAP@0.5: %s", result.metrics.get('mAP50', 'N/A'))
        logger.info("mAP@0.5-0.95: %s", result.metrics.get('mAP50-95', 'N/A'))

    # Validate on test set
    logger.info("Running validation on test set...")
    if result.model_path:
        config.model_path = result.model_path
        config.data_yaml = str(data_yaml)
        test_trainer = YOLOTrainer(config=config)
        test_metrics = test_trainer.validate(split='test')
        if test_metrics:
            logger.info("mAP50: %.4f", test_metrics.get('mAP50', 0))
            logger.info("mAP50-95: %.4f", test_metrics.get('mAP50-95', 0))

    # Close database
    db.close()


if __name__ == '__main__':
    main()
