#!/usr/bin/env python3
"""
Simple Prompt Tester - Test current prompt without optimization
"""
import click
import json
import logging
import json
import PyPDF2
import os

from dotenv import load_dotenv
from pathlib import Path
from typing import List, Dict, Tuple
from datetime import datetime
from difflib import SequenceMatcher
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"
    )

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 PromptTester:
    """Test a single prompt on PDFs"""
    def __init__(self):
        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
        )
        self.ocr_cache = {}

    def get_current_prompt(self) -> Tuple[str, str]:
        """Get the current working prompt"""
        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"""
        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 in table_array:
            row_text = " | ".join(str(cell).strip() for cell in row)
            if row_text.strip():
                lines.append(row_text)
        return "\n".join(lines)

    def test_prompt_on_pdf(self, system_prompt: str, user_prompt_template: str, pdf_path: Path) -> List[Dict]:
        """Test prompt on a single PDF"""
        pages_text = self.extract_text_from_pdf(pdf_path)
        all_layers = []

        for page_idx, page_text in enumerate(pages_text):
            if not page_text.strip():
                continue

            user_prompt = user_prompt_template.format(
                page_num=page_idx + 1,
                text=page_text[:4000]
            )

            try:
                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=3000,
                    response_format={"type": "json_object"}
                )

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

                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):
                            all_layers.append({
                                "page": page_idx + 1,
                                "depth_start": depth_start,
                                "depth_end": depth_end,
                                "description": description
                            })
                    except:
                        continue

            except Exception as e:
                logger.error(f"GPT extraction failed for page {page_idx + 1}: {e}")

        return sorted(all_layers, key=lambda x: (x['page'], x['depth_start']))

    def compare_results(self, extracted: List[Dict], expected: List[Dict]) -> Dict:
        """Compare extracted layers with expected results"""
        exact_matches = 0
        depth_matches = 0

        for exp in expected:
            for ext in extracted:
                if (ext['page'] == exp['page'] and
                    abs(ext['depth_start'] - exp['depth_start']) < 0.01 and
                    abs(ext['depth_end'] - exp['depth_end']) < 0.01):
                    depth_matches += 1

                    desc_sim = SequenceMatcher(None,
                        ext['description'].lower(),
                        exp['description'].lower()
                    ).ratio()

                    if desc_sim == 1:
                        exact_matches += 1
                    break

        if len(expected) > 0:
            similarity = len([e for e in extracted if any(
                abs(e['depth_start'] - x['depth_start']) < 0.5 and
                abs(e['depth_end'] - x['depth_end']) < 0.5 and
                e['page'] == x['page'] for x in expected
            )]) / len(expected)
        else:
            similarity = 1.0 if len(extracted) == 0 else 0.0

        return {
            'similarity': similarity,
            'exact_matches': exact_matches,
            'depth_matches': depth_matches,
            'extracted_count': len(extracted),
            'expected_count': len(expected)
        }

    def test_all_pdfs(self, examples_dir: Path, expected_dir: Path):
        """Test prompt on all PDFs"""
        test_pairs = []
        for json_file in expected_dir.glob("*_layers.json"):
            pdf_name = json_file.stem.replace("_layers", "") + ".pdf"
            pdf_path = examples_dir / pdf_name
            if pdf_path.exists():
                test_pairs.append((pdf_path, json_file))

        if not test_pairs:
            click.echo("❌ No valid PDF/JSON pairs found!")
            return

        click.echo(f"\n📚 Found {len(test_pairs)} test PDFs")

        system_prompt, user_prompt_template = self.get_current_prompt()

        results = {}
        total_similarity = 0

        with click.progressbar(test_pairs, label='Testing PDFs') as pairs:
            for pdf_path, json_path in pairs:
                try:
                    with open(json_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)

                    if isinstance(data, dict) and 'layers' in data:
                        expected_layers = data['layers']
                    elif isinstance(data, list):
                        expected_layers = data
                    else:
                        logger.warning(f"Unknown format in {json_path}")
                        continue

                    extracted_layers = self.test_prompt_on_pdf(
                        system_prompt, user_prompt_template, pdf_path
                    )

                    comparison = self.compare_results(extracted_layers, expected_layers)
                    results[pdf_path.name] = comparison
                    total_similarity += comparison['similarity']

                    single_result = {
                        "metadata": {
                            "source_file": str(pdf_path.resolve()),
                            "expected_file": str(json_path.resolve()),
                            "extraction_time": datetime.now().isoformat(),
                            "extracted_count": comparison['extracted_count'],
                            "expected_count": comparison['expected_count']
                        },
                        "layers": extracted_layers
                    }

                    output_dir = Path("output")
                    output_dir.mkdir(parents=True, exist_ok=True)
                    output_file_path = output_dir / f"{pdf_path.stem}_test_result.json"

                    with open(output_file_path, 'w', encoding='utf-8') as out_f:
                        json.dump(single_result, out_f, indent=2, ensure_ascii=False)

                    click.echo(f"💾 Résultat sauvegardé : {output_file_path}")


                except Exception as e:
                    logger.error(f"Error with {pdf_path.name}: {e}")
                    results[pdf_path.name] = {
                        'similarity': 0,
                        'error': str(e)
                    }

        click.echo("\n" + "="*60)
        click.echo("📊 TEST RESULTS")
        click.echo("="*60)

        for pdf_name, result in results.items():
            if 'error' in result:
                click.echo(f"\n❌ {pdf_name}: ERROR - {result['error']}")
            else:
                click.echo(f"\n📄 {pdf_name}:")
                click.echo(f"   Similarity: {result['similarity']:.1%}")
                click.echo(f"   Extracted: {result['extracted_count']} layers")
                click.echo(f"   Expected: {result['expected_count']} layers")
                click.echo(f"   Exact matches: {result['exact_matches']}")
                click.echo(f"   Depth matches: {result['depth_matches']}")

        avg_similarity = total_similarity / len(test_pairs) if test_pairs else 0
        click.echo("\n" + "="*60)
        click.echo(f"📈 OVERALL: {avg_similarity:.1%} average similarity")
        click.echo("="*60)

        output_file = Path("output/test_results.json")
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump({
                'timestamp': datetime.now().isoformat(),
                'average_similarity': avg_similarity,
                'results': results
            }, f, indent=2)

        click.echo(f"\n💾 Detailed results saved to: {output_file}")

@click.command()
@click.option('--examples-dir', '-e', default='examples',help='Directory containing test PDFs')
@click.option('--expected-dir', '-x', default='output_man', help='Directory containing expected results')
@click.option('--pdf', '-p', help='Test a single PDF only')
def main(examples_dir, expected_dir, pdf):
    """Test current prompt on PDFs without optimization"""

    click.echo("🧪 Simple Prompt Tester")
    click.echo("=" * 30)

    tester = PromptTester()

    if pdf:
        pdf_path = Path(pdf)
        if not pdf_path.exists():
            raise click.ClickException(f"PDF not found: {pdf}")

        click.echo(f"\n📄 Testing single PDF: {pdf_path.name}")

        system_prompt, user_prompt_template = tester.get_current_prompt()
        layers = tester.test_prompt_on_pdf(system_prompt, user_prompt_template, pdf_path)

        click.echo(f"\n✅ Extracted {len(layers)} layers:")
        for i, layer in enumerate(layers, 1):
            click.echo(f"\n{i}. Depth: {layer['depth_start']:.2f} - {layer['depth_end']:.2f}m")
            click.echo(f"   Page: {layer['page']}")
            click.echo(f"   Description: {layer['description']}")
    else:
        tester.test_all_pdfs(Path(examples_dir), Path(expected_dir))

if __name__ == '__main__':
    main()
