import os
import tempfile
from docx import Document
from docx.shared import Inches

class DocumentGenerator:
    def __init__(self):
        self.temp_dir = tempfile.gettempdir()
    
    def create_word_document(self, content, content_type, filename_prefix="ArtifaX"):
        """Create a Word document from content"""
        doc = Document()
        
        # Add title based on content type
        title_map = {
            'mapping_document': 'API Mapping Document',
            'testcases': 'Test Cases Document', 
            'sequence_diagram': 'Sequence Diagram Document'
        }
        
        title = title_map.get(content_type, 'Generated Document')
        doc.add_heading(title, 0)
        doc.add_heading('Generated by ArtifaX', 2)
        
        # Add content
        if content_type == 'mapping_document':
            self._add_mapping_content(doc, content)
        elif content_type == 'testcases':
            self._add_testcase_content(doc, content)
        elif content_type == 'sequence_diagram':
            self._add_sequence_content(doc, content)
        else:
            doc.add_paragraph(content)
        
        # Save to temp file
        filename = f"{filename_prefix}_{content_type}.docx"
        filepath = os.path.join(self.temp_dir, filename)
        doc.save(filepath)
        
        return filepath, filename
    
    def _add_mapping_content(self, doc, content):
        """Format mapping document content"""
        doc.add_heading('Field Mappings', 1)
        
        # Split content into lines and create table if it looks like tabular data
        lines = content.split('\n')
        table_lines = [line for line in lines if '|' in line and line.strip()]
        
        if len(table_lines) > 1:
            # Create table
            headers = [cell.strip() for cell in table_lines[0].split('|')]
            table = doc.add_table(rows=1, cols=len(headers))
            table.style = 'Table Grid'
            
            # Add headers
            hdr_cells = table.rows[0].cells
            for i, header in enumerate(headers):
                if i < len(hdr_cells):
                    hdr_cells[i].text = header
            
            # Add data rows
            for line in table_lines[1:]:
                cells_data = [cell.strip() for cell in line.split('|')]
                row_cells = table.add_row().cells
                for i, cell_data in enumerate(cells_data):
                    if i < len(row_cells):
                        row_cells[i].text = cell_data
        else:
            doc.add_paragraph(content)
    
    def _add_testcase_content(self, doc, content):
        """Format test case content"""
        doc.add_heading('Test Cases', 1)
        
        # Split content into test cases
        sections = content.split('\n\n')
        for section in sections:
            if section.strip():
                if section.startswith('Test Case'):
                    doc.add_heading(section.split('\n')[0], 2)
                    doc.add_paragraph('\n'.join(section.split('\n')[1:]))
                else:
                    doc.add_paragraph(section)
    
    def _add_sequence_content(self, doc, content):
        """Format sequence diagram content"""
        doc.add_heading('Sequence Diagram', 1)
        doc.add_paragraph(content)