"""
Test Japanese segmentation with real customer documents.

Extracts text samples from customer files and compares segmentation approaches.
"""

import sys
import os
import time
from pathlib import Path
from typing import List, Dict, Any
import warnings
warnings.filterwarnings('ignore')

sys.path.insert(0, '/home/neil/Documents/textiq-doc-extraction')

from src.extraction_v2.domain_term_extractor import DomainTermExtractor


def extract_text_from_docx(file_path: str, max_chars: int = 1000) -> str:
    """Extract text from DOCX file."""
    try:
        from docx import Document
        doc = Document(file_path)
        text = "\n".join([para.text for para in doc.paragraphs if para.text.strip()])
        return text[:max_chars]
    except Exception as e:
        print(f"  Failed to read DOCX: {e}")
        return None


def extract_text_from_xlsx(file_path: str, max_chars: int = 1000) -> str:
    """Extract text from Excel file."""
    try:
        import openpyxl
        wb = openpyxl.load_workbook(file_path, read_only=True, data_only=True)
        text_parts = []

        # Read first sheet
        ws = wb.active
        for row in ws.iter_rows(max_row=50, max_col=10, values_only=True):
            row_text = " ".join([str(cell) for cell in row if cell is not None])
            if row_text.strip():
                text_parts.append(row_text)
                if len(" ".join(text_parts)) > max_chars:
                    break

        wb.close()
        return " ".join(text_parts)[:max_chars]
    except Exception as e:
        print(f"  Failed to read XLSX: {e}")
        return None


def find_customer_documents(base_dir: str = "CustomerDocument", limit: int = 10) -> List[str]:
    """Find Japanese customer documents."""
    documents = []
    base_path = Path(base_dir)

    if not base_path.exists():
        print(f"CustomerDocument directory not found: {base_path}")
        return []

    # Find DOCX and XLSX files
    for pattern in ["**/*.docx", "**/*.xlsx"]:
        for file_path in base_path.glob(pattern):
            if len(documents) >= limit:
                break
            # Skip temp files
            if file_path.name.startswith("~"):
                continue
            documents.append(str(file_path))

    return documents[:limit]


def segment_with_mecab(text: str) -> str:
    """Segment Japanese text using MeCab."""
    try:
        import MeCab
        tagger = MeCab.Tagger("-Owakati")
        return tagger.parse(text).strip()
    except Exception:
        return None


def segment_with_janome(text: str) -> str:
    """Segment Japanese text using Janome."""
    try:
        from janome.tokenizer import Tokenizer
        tokenizer = Tokenizer()
        tokens = [token.surface for token in tokenizer.tokenize(text)]
        return " ".join(tokens)
    except Exception:
        return None


def segment_with_sudachi(text: str) -> str:
    """Segment Japanese text using SudachiPy."""
    try:
        from sudachipy import tokenizer, dictionary
        tokenizer_obj = dictionary.Dictionary().create()
        mode = tokenizer.Tokenizer.SplitMode.C
        tokens = [m.surface() for m in tokenizer_obj.tokenize(text, mode)]
        return " ".join(tokens)
    except Exception:
        return None


def test_document(file_path: str, doc_index: int) -> Dict[str, Any]:
    """Test segmentation on a single document."""
    print(f"\n{'=' * 70}")
    print(f"Document {doc_index}: {Path(file_path).name}")
    print(f"{'=' * 70}")

    # Extract text
    if file_path.endswith('.docx'):
        text = extract_text_from_docx(file_path)
    elif file_path.endswith('.xlsx'):
        text = extract_text_from_xlsx(file_path)
    else:
        return None

    if not text or len(text.strip()) < 50:
        print("  Skipped: Insufficient text content")
        return None

    print(f"Extracted text ({len(text)} chars):")
    print(f"  {text[:200]}...")

    extractor = DomainTermExtractor()

    # Test baseline
    print(f"\n{'─' * 70}")
    print("Baseline (No Segmentation)")
    print(f"{'─' * 70}")
    start = time.time()
    baseline_result = extractor.extract_from_chunk(text)
    baseline_time = time.time() - start
    print(f"  Terms: {baseline_result['domain_terms'][:5]}{'...' if len(baseline_result['domain_terms']) > 5 else ''}")
    print(f"  Count: {len(baseline_result['domain_terms'])}")
    print(f"  Category: {baseline_result['domain_category']}")
    print(f"  Time: {int(baseline_time * 1000)}ms")

    # Test segmentation methods
    results = {
        'file_path': file_path,
        'text_length': len(text),
        'baseline': {
            'terms': baseline_result['domain_terms'],
            'count': len(baseline_result['domain_terms']),
            'category': baseline_result['domain_category'],
            'time_ms': int(baseline_time * 1000)
        }
    }

    segmenters = [
        ("MeCab", segment_with_mecab),
        ("Janome", segment_with_janome),
        ("SudachiPy", segment_with_sudachi),
    ]

    for name, segmenter in segmenters:
        segmented_text = segmenter(text)
        if segmented_text is None:
            continue

        print(f"\n{'─' * 70}")
        print(f"{name}")
        print(f"{'─' * 70}")

        start = time.time()
        result = extractor.extract_from_chunk(segmented_text)
        elapsed = time.time() - start

        print(f"  Terms: {result['domain_terms'][:5]}{'...' if len(result['domain_terms']) > 5 else ''}")
        print(f"  Count: {len(result['domain_terms'])}")
        print(f"  Category: {result['domain_category']}")
        print(f"  Time: {int(elapsed * 1000)}ms")

        results[name.lower()] = {
            'terms': result['domain_terms'],
            'count': len(result['domain_terms']),
            'category': result['domain_category'],
            'time_ms': int(elapsed * 1000)
        }

    return results


def main():
    """Run comprehensive comparison on customer documents."""
    print("=" * 70)
    print("COMPREHENSIVE JAPANESE SEGMENTATION TEST")
    print("Testing with Real Customer Documents")
    print("=" * 70)

    # Find documents
    print("\nSearching for customer documents...")
    documents = find_customer_documents(limit=10)

    if not documents:
        print("No customer documents found!")
        return

    print(f"Found {len(documents)} documents to test\n")

    # Test each document
    all_results = []
    successful_tests = 0

    for i, doc_path in enumerate(documents, 1):
        try:
            result = test_document(doc_path, i)
            if result:
                all_results.append(result)
                successful_tests += 1
        except Exception as e:
            print(f"\n  ERROR processing document: {e}")
            continue

    # Summary
    print("\n\n" + "=" * 70)
    print("SUMMARY - Customer Document Analysis")
    print("=" * 70)
    print(f"\nDocuments analyzed: {successful_tests}")

    if not all_results:
        print("No results to summarize")
        return

    # Calculate averages
    methods = ['baseline', 'mecab', 'janome', 'sudachipy']
    method_stats = {m: {'count': 0, 'total_terms': 0, 'total_time': 0} for m in methods}

    for result in all_results:
        for method in methods:
            if method in result:
                method_stats[method]['count'] += 1
                method_stats[method]['total_terms'] += result[method]['count']
                method_stats[method]['total_time'] += result[method]['time_ms']

    print("\n┌─────────────────┬──────────────┬────────────────┬──────────────┐")
    print("│ Method          │ Avg Terms    │ Avg Time       │ Improvement  │")
    print("├─────────────────┼──────────────┼────────────────┼──────────────┤")

    baseline_avg = (method_stats['baseline']['total_terms'] / method_stats['baseline']['count']
                   if method_stats['baseline']['count'] > 0 else 0)

    for method in methods:
        stats = method_stats[method]
        if stats['count'] == 0:
            continue

        avg_terms = stats['total_terms'] / stats['count']
        avg_time = stats['total_time'] / stats['count']
        improvement = ((avg_terms - baseline_avg) / baseline_avg * 100) if baseline_avg > 0 else 0

        method_name = method.capitalize().ljust(15)
        print(f"│ {method_name} │ {avg_terms:>12.1f} │ {avg_time:>12.0f}ms │ {improvement:>11.1f}% │")

    print("└─────────────────┴──────────────┴────────────────┴──────────────┘")

    # Find best performer
    best_method = max(
        [(m, method_stats[m]['total_terms'] / method_stats[m]['count'])
         for m in methods if method_stats[m]['count'] > 0],
        key=lambda x: x[1]
    )

    print(f"\n✓ Best performer: {best_method[0].upper()} ({best_method[1]:.1f} avg terms)")

    # Category detection analysis
    print("\n" + "=" * 70)
    print("Category Detection Analysis")
    print("=" * 70)

    for method in methods:
        if method_stats[method]['count'] == 0:
            continue

        categories = {}
        for result in all_results:
            if method in result:
                cat = result[method]['category']
                categories[cat] = categories.get(cat, 0) + 1

        print(f"\n{method.capitalize()}:")
        for cat, count in sorted(categories.items(), key=lambda x: -x[1]):
            print(f"  {cat}: {count} documents")

    print("\n" + "=" * 70)


if __name__ == "__main__":
    main()