#!/usr/bin/env python3
"""
Pipeline OCR pour extraire le texte des manuels de chinois (PDF scannés)
Utilise PaddleOCR pour la reconnaissance de caractères chinois
"""

import os
import sys
import argparse
from pathlib import Path
from typing import List, Tuple
import json
from datetime import datetime
from dotenv import load_dotenv

try:
    from paddleocr import PaddleOCR
    from pdf2image import convert_from_path
    from PIL import Image
    from openai import OpenAI
except ImportError as e:
    print(f"Erreur d'import: {e}")
    print("\nInstallez les dépendances avec:")
    print("pip install paddlepaddle paddleocr pdf2image Pillow openai python-dotenv")
    print("\nSur Linux, installez aussi: sudo apt-get install poppler-utils")
    sys.exit(1)

# Charger les variables d'environnement
load_dotenv()


class ChineseOCRPipeline:
    """Pipeline OCR pour traiter les PDFs de manuels chinois"""

    def __init__(self, lang: str = 'ch', use_gpu: bool = False, use_ai_correction: bool = False):
        """
        Initialise le pipeline OCR

        Args:
            lang: Langue pour l'OCR ('ch' pour chinois, 'en' pour anglais)
            use_gpu: Utiliser le GPU si disponible
            use_ai_correction: Activer la correction IA avec GPT-4o-mini
        """
        print(f"Initialisation de PaddleOCR (langue: {lang})...")
        self.ocr = PaddleOCR(
            use_textline_orientation=True,  # Détection d'angle pour images tournées
            lang=lang
        )
        print("✓ PaddleOCR initialisé")

        # Initialiser OpenAI si correction activée
        self.use_ai_correction = use_ai_correction
        self.openai_client = None
        if use_ai_correction:
            api_key = os.getenv('OPENAI_API_KEY')
            if not api_key:
                print("⚠️  OPENAI_API_KEY non trouvée dans .env - correction IA désactivée")
                self.use_ai_correction = False
            else:
                self.openai_client = OpenAI(api_key=api_key)
                print("✓ Correction IA (GPT-4o-mini) activée")

    def pdf_to_images(self, pdf_path: Path, dpi: int = 300) -> List[Image.Image]:
        """
        Convertit un PDF en images

        Args:
            pdf_path: Chemin vers le PDF
            dpi: Résolution de conversion (300 recommandé pour l'OCR)

        Returns:
            Liste d'images PIL
        """
        print(f"  Conversion du PDF en images (DPI: {dpi})...")
        try:
            images = convert_from_path(str(pdf_path), dpi=dpi)
            print(f"  ✓ {len(images)} pages converties")
            return images
        except Exception as e:
            print(f"  ✗ Erreur lors de la conversion: {e}")
            return []

    def extract_text_from_image(self, image: Image.Image) -> Tuple[str, List[dict]]:
        """
        Extrait le texte d'une image avec PaddleOCR

        Args:
            image: Image PIL

        Returns:
            Tuple (texte_brut, résultats_détaillés)
        """
        import numpy as np

        # Convertir PIL Image en numpy array
        img_array = np.array(image)

        # Lancer l'OCR (utilise la nouvelle API predict)
        result = self.ocr.predict(img_array)

        # Extraire le texte
        # Le format de retour est une liste contenant un objet OCRResult
        if result is None or not isinstance(result, list) or len(result) == 0:
            return "", []

        # Prendre le premier résultat (un seul résultat par image)
        ocr_result = result[0]

        # Accéder aux attributs comme un dictionnaire
        rec_texts = ocr_result.get('rec_texts', [])
        rec_scores = ocr_result.get('rec_scores', [])
        rec_polys = ocr_result.get('rec_polys', [])

        if not rec_texts:
            return "", []

        text_lines = []
        detailed_results = []

        # Itérer sur les textes, scores et polygones
        for i, text in enumerate(rec_texts):
            confidence = rec_scores[i] if i < len(rec_scores) else 0.0
            bbox = rec_polys[i].tolist() if i < len(rec_polys) else []

            text_lines.append(text)
            detailed_results.append({
                'text': text,
                'confidence': float(confidence),
                'bbox': bbox
            })

        return '\n'.join(text_lines), detailed_results

    def correct_text_with_ai(self, text: str, page_num: int) -> str:
        """
        Corrige le texte OCR avec GPT-4o-mini

        Args:
            text: Texte brut de l'OCR
            page_num: Numéro de la page

        Returns:
            Texte corrigé
        """
        if not self.use_ai_correction or not self.openai_client:
            return text

        if not text.strip():
            return text

        try:
            response = self.openai_client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {
                        "role": "system",
                        "content": """Tu es un expert en chinois qui corrige les erreurs OCR dans des manuels de chinois.

Ton travail:
1. Corriger les erreurs OCR évidentes (caractères manquants, mal reconnus)
2. Vérifier la cohérence pinyin/caractères chinois
3. Corriger la ponctuation chinoise si nécessaire
4. Si une phrase semble incomplète (ex: "心里有说不出的兴" au lieu de "心里有说不出的高兴"), la compléter
5. Corriger les questions qui citent le texte principal

IMPORTANT:
- Ne PAS modifier la structure ou l'ordre du texte
- Ne PAS ajouter de contenu qui n'était pas là
- Ne PAS traduire
- Garder EXACTEMENT le même formatage (sauts de ligne, espaces)
- Seulement corriger les erreurs évidentes"""
                    },
                    {
                        "role": "user",
                        "content": f"Corrige les erreurs OCR dans ce texte de la page {page_num} d'un manuel de chinois:\n\n{text}"
                    }
                ],
                temperature=0.1,
                max_tokens=4000
            )

            corrected_text = response.choices[0].message.content
            return corrected_text

        except Exception as e:
            print(f"    ⚠️  Erreur correction IA: {e}")
            return text

    def process_pdf(self, pdf_path: Path, output_dir: Path) -> bool:
        """
        Traite un fichier PDF et sauvegarde les résultats

        Args:
            pdf_path: Chemin vers le PDF
            output_dir: Dossier de sortie

        Returns:
            True si succès, False sinon
        """
        print(f"\n{'='*60}")
        print(f"Traitement: {pdf_path.name}")
        print(f"{'='*60}")

        # Convertir PDF en images
        images = self.pdf_to_images(pdf_path)
        if not images:
            return False

        # Créer le dossier de sortie
        output_dir.mkdir(parents=True, exist_ok=True)

        # Nom de base pour les fichiers de sortie
        base_name = pdf_path.stem

        # Traiter chaque page
        all_text = []
        all_text_corrected = []
        all_detailed_results = []

        for page_num, image in enumerate(images, start=1):
            print(f"\n  {'─'*50}")
            print(f"  📄 Page {page_num}/{len(images)}")
            print(f"  {'─'*50}")

            text, detailed = self.extract_text_from_image(image)

            # Afficher un aperçu du texte extrait
            if text:
                lines = text.split('\n')
                num_lines = len(lines)
                preview = '\n    '.join(lines[:5])  # Afficher les 5 premières lignes
                print(f"  ✓ {num_lines} ligne(s) détectée(s)")
                print(f"  📝 Aperçu OCR brut:")
                print(f"    {preview}")
                if num_lines > 5:
                    print(f"    ... ({num_lines - 5} ligne(s) supplémentaire(s))")
            else:
                print(f"  ⚠️  Aucun texte détecté sur cette page")

            # Correction IA si activée
            corrected_text = text
            if self.use_ai_correction and text:
                print(f"  🤖 Correction IA en cours...")
                corrected_text = self.correct_text_with_ai(text, page_num)
                print(f"  ✓ Correction terminée")

            all_text.append(f"=== Page {page_num} ===\n{text}\n")
            all_text_corrected.append(f"=== Page {page_num} ===\n{corrected_text}\n")
            all_detailed_results.append({
                'page': page_num,
                'text': text,
                'text_corrected': corrected_text if self.use_ai_correction else None,
                'details': detailed
            })

        # Sauvegarder le texte brut
        txt_output = output_dir / f"{base_name}.txt"
        with open(txt_output, 'w', encoding='utf-8') as f:
            f.write('\n'.join(all_text))
        print(f"  ✓ Texte brut sauvegardé: {txt_output}")

        # Sauvegarder le texte corrigé si correction IA activée
        if self.use_ai_correction:
            txt_corrected_output = output_dir / f"{base_name}_corrected.txt"
            with open(txt_corrected_output, 'w', encoding='utf-8') as f:
                f.write('\n'.join(all_text_corrected))
            print(f"  ✓ Texte corrigé sauvegardé: {txt_corrected_output}")

        # Sauvegarder les résultats détaillés (JSON)
        json_output = output_dir / f"{base_name}.json"
        with open(json_output, 'w', encoding='utf-8') as f:
            json.dump({
                'source_pdf': str(pdf_path),
                'processed_at': datetime.now().isoformat(),
                'ai_correction_enabled': self.use_ai_correction,
                'pages': all_detailed_results
            }, f, ensure_ascii=False, indent=2)
        print(f"  ✓ Détails sauvegardés: {json_output}")

        return True

    def process_directory(self, base_dir: Path = Path("Raw")) -> dict:
        """
        Traite tous les PDFs dans Raw/*/PDF/

        Args:
            base_dir: Dossier de base (défaut: Raw)

        Returns:
            Dictionnaire avec statistiques de traitement
        """
        stats = {
            'total': 0,
            'success': 0,
            'failed': 0,
            'files': []
        }

        # Trouver tous les dossiers PDF
        pdf_dirs = list(base_dir.glob("*/PDF"))

        print(f"\nTrouvé {len(pdf_dirs)} dossiers de PDFs")

        for pdf_dir in pdf_dirs:
            # Dossier de sortie (même niveau que PDF)
            output_dir = pdf_dir.parent / "OCR"

            # Trouver tous les PDFs
            pdf_files = list(pdf_dir.glob("*.pdf")) + list(pdf_dir.glob("*.PDF"))

            print(f"\n{pdf_dir.relative_to(base_dir)}: {len(pdf_files)} fichiers")

            for pdf_file in pdf_files:
                stats['total'] += 1
                success = self.process_pdf(pdf_file, output_dir)

                if success:
                    stats['success'] += 1
                    stats['files'].append(str(pdf_file.relative_to(base_dir)))
                else:
                    stats['failed'] += 1

        return stats


def main():
    parser = argparse.ArgumentParser(
        description="Pipeline OCR pour manuels de chinois",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Exemples:
  # Mode auto - traiter tous les PDFs
  python ocr_pipeline.py --mode auto

  # Mode manuel - traiter un fichier spécifique
  python ocr_pipeline.py --mode manual --input "Raw/DevelopChinese/PDF/听力Chapter1.pdf"

  # Avec GPU
  python ocr_pipeline.py --mode auto --gpu
        """
    )

    parser.add_argument(
        '--mode',
        choices=['auto', 'manual'],
        default='auto',
        help='Mode de traitement (auto: tous les PDFs, manual: un fichier)'
    )

    parser.add_argument(
        '--input',
        type=str,
        help='Chemin du fichier PDF (requis en mode manual)'
    )

    parser.add_argument(
        '--output',
        type=str,
        help='Dossier de sortie personnalisé (optionnel en mode manual)'
    )

    parser.add_argument(
        '--gpu',
        action='store_true',
        help='Utiliser le GPU si disponible'
    )

    parser.add_argument(
        '--dpi',
        type=int,
        default=300,
        help='Résolution pour la conversion PDF (défaut: 300)'
    )

    parser.add_argument(
        '--ai-correction',
        action='store_true',
        help='Activer la correction IA avec GPT-4o-mini (nécessite OPENAI_API_KEY dans .env)'
    )

    args = parser.parse_args()

    # Validation
    if args.mode == 'manual' and not args.input:
        parser.error("--input est requis en mode manual")

    # Initialiser le pipeline
    pipeline = ChineseOCRPipeline(
        lang='ch',
        use_gpu=args.gpu,
        use_ai_correction=args.ai_correction
    )

    # Mode auto
    if args.mode == 'auto':
        print("\n🚀 Mode AUTO - Traitement de tous les PDFs")
        stats = pipeline.process_directory()

        print(f"\n{'='*60}")
        print("RÉSUMÉ")
        print(f"{'='*60}")
        print(f"Total:   {stats['total']} fichiers")
        print(f"Succès:  {stats['success']} fichiers")
        print(f"Échecs:  {stats['failed']} fichiers")
        print(f"{'='*60}")

    # Mode manuel
    else:
        print(f"\n🎯 Mode MANUAL - Traitement de: {args.input}")

        pdf_path = Path(args.input)
        if not pdf_path.exists():
            print(f"Erreur: Le fichier {args.input} n'existe pas")
            sys.exit(1)

        # Déterminer le dossier de sortie
        if args.output:
            output_dir = Path(args.output)
        else:
            # Par défaut: Raw/*/OCR/
            output_dir = pdf_path.parent.parent / "OCR"

        success = pipeline.process_pdf(pdf_path, output_dir)

        if success:
            print("\n✓ Traitement réussi!")
        else:
            print("\n✗ Échec du traitement")
            sys.exit(1)


if __name__ == "__main__":
    main()