# services/docx_agent.py

import logging
import os
import uuid
import ast
from docx import Document
from docx.shared import Inches
from typing import List, Any, Dict

logger = logging.getLogger(__name__)
TEMP_DIR = "temp_docs"


class DocxAgent:
    """Agent to handle .docx file generation."""

    def __init__(self):
        if not os.path.exists(TEMP_DIR):
            os.makedirs(TEMP_DIR)
            logger.info("Created temporary directory for documents: %s", TEMP_DIR)

    def create_document_from_sql_result(self, sql_result: str, title: str) -> Dict[str, str]:
        """
        Creates a .docx file from a raw SQL result string.

        Args:
            sql_result: The string representation of the SQL result (e.g., "[('John Doe', 30), ('Jane Smith', 25)]").
            title: The title for the document.

        Returns:
            A dictionary containing the file path and a user-friendly message.
        """
        try:
            # Safely parse the string representation of the list of tuples
            data = ast.literal_eval(sql_result)
            if not isinstance(data, list) or (data and not isinstance(data[0], tuple)):
                raise ValueError("SQL result is not in the expected format (list of tuples).")
        except (ValueError, SyntaxError) as e:
            logger.error("Could not parse SQL result for docx generation: %s", e)
            # If parsing fails, create a document with the raw text
            return self._create_document_with_raw_text(sql_result, title)

        if not data:
            return self._create_document_with_raw_text("No data found for this query.", title)

        document = Document()
        document.add_heading(title, level=1)

        # Assuming the first row can be used for headers, though they are not explicitly provided
        # This part might need enhancement if you have headers
        num_columns = len(data[0])
        table = document.add_table(rows=1, cols=num_columns)
        table.style = 'Table Grid'

        # Placeholder for headers - you might want to generate these with another LLM call
        hdr_cells = table.rows[0].cells
        for i in range(num_columns):
            hdr_cells[i].text = f'Column {i + 1}'

        # Add data rows
        for row_data in data:
            row_cells = table.add_row().cells
            for i, cell_value in enumerate(row_data):
                row_cells[i].text = str(cell_value)

        file_id = str(uuid.uuid4())
        filename = f"{file_id}.docx"
        filepath = os.path.join(TEMP_DIR, filename)

        document.save(filepath)
        logger.info("Successfully created docx file: %s", filepath)

        return {
            "filepath": filepath,
            "filename": f"{title.replace(' ', '_')}.docx",
            "message": f"Your document '{title}' has been generated."
        }

    def _create_document_with_raw_text(self, text: str, title: str) -> Dict[str, str]:
        """Fallback to create a document with raw text if structured parsing fails."""
        document = Document()
        document.add_heading(title, level=1)
        document.add_paragraph(text)

        file_id = str(uuid.uuid4())
        filename = f"{file_id}.docx"
        filepath = os.path.join(TEMP_DIR, filename)

        document.save(filepath)
        logger.info("Created docx file with raw text fallback: %s", filepath)

        return {
            "filepath": filepath,
            "filename": f"{title.replace(' ', '_')}.docx",
            "message": f"Your document '{title}' has been generated with the available data."
        }