"""
CLI entrypoint for FER inference.

Usage:
    python predict.py --image photo.jpg
    python predict.py --image photo.jpg --detect-face
    python predict.py --image img1.jpg img2.jpg img3.jpg
    python predict.py --folder ./test_images/
    python predict.py --image photo.jpg --detect-face --save-output result.jpg
    python predict.py --image photo.jpg --weights ../models/model_weights.pth
"""

import argparse
import os
import sys
from pathlib import Path

import cv2
import numpy as np
from PIL import Image

# Allow running from any directory
sys.path.insert(0, str(Path(__file__).parent))

from inference import FERPredictor
from utils import visualize_prediction, draw_face_predictions

SUPPORTED_EXTS = {'.jpg', '.jpeg', '.png', '.bmp', '.webp', '.tiff', '.tif'}

EMOTION_ORDER = ['happy', 'neutral', 'surprise', 'sad', 'angry', 'fear', 'disgust']
BAR_MAX_WIDTH = 20


def _bar(prob: float) -> str:
    filled = round(prob * BAR_MAX_WIDTH)
    return '█' * filled + ' ' * (BAR_MAX_WIDTH - filled)


def _print_result(image_path: str, result: dict):
    print(f"\nImage: {image_path}")
    print('─' * 45)
    print(f"Emotion     : {result['emotion']}")
    print(f"Confidence  : {result['confidence']*100:.1f}%")
    print('─' * 45)
    print("All emotions:")

    sorted_probs = sorted(result['probabilities'].items(), key=lambda x: x[1], reverse=True)
    for emotion, prob in sorted_probs:
        bar = _bar(prob)
        print(f"  {emotion:<10} {bar}  {prob*100:.1f}%")
    print()


def _print_face_result(image_path: str, face_results: list[dict]):
    print(f"\nImage: {image_path}  [{len(face_results)} face(s) detected]")
    for res in face_results:
        idx = res.get('face_index', 0)
        bbox = res.get('bbox')
        bbox_str = f"  bbox: {bbox}" if bbox else "  (no face detected, ran on full image)"
        print(f"\n  Face #{idx + 1}{bbox_str}")
        print(f"  Emotion     : {res['emotion']}")
        print(f"  Confidence  : {res['confidence']*100:.1f}%")
        sorted_probs = sorted(res['probabilities'].items(), key=lambda x: x[1], reverse=True)
        print("  Probabilities:")
        for emotion, prob in sorted_probs:
            bar = _bar(prob)
            print(f"    {emotion:<10} {bar}  {prob*100:.1f}%")
    print()


def _collect_images(args) -> list[str]:
    paths = []
    if args.image:
        paths.extend(args.image)
    if args.folder:
        folder = Path(args.folder)
        if not folder.is_dir():
            print(f"[ERROR] Folder not found: {args.folder}")
            sys.exit(1)
        for p in sorted(folder.iterdir()):
            if p.suffix.lower() in SUPPORTED_EXTS:
                paths.append(str(p))
        if not paths:
            print(f"[WARNING] No supported image files found in: {args.folder}")
    return paths


def _save_annotated(image_path: str, face_results: list[dict], save_path: str):
    img_bgr = cv2.imread(image_path)
    if img_bgr is None:
        print(f"[WARNING] Could not reload image for annotation: {image_path}")
        return
    annotated = draw_face_predictions(img_bgr, face_results)
    cv2.imwrite(save_path, annotated)
    print(f"[INFO] Annotated image saved to: {save_path}")


def main():
    parser = argparse.ArgumentParser(
        description='FER Inference — predict facial expressions from images.'
    )
    parser.add_argument('--image', nargs='+', help='Path(s) to input image(s).')
    parser.add_argument('--folder', help='Folder of images to process.')
    parser.add_argument('--detect-face', action='store_true',
                        help='Run face detection before prediction.')
    parser.add_argument('--face-method', default='mtcnn', choices=['mtcnn', 'haar'],
                        help='Face detection method (default: mtcnn).')
    parser.add_argument('--weights', default='../models/model_weights.pth',
                        help='Path to model_weights.pth (default: ../models/model_weights.pth).')
    parser.add_argument('--device', default='auto',
                        help='Device: auto | cpu | cuda | cuda:0 (default: auto).')
    parser.add_argument('--save-output', help='Save annotated image to this path.')
    parser.add_argument('--save-plot', help='Save probability bar chart to this path.')

    args = parser.parse_args()

    if not args.image and not args.folder:
        parser.print_help()
        sys.exit(1)

    image_paths = _collect_images(args)
    if not image_paths:
        print("[ERROR] No images to process.")
        sys.exit(1)

    predictor = FERPredictor(weights_path=args.weights, device=args.device)

    for img_path in image_paths:
        if not os.path.exists(img_path):
            print(f"[WARNING] File not found, skipping: {img_path}")
            continue

        if args.detect_face:
            face_results = predictor.predict_with_face_detection(
                img_path, method=args.face_method
            )
            _print_face_result(img_path, face_results)

            if args.save_output:
                save_path = args.save_output if len(image_paths) == 1 else \
                    f"{Path(args.save_output).stem}_{Path(img_path).stem}{Path(args.save_output).suffix}"
                _save_annotated(img_path, face_results, save_path)

            if args.save_plot and face_results:
                # Plot probabilities of first face
                from utils import plot_emotion_bars
                plot_emotion_bars(face_results[0]['probabilities'],
                                  title=f"{face_results[0]['emotion']} ({face_results[0]['confidence']*100:.1f}%)",
                                  save_path=args.save_plot)
        else:
            result = predictor.predict_image(img_path)
            _print_result(img_path, result)

            if args.save_output:
                img = Image.open(img_path)
                save_path = args.save_output if len(image_paths) == 1 else \
                    f"{Path(args.save_output).stem}_{Path(img_path).stem}{Path(args.save_output).suffix}"
                visualize_prediction(img, result, save_path=save_path)

            if args.save_plot:
                from utils import plot_emotion_bars
                plot_emotion_bars(result['probabilities'],
                                  title=f"{result['emotion']} ({result['confidence']*100:.1f}%)",
                                  save_path=args.save_plot)


if __name__ == '__main__':
    main()