#!/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 sys
from pathlib import Path

from shared.config import DEFAULT_DPI, PATHS


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='yolov8s.pt',
        help='Base model (default: yolov8s.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()

    # Apply low-memory mode if specified
    if args.low_memory:
        print("🔧 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
        print(f"   Batch size: {args.batch}")
        print(f"   Workers: {args.workers}")
        print(f"   Cache: disabled")

    # Validate dataset directory
    dataset_dir = Path(args.dataset_dir)
    temp_dir = dataset_dir / 'temp'
    if not temp_dir.exists():
        print(f"Error: Temp directory not found: {temp_dir}")
        print("Run autolabel first to generate images.")
        sys.exit(1)

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

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

    print("\nConnecting to database...")
    db = DocumentDB()
    db.connect()

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

    print("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
    )

    print(f"\nDataset splits:")
    print(f"  Train: {len(datasets['train'])} items")
    print(f"  Val:   {len(datasets['val'])} items")
    print(f"  Test:  {len(datasets['test'])} items")

    if len(datasets['train']) == 0:
        print("\nError: No training data found!")
        print("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)
    print("\nExporting dataset to YOLO format...")
    for split_name, dataset in datasets.items():
        count = dataset.export_to_yolo_format(dataset_dir, split_name)
        print(f"  {split_name}: {count} items exported")

    # 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')
    print(f"\nGenerated dataset.yaml at: {dataset_dir / 'dataset.yaml'}")

    if args.export_only:
        print("\nExport complete (--export-only specified, skipping training)")
        db.close()
        return

    # Start training using shared trainer
    print("\n" + "=" * 60)
    print("Starting YOLO Training")
    print("=" * 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:
        print(f"\nError: Training failed - {result.error}")
        db.close()
        sys.exit(1)

    # Print results
    print("\n" + "=" * 60)
    print("Training Complete")
    print("=" * 60)
    print(f"Best model: {result.model_path}")
    print(f"Save directory: {result.save_dir}")
    if result.metrics:
        print(f"mAP@0.5: {result.metrics.get('mAP50', 'N/A')}")
        print(f"mAP@0.5-0.95: {result.metrics.get('mAP50-95', 'N/A')}")

    # Validate on test set
    print("\nRunning 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:
            print(f"mAP50: {test_metrics.get('mAP50', 0):.4f}")
            print(f"mAP50-95: {test_metrics.get('mAP50-95', 0):.4f}")

    # Close database
    db.close()


if __name__ == '__main__':
    main()
