from docx import Document
from docx.shared import Inches, Pt
from docx.enum.text import WD_ALIGN_PARAGRAPH
from typing import Callable, List, Any
import os
from dataclasses import dataclass

@dataclass
class PageConfig:
    page_function: Callable
    add_page_break: bool = True

class WordDocumentBuilder:
    """Build Word documents using configurable page functions."""
    
    def __init__(self):
        self.doc = Document()
        self.current_section = 1
        self.figure_number = 0;
        self.table_number = 0;
    
    def add_page_break(self):
        """Add a page break to the document."""
        self.doc.add_page_break()
    
    def create_document(self, page_functions: List[PageConfig], output_path: str, **kwargs):
        """
        Create a Word document by executing page functions.
        
        Args:
            page_functions: List of functions that add content to the document
            output_path: Path where the document will be saved
            **kwargs: Additional parameters passed to all page functions
        
        Returns:
            The created Document object
        """
        for i, page_conf in enumerate(page_functions):
            # Execute page function with builder and any additional kwargs
            page_conf.page_function(self, **kwargs)
            
            # Add page break after each page except the last
            if i < len(page_functions) - 1 and page_conf.add_page_break:
                self.add_page_break()
            
            # Increment the section state
            self.current_section += 1
        
        # Ensure directory exists
        os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else '.', exist_ok=True)
        
        # Save document
        self.doc.save(output_path)
        return self.doc


# Example page functions
def title_page(builder: WordDocumentBuilder, title: str = "Document Title", author: str = "Author", **kwargs):
    """Create a title page."""
    title_para = builder.doc.add_paragraph()
    title_run = title_para.add_run(title)
    title_run.font.size = Pt(28)
    title_run.bold = True
    title_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
    
    builder.doc.add_paragraph()  # Spacing
    
    author_para = builder.doc.add_paragraph(author)
    author_para.alignment = WD_ALIGN_PARAGRAPH.CENTER


def content_page(builder: WordDocumentBuilder, heading: str = "Section", content: str = "", **kwargs):
    """Create a content page with heading and body text."""
    builder.doc.add_heading(heading, level=1)
    builder.doc.add_paragraph(content)


def table_page(builder: WordDocumentBuilder, table_data: List[List[str]] = [], table_title: str = "Data Table", **kwargs):
    """Create a page with a table."""
    if table_data is None:
        table_data = [["Header 1", "Header 2"], ["Data 1", "Data 2"]]
    
    builder.doc.add_heading(table_title, level=1)
    
    table = builder.doc.add_table(rows=len(table_data), cols=len(table_data[0]))
    table.style = 'Light Grid Accent 1'
    
    for i, row in enumerate(table_data):
        for j, cell_value in enumerate(row):
            table.rows[i].cells[j].text = str(cell_value)


def image_page(builder: WordDocumentBuilder, image_path: str = "", caption: str = "", **kwargs):
    """Create a page with an image."""
    if image_path and os.path.exists(image_path):
        builder.doc.add_paragraph(caption)
        builder.doc.add_picture(image_path, width=Inches(5))

# Usage example
if __name__ == "__main__":
    # Create document builder
    builder = WordDocumentBuilder()
    
    # Define pages
    pages = [
        PageConfig(title_page, add_page_break=True),
        PageConfig(content_page, add_page_break=True),
        PageConfig(table_page, add_page_break=False),
    ]
    
    # Create document with custom parameters
    doc = builder.create_document(
        pages,
        "output/my_document.docx",
        title="My Custom Report",
        author="John Doe",
        heading="Introduction",
        content="This is the introduction section with detailed information.",
        table_data=[
            ["Product", "Price", "Quantity"],
            ["Widget", "$10", "100"],
            ["Gadget", "$25", "50"]
        ],
        table_title="Sales Data"
    )
    
    print("Document created successfully!")