#!/usr/bin/env python3
"""
Single PDF Layer Extractor - Extract geological layers from a single PDF
"""
import click
import json
import logging
import PyPDF2
import os

from dotenv import load_dotenv
from pathlib import Path
from typing import List, Dict, Tuple
from datetime import datetime
from azure.ai.formrecognizer import DocumentAnalysisClient
from azure.core.credentials import AzureKeyCredential
from openai import AzureOpenAI

load_dotenv()

# Configuration Azure Document Intelligence
AZURE_DOC_ENDPOINT = os.getenv("AZURE_DOC_ENDPOINT")
AZURE_DOC_KEY = os.getenv("AZURE_DOC_KEY")

# Configuration Azure OpenAI
AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
AZURE_OPENAI_KEY = os.getenv("AZURE_OPENAI_KEY")
AZURE_DEPLOYMENT = os.getenv("AZURE_DEPLOYMENT", "gpt-4.1")
API_VERSION = os.getenv("API_VERSION", "2024-12-01-preview")

if not all([AZURE_DOC_ENDPOINT, AZURE_DOC_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_KEY]):
    raise ValueError(
        "❌ Missing Azure credentials! Please check your .env file.\n"
        "Required variables: AZURE_DOC_ENDPOINT, AZURE_DOC_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_KEY"
    )

# Configuration du logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

logger.info("✅ Environment variables loaded successfully")
logger.info(f"📍 Document Intelligence endpoint: {AZURE_DOC_ENDPOINT}")
logger.info(f"🤖 OpenAI deployment: {AZURE_DEPLOYMENT}")

class SinglePDFExtractor:
    """Extract geological layers from a single PDF"""

    def __init__(self):
        """Initialize Azure clients"""
        self.doc_client = DocumentAnalysisClient(
            endpoint=AZURE_DOC_ENDPOINT,
            credential=AzureKeyCredential(AZURE_DOC_KEY)
        )
        self.openai_client = AzureOpenAI(
            azure_endpoint=AZURE_OPENAI_ENDPOINT,
            api_key=AZURE_OPENAI_KEY,
            api_version=API_VERSION
        )

    def get_prompts(self) -> Tuple[str, str]:
        """Get the system and user prompts"""
        system_prompt = """Tu es un expert géotechnicien spécialisé dans l'analyse de logs de forage.
Ta tâche est d'extraire UNIQUEMENT les couches géologiques avec leurs profondeurs et descriptions COMPLÈTES.

RÈGLES CRITIQUES:
1. Chaque couche DOIT avoir une profondeur de début et de fin en mètres
2. La description doit contenir TOUT le texte associé à la couche (même sur plusieurs lignes)
3. Notation décimale: 0,3 ou 0.3 = 0.3m (PAS 3m!)
4. Profondeurs typiques: 0-0.3m, 0.3-1m, 1-4m, 4-10m, etc.
5. NE PAS séparer ou classifier le matériau - tout dans la description
6. IGNORER les mots isolés qui ne sont pas dans la colonne description

RÈGLE DE CONTINUITÉ DES COUCHES:
⚠️ Une couche géologique est CONTINUE entre deux profondeurs
⚠️ Si tu vois une profondeur isolée suivie de description, cherche OBLIGATOIREMENT la profondeur de fin
⚠️ Les profondeurs de fin peuvent être:
   - Sur la même ligne après la description
   - Sur une ligne suivante
   - Au début de la couche suivante (la fin d'une couche = début de la suivante)
⚠️ NE JAMAIS fragmenter une description continue en plusieurs couches

GESTION DES PAGES:
⚠️ Les couches peuvent continuer d'une page à l'autre
⚠️ Sur une nouvelle page, si la première profondeur n'est pas 0.0m, c'est la continuation de la page précédente
⚠️ Ne pas dupliquer les couches identiques sur plusieurs pages

SORTIE: JSON valide uniquement"""

        user_prompt_template = """Extrait toutes les couches géologiques de cette page de log de forage.

CE QUI EST UNE COUCHE GÉOLOGIQUE:
✓ Descriptions de matériaux: sol, roche, sable, argile, limon, remblai, terre végétale, etc.
✓ A une plage de profondeur ET des propriétés de matériau
✓ Décrit les caractéristiques physiques: couleur, consistance, humidité, altération, granulométrie

CE QUI N'EST PAS UNE COUCHE (IGNORER):
✗ Descriptions d'images ou légendes
✗ Métadonnées du projet (coordonnées, dates, noms d'entreprise)
✗ Descriptions d'équipement ou de méthodologie
✗ Données administratives ou commentaires sur la procédure
✗ Mots isolés dans d'autres colonnes (classifications, formations géologiques, etc.)

EXCLUSIONS ABSOLUES - NE JAMAIS extraire:
✗ Annotations techniques isolées: CZ, JT, JTx5, etc.
✗ Mesures techniques sans description de matériau
✗ Échantillons ("Undisturbed sample", "Disturbed sample", "sample", "échantillon")
✗ Lignes contenant uniquement des angles, pressions ou mesures
✗ Références à des essais ou tests sans description de sol

ALGORITHME DE DÉTECTION DES PROFONDEURS:
1. Pour chaque nombre qui ressemble à une profondeur (X.X m ou X,X m):
   - Si suivi de "-" ou "à" ou "bis" → profondeur de début, chercher la fin
   - Si précédé de "-" ou "à" ou "bis" → profondeur de fin
   - Si isolé → vérifier le contexte:
     * Début de description → profondeur de début
     * Fin de description → profondeur de fin
     * Nouvelle ligne → début de nouvelle couche
2. TOUJOURS vérifier la cohérence: fin couche N = début couche N+1
3. Si une couche n'a pas de profondeur de fin explicite, elle va jusqu'au début de la couche suivante

MOTIFS À RECHERCHER:
- "0.00-0.30m: Terre végétale, brun foncé..."
- "0,3 à 1,0m Sable fin beige..."
- "de 1m à 4m: Argile plastique grise..."
- "1.00 - 4.00 ARGILE limoneuse..."
- "0.3 m Description... 4.0 m" (profondeurs séparées)
- Tableaux avec colonnes profondeur/description

IDENTIFICATION DES COLONNES DE TABLEAU:
⚠️ Dans un tableau, identifier VISUELLEMENT les colonnes avant d'extraire
⚠️ Colonnes à IGNORER complètement:
   - Formation géologique (Quaternaire, Tertiaire, Quaternaire, etc.)
   - Stratigraphie
   - Classification (Colluvions, Alluvions, etc.)
   - Symboles ou codes isolés
   - Âge géologique
⚠️ Colonnes à UTILISER:
   - Description du matériau
   - Propriétés physiques (couleur, texture, consistance)
⚠️ Si doute sur une colonne, vérifier si elle contient des propriétés physiques détaillées

RÈGLE CRITIQUE DE SÉPARATION:
⚠️ Les profondeurs (0-0.3m, 0.30m, etc.) ne font JAMAIS partie de la description
⚠️ Si une ligne contient "0,30 m - Mutterboden, braun", alors:
   - depth_start: 0.0, depth_end: 0.3
   - description: "Mutterboden, braun" (SANS les profondeurs)
⚠️ Retirer TOUS les préfixes de profondeur du champ description
⚠️ NE PAS inclure les profondeurs répétées à la fin de la description

IMPORTANT:
- Capture TOUT le texte de description, même s'il est long ou sur plusieurs lignes
- Ne pas résumer ou reformuler - garder le texte original
- Inclure TOUTES les propriétés mentionnées (couleur, texture, humidité, inclusions, etc.)
- EXCLURE ABSOLUMENT les profondeurs du champ description
- Si le texte commence par "0,30 m - Mutterboden", la description est SEULEMENT "Mutterboden"
- Séparer clairement: profondeurs dans depth_start/depth_end, matériau et propriétés dans description
- NE PAS mélanger les colonnes d'un tableau - prendre uniquement la colonne description principale

VALIDATION OBLIGATOIRE avant de retourner le JSON:
□ Toutes les couches ont une profondeur de début ET de fin
□ Les profondeurs sont continues (pas de trous entre les couches)
□ Les profondeurs sont cohérentes (début < fin)
□ Les descriptions ne contiennent PAS les profondeurs
□ Les descriptions ne contiennent PAS de noms de colonnes de classification
□ Aucune annotation technique isolée n'est présente
□ Pas de duplication de couches identiques

Retourne ce JSON:
{{
    "layers": [
        {{
            "depth_start": 0.0,
            "depth_end": 0.3,
            "description": "Texte de description SANS les profondeurs - uniquement matériau et propriétés"
        }}
    ]
}}

EXEMPLE CONCRET:
Texte original: "0,00 - 0,30 m Mutterboden, dunkelbraun, humos"
Résultat attendu:
{{
    "depth_start": 0.0,
    "depth_end": 0.3,
    "description": "Mutterboden, dunkelbraun, humos"
}}
PAS: "description": "0,00 - 0,30 m Mutterboden, dunkelbraun, humos" ❌

TEXTE DE LA PAGE {page_num}:
{text}"""

        return system_prompt, user_prompt_template

    def extract_text_from_pdf(self, pdf_path: Path) -> List[str]:
        """Extract text from PDF with chunking"""
        logger.info(f"🔍 Extracting text from {pdf_path.name}")

        total_pages = self.get_page_count(pdf_path)
        logger.info(f"📄 PDF has {total_pages} pages")

        if total_pages == 0:
            logger.error("Could not determine page count")
            return []

        # Extraire par chunks de 2 pages
        pages_text = []
        chunk_size = 2
        for start_page in range(1, total_pages + 1, chunk_size):
            end_page = min(start_page + chunk_size - 1, total_pages)
            with open(pdf_path, "rb") as f:
                poller = self.doc_client.begin_analyze_document(
                    model_id="prebuilt-document",
                    document=f,
                    pages=f"{start_page}-{end_page}"
                )
            result = poller.result()
            for page in result.pages:
                page_lines = [line.content for line in (page.lines or [])]
                pages_text.append('\n'.join(page_lines))

        return pages_text

    def get_page_count(self, pdf_path: Path) -> int:
        """Get the number of pages in a PDF"""
        try:
            with open(pdf_path, "rb") as f:
                pdf_reader = PyPDF2.PdfReader(f)
                return len(pdf_reader.pages)
        except:
            pass
            return 0

    def _extract_table_text(self, table) -> str:
        """Extract text from table with better structure preservation"""
        if not hasattr(table, 'cells') or not table.cells:
            return ""

        max_row = max(cell.row_index for cell in table.cells)
        max_col = max(cell.column_index for cell in table.cells)
        table_array = [["" for _ in range(max_col + 1)] for _ in range(max_row + 1)]

        for cell in table.cells:
            if hasattr(cell, 'content'):
                table_array[cell.row_index][cell.column_index] = cell.content

        lines = []
        for row_idx, row in enumerate(table_array):
            row_text = " | ".join(str(cell).strip() for cell in row)
            if row_text.strip():
                lines.append(row_text)
                # Add separator after header row
                if row_idx == 0 and len(lines) == 1:
                    lines.append("-" * len(row_text))

        return "\n".join(lines)

    def extract_layers_from_page(self, page_text: str, page_num: int, system_prompt: str, user_prompt_template: str) -> List[Dict]:
        """Extract layers from a single page using GPT"""
        if not page_text.strip():
            logger.warning(f"⚠️  Page {page_num} is empty")
            return []

        max_chars = 8000 # todo
        user_prompt = user_prompt_template.format(
            page_num=page_num,
            text=page_text[:max_chars]
        )

        try:
            logger.info(f"🤖 Processing page {page_num} with GPT-4...")

            response = self.openai_client.chat.completions.create(
                model=AZURE_DEPLOYMENT,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0.1,
                max_tokens=4000, # todo
                response_format={"type": "json_object"}
            )

            result = json.loads(response.choices[0].message.content)
            valid_layers = []

            for layer in result.get("layers", []):
                try:
                    depth_start = float(str(layer.get("depth_start", 0)).replace(",", "."))
                    depth_end = float(str(layer.get("depth_end", 0)).replace(",", "."))
                    description = layer.get("description", "").strip()

                    if (depth_start < depth_end and depth_end <= 100 and depth_start >= 0 and description):
                        valid_layers.append({
                            "page": page_num,
                            "depth_start": depth_start,
                            "depth_end": depth_end,
                            "description": description
                        })
                except Exception as e:
                    logger.warning(f"   Invalid layer data: {e}")
                    continue
            logger.info(f"   ✓ Found {len(valid_layers)} valid layers on page {page_num}")
            return valid_layers

        except Exception as e:
            logger.error(f"❌ GPT extraction failed for page {page_num}: {e}")
            return []

    def merge_duplicate_layers(self, layers: List[Dict]) -> List[Dict]:
        """Merge duplicate layers across pages"""
        if not layers:
            return []

        layers.sort(key=lambda x: (x['depth_start'], x['page']))

        merged = []
        current = None

        for layer in layers:
            if current is None:
                current = layer.copy()
            else:
                if (abs(current['depth_start'] - layer['depth_start']) < 0.01 and
                    abs(current['depth_end'] - layer['depth_end']) < 0.01):
                    if len(layer['description']) > len(current['description']):
                        current['description'] = layer['description']
                    current['page'] = min(current['page'], layer['page'])
                else:
                    merged.append(current)
                    current = layer.copy()

        if current:
            merged.append(current)

        return merged

    def extract_layers_from_pdf(self, pdf_path: Path) -> Dict:
        """Main extraction function"""
        logger.info(f"\n🚀 Starting extraction for: {pdf_path.name}")
        start_time = datetime.now()

        pages_text = self.extract_text_from_pdf(pdf_path)

        if not pages_text:
            logger.error("❌ No text extracted from PDF")
            return {
                "metadata": {
                    "source_file": str(pdf_path.resolve()),
                    "extraction_date": datetime.now().isoformat(),
                    "extraction_method": "Azure Document Intelligence + GPT-4",
                    "pages": 0,
                    "layer_count": 0,
                    "error": "No text extracted from PDF"
                },
                "layers": []
            }

        system_prompt, user_prompt_template = self.get_prompts()

        all_layers = []
        for page_idx, page_text in enumerate(pages_text):
            page_layers = self.extract_layers_from_page(
                page_text, page_idx + 1, system_prompt, user_prompt_template
            )
            all_layers.extend(page_layers)

        logger.info(f"\n🔄 Merging duplicate layers...")
        merged_layers = self.merge_duplicate_layers(all_layers)

        extraction_time = (datetime.now() - start_time).total_seconds()

        logger.info(f"\n✅ Extraction completed!")
        logger.info(f"   Total layers found: {len(merged_layers)}")
        logger.info(f"   Extraction time: {extraction_time:.1f} seconds")

        return {
            "metadata": {
                "source_file": str(pdf_path.resolve()),
                "extraction_date": datetime.now().isoformat(),
                "extraction_method": "Azure Document Intelligence + GPT-4",
                "pages": len(pages_text),
                "layer_count": len(merged_layers),
                "extraction_time_seconds": round(extraction_time, 1)
            },
            "layers": merged_layers
        }

@click.command()
@click.argument('pdf_path', type=click.Path(exists=True, path_type=Path))
@click.option('--output', '-o', type=click.Path(path_type=Path), help='Output JSON file path (default: same name as PDF)')
@click.option('--pretty', is_flag=True, help='Pretty print the JSON output')
def main(pdf_path: Path, output: Path, pretty: bool):
    """Extract geological layers from a single PDF file."""
    click.echo("🪨 Geological Layer Extractor")
    click.echo("=" * 40)

    extractor = SinglePDFExtractor()
    result = extractor.extract_layers_from_pdf(pdf_path)

    if output is None:
        output_dir = Path('output')
        output_dir.mkdir(exist_ok=True)
        output = output_dir / pdf_path.with_suffix('.json').name
    with open(output, 'w', encoding='utf-8') as f:
        if pretty:
            json.dump(result, f, indent=2, ensure_ascii=False)
        else:
            json.dump(result, f, ensure_ascii=False)

    click.echo(f"\n💾 Results saved to: {output}")
    click.echo("\n📊 Extraction Summary:")
    click.echo(f"   Source: {pdf_path.name}")
    click.echo(f"   Pages processed: {result['metadata']['pages']}")
    click.echo(f"   Layers extracted: {result['metadata']['layer_count']}")
    click.echo(f"   Time taken: {result['metadata']['extraction_time_seconds']}s")

    if result['layers']:
        click.echo("\n📋 Preview of extracted layers:")
        for i, layer in enumerate(result['layers'][:3], 1):
            click.echo(f"\n   {i}. {layer['depth_start']:.1f} - {layer['depth_end']:.1f}m")
            click.echo(f"      {layer['description'][:80]}{'...' if len(layer['description']) > 80 else ''}")
        if len(result['layers']) > 3:
            click.echo(f"\n   ... and {len(result['layers']) - 3} more layers")

if __name__ == '__main__':
    main()