"""
training_doc.py
===============
Training module for DOC/DOCX document formats.

Handles text extraction from Word documents (.doc, .docx),
semantic chunking, and FAISS embedding creation.

- .docx: Extracted directly using python-docx
- .doc:  Converted to .docx via LibreOffice, then extracted

Shared utilities (semantic chunking, embedding creation, FAISS merging)
are imported from training.py to maintain DRY principles.
"""

import os
import subprocess
import tempfile
import time

from docx import Document as DocxDocument
from langchain_core.documents import Document
from langchain_ollama import OllamaEmbeddings

# Import shared utilities from the existing training module
from training import (
    split_text_with_semantic_chunker,
    count_tokens_in_documents,
    save_documents_to_txt,
    create_and_save_embeddings,
    merge_all_faiss,
)


# ---------------------------------------------------------------------------
# Document Loading Helpers
# ---------------------------------------------------------------------------

def _extract_text_from_docx(file_path: str) -> str:
    """
    Extract all text content from a .docx file using python-docx.

    Iterates through all paragraphs and table cells to capture
    the complete document content.
    """
    doc = DocxDocument(file_path)
    text_parts = []

    # Extract text from paragraphs
    for paragraph in doc.paragraphs:
        stripped = paragraph.text.strip()
        if stripped:
            text_parts.append(stripped)

    # Extract text from tables (often missed by simple paragraph extraction)
    for table in doc.tables:
        for row in table.rows:
            for cell in row.cells:
                cell_text = cell.text.strip()
                if cell_text:
                    text_parts.append(cell_text)

    return "\n".join(text_parts)


def _convert_doc_to_docx(doc_file_path: str) -> str:
    """
    Convert a legacy .doc file to .docx using LibreOffice headless mode.

    Returns the path to the newly created .docx file.
    Raises RuntimeError if LibreOffice conversion fails.
    """
    output_dir = os.path.dirname(doc_file_path) or "."

    try:
        result = subprocess.run(
            [
                "libreoffice",
                "--headless",
                "--convert-to", "docx",
                "--outdir", output_dir,
                doc_file_path,
            ],
            capture_output=True,
            text=True,
            timeout=120,  # 2-minute timeout for large documents
        )

        if result.returncode != 0:
            raise RuntimeError(
                f"LibreOffice conversion failed (exit code {result.returncode}): "
                f"{result.stderr}"
            )

    except FileNotFoundError:
        raise RuntimeError(
            "LibreOffice is not installed or not in PATH. "
            "It is required to convert legacy .doc files."
        )

    # Derive the expected output path
    base_name = os.path.splitext(os.path.basename(doc_file_path))[0]
    docx_path = os.path.join(output_dir, f"{base_name}.docx")

    if not os.path.exists(docx_path):
        raise RuntimeError(
            f"LibreOffice conversion produced no output file. "
            f"Expected: {docx_path}"
        )

    return docx_path


# ---------------------------------------------------------------------------
# Main Document Loader
# ---------------------------------------------------------------------------

def load_doc(file_path: str):
    """
    Load a DOC or DOCX file and return a list of LangChain Document objects.

    For .docx files, text is extracted directly using python-docx.
    For .doc files, the document is first converted to .docx via LibreOffice.

    Returns:
        list[Document]: List of LangChain Document objects with page_content set.
    """
    file_extension = os.path.splitext(file_path)[1].lower()
    converted_docx_path = None

    try:
        if file_extension == ".docx":
            text = _extract_text_from_docx(file_path)
        elif file_extension == ".doc":
            # Convert .doc → .docx first, then extract
            converted_docx_path = _convert_doc_to_docx(file_path)
            text = _extract_text_from_docx(converted_docx_path)
        else:
            raise ValueError(
                f"Unsupported file type: {file_extension}. "
                f"Only .doc and .docx are supported."
            )

        if not text.strip():
            print(f"Warning: No text content extracted from {file_path}")
            return []

        # Wrap in LangChain Document objects (one document per file)
        documents = [
            Document(
                page_content=text,
                metadata={"source": file_path, "format": file_extension}
            )
        ]

        print(f"Successfully loaded {file_extension} document: {len(text)} characters extracted.")
        return documents

    finally:
        # Clean up converted file if it was created
        if converted_docx_path and os.path.exists(converted_docx_path):
            os.remove(converted_docx_path)
            print(f"Cleaned up temporary converted file: {converted_docx_path}")


# ---------------------------------------------------------------------------
# Token Counting (compatible with Document objects from load_doc)
# ---------------------------------------------------------------------------

def count_tokens_in_doc_documents(documents):
    """
    Count total tokens across all Document objects.
    Uses whitespace tokenization for consistency with the existing system.
    """
    total_tokens = 0
    for doc in documents:
        tokens = doc.page_content.split()
        total_tokens += len(tokens)
    return total_tokens


# ---------------------------------------------------------------------------
# Main (for standalone testing)
# ---------------------------------------------------------------------------

def main():
    """Test the DOC/DOCX training pipeline standalone."""
    file_path = "test_document.docx"  # Replace with actual test file
    client_id = "test_client"
    reference_id = "test_ref"
    output_dir = f"temp/{client_id}_{reference_id}"

    start_time = time.time()
    print(f"Start Time: {start_time}")

    # Step 1: Load the document
    docs = load_doc(file_path)
    if not docs:
        print("No documents loaded. Exiting.")
        return

    token_count = count_tokens_in_doc_documents(docs)
    print(f"Total tokens: {token_count}")

    # Step 2: Create embeddings and split
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    split_documents = split_text_with_semantic_chunker(docs, embeddings)

    # Step 3: Save text chunks
    save_documents_to_txt(split_documents, output_dir)

    # Step 4: Create and save embeddings
    create_and_save_embeddings(split_documents, client_id, reference_id)

    # Step 5: Merge FAISS indices
    merge_all_faiss(client_id, reference_id)

    end_time = time.time()
    print(f"Training process took {end_time - start_time:.2f} seconds")
    print(f"Total tokens processed: {token_count}")


if __name__ == "__main__":
    main()