"""
Legal Document Formatter
========================
Generates properly formatted legal documents (.docx) with:
- 12pt Times New Roman
- No heading styles (just bold text)
- All black text
- Single-spaced with single blank line between paragraphs

For use with memo drafts, briefs, and other legal writing output.
"""

import re
import os
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
from pathlib import Path

# Try to import python-docx for Word document generation
try:
    from docx import Document
    from docx.shared import Pt, Inches, RGBColor
    from docx.enum.text import WD_ALIGN_PARAGRAPH, WD_LINE_SPACING
    from docx.enum.style import WD_STYLE_TYPE
    DOCX_AVAILABLE = True
except ImportError:
    DOCX_AVAILABLE = False


def get_desktop_path() -> Path:
    """
    Get the correct Desktop path, handling OneDrive folder redirection.

    Returns:
        Path to Desktop folder
    """
    # Check OneDrive Desktop first (common on Windows with OneDrive)
    onedrive_desktop = Path.home() / "OneDrive" / "Desktop"
    if onedrive_desktop.exists():
        return onedrive_desktop

    # Fall back to regular Desktop
    regular_desktop = Path.home() / "Desktop"
    if regular_desktop.exists():
        return regular_desktop

    # If neither exists, try to create regular Desktop
    regular_desktop.mkdir(parents=True, exist_ok=True)
    return regular_desktop


@dataclass
class LegalDocumentStyle:
    """Style settings for legal documents."""
    font_name: str = "Times New Roman"
    font_size: int = 12
    line_spacing: float = 1.0  # Single-spaced
    paragraph_spacing_after: int = 12  # Points (roughly one blank line)
    margin_top: float = 1.0  # Inches
    margin_bottom: float = 1.0
    margin_left: float = 1.0
    margin_right: float = 1.0
    text_color: Tuple[int, int, int] = (0, 0, 0)  # Black


DEFAULT_STYLE = LegalDocumentStyle()


def create_legal_document(
    content: Dict[str, str],
    output_path: str,
    title: Optional[str] = None,
    style: LegalDocumentStyle = DEFAULT_STYLE
) -> str:
    """
    Create a formatted Word document from content sections.

    Args:
        content: Dictionary mapping section name -> section content
        output_path: Path to save the .docx file
        title: Optional document title
        style: Style settings to apply

    Returns:
        Path to the created document, or error message
    """
    if not DOCX_AVAILABLE:
        return format_as_plain_text(content, title)

    doc = Document()

    # Set document margins
    for section in doc.sections:
        section.top_margin = Inches(style.margin_top)
        section.bottom_margin = Inches(style.margin_bottom)
        section.left_margin = Inches(style.margin_left)
        section.right_margin = Inches(style.margin_right)

    # Add title if provided (bold, centered, but NOT a heading style)
    if title:
        title_para = doc.add_paragraph()
        title_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
        title_run = title_para.add_run(title.upper())
        title_run.bold = True
        title_run.font.name = style.font_name
        title_run.font.size = Pt(style.font_size)
        title_run.font.color.rgb = RGBColor(*style.text_color)
        # Single spacing
        title_para.paragraph_format.line_spacing_rule = WD_LINE_SPACING.SINGLE
        title_para.paragraph_format.space_after = Pt(style.paragraph_spacing_after)

    # Add each section
    for section_name, section_content in content.items():
        # Section header (bold, but NOT a heading style)
        if section_name and section_name.strip():
            header_para = doc.add_paragraph()
            header_run = header_para.add_run(section_name.upper())
            header_run.bold = True
            header_run.font.name = style.font_name
            header_run.font.size = Pt(style.font_size)
            header_run.font.color.rgb = RGBColor(*style.text_color)
            header_para.paragraph_format.line_spacing_rule = WD_LINE_SPACING.SINGLE
            header_para.paragraph_format.space_after = Pt(style.paragraph_spacing_after)

        # Section content - split by double newlines for paragraphs
        paragraphs = re.split(r'\n\n+', section_content.strip())

        for para_text in paragraphs:
            if not para_text.strip():
                continue

            para = doc.add_paragraph()

            # Handle inline formatting (bold with **)
            parts = re.split(r'(\*\*[^*]+\*\*)', para_text)

            for part in parts:
                if part.startswith('**') and part.endswith('**'):
                    # Bold text
                    run = para.add_run(part[2:-2])
                    run.bold = True
                else:
                    run = para.add_run(part)

                run.font.name = style.font_name
                run.font.size = Pt(style.font_size)
                run.font.color.rgb = RGBColor(*style.text_color)

            # Apply paragraph formatting
            para.paragraph_format.line_spacing_rule = WD_LINE_SPACING.SINGLE
            para.paragraph_format.space_after = Pt(style.paragraph_spacing_after)
            para.paragraph_format.space_before = Pt(0)

    # Save document
    doc.save(output_path)
    return f"Document saved to: {output_path}"


def format_as_plain_text(content: Dict[str, str], title: Optional[str] = None) -> str:
    """
    Format content as plain text when docx is not available.
    Uses formatting that can be easily copied into Word.

    Args:
        content: Dictionary mapping section name -> section content
        title: Optional document title

    Returns:
        Formatted plain text string
    """
    output = []

    if title:
        output.append(title.upper())
        output.append("")  # Blank line after title

    for section_name, section_content in content.items():
        if section_name and section_name.strip():
            output.append(section_name.upper())
            output.append("")  # Blank line after section header

        # Clean up content - normalize whitespace
        paragraphs = re.split(r'\n\n+', section_content.strip())
        for para in paragraphs:
            if para.strip():
                # Remove markdown-style bold markers for plain text
                clean_para = re.sub(r'\*\*([^*]+)\*\*', r'\1', para)
                # Normalize internal whitespace
                clean_para = ' '.join(clean_para.split())
                output.append(clean_para)
                output.append("")  # Blank line between paragraphs

    return '\n'.join(output)


def format_memo_for_export(memo_dict: Dict) -> Dict[str, str]:
    """
    Convert a memo draft dictionary to export-ready sections.

    Args:
        memo_dict: Dictionary with memo components

    Returns:
        Dictionary ready for create_legal_document()
    """
    sections = {}

    if memo_dict.get('question_presented'):
        sections['Question Presented'] = memo_dict['question_presented']

    if memo_dict.get('brief_answer'):
        sections['Brief Answer'] = memo_dict['brief_answer']

    if memo_dict.get('facts_template') or memo_dict.get('facts'):
        sections['Statement of Facts'] = memo_dict.get('facts_template') or memo_dict.get('facts', '')

    # Handle IRAC discussion
    if memo_dict.get('discussion_irac'):
        irac = memo_dict['discussion_irac']
        discussion_parts = []

        if irac.get('issue'):
            discussion_parts.append(irac['issue'])
        if irac.get('rule'):
            discussion_parts.append(irac['rule'])
        if irac.get('application'):
            discussion_parts.append(irac['application'])
        if irac.get('conclusion'):
            discussion_parts.append(irac['conclusion'])

        sections['Discussion'] = '\n\n'.join(discussion_parts)

    if memo_dict.get('conclusion'):
        sections['Conclusion'] = memo_dict['conclusion']

    return sections


def format_brief_section(
    section_type: str,
    content: str,
    citations: List[str] = None
) -> str:
    """
    Format a brief section with proper legal formatting.

    Args:
        section_type: Type of section (e.g., "argument", "facts")
        content: Raw content
        citations: Optional list of citations to include

    Returns:
        Formatted section text
    """
    output = []

    # Clean up content
    paragraphs = re.split(r'\n\n+', content.strip())

    for para in paragraphs:
        if para.strip():
            # Preserve any inline citations
            output.append(para.strip())

    return '\n\n'.join(output)


# ============================================================================
# MCP INTEGRATION FUNCTIONS
# ============================================================================

def export_memo_to_docx(
    memo_content: str,
    output_path: str = None,
    title: str = "LEGAL MEMORANDUM"
) -> str:
    """
    Export memo content to a properly formatted Word document.

    Args:
        memo_content: The formatted memo text (from generate_memo_draft)
        output_path: Where to save (default: Desktop)
        title: Document title

    Returns:
        Status message with file path or plain text if docx unavailable
    """
    if output_path is None:
        output_path = str(Path.home() / "Desktop" / "legal_memo.docx")

    # Parse the memo content into sections
    sections = parse_memo_sections(memo_content)

    if DOCX_AVAILABLE:
        return create_legal_document(sections, output_path, title)
    else:
        return format_as_plain_text(sections, title)


def parse_memo_sections(memo_text: str) -> Dict[str, str]:
    """
    Parse formatted memo text into sections.

    Args:
        memo_text: Raw memo text with section headers

    Returns:
        Dictionary of section name -> content
    """
    sections = {}

    # Common section patterns
    section_patterns = [
        r'(?:##\s*)?QUESTION PRESENTED\s*[-=]*\s*\n(.*?)(?=(?:##\s*)?(?:BRIEF ANSWER|STATEMENT|DISCUSSION|CONCLUSION|$))',
        r'(?:##\s*)?BRIEF ANSWER\s*[-=]*\s*\n(.*?)(?=(?:##\s*)?(?:STATEMENT|DISCUSSION|CONCLUSION|$))',
        r'(?:##\s*)?STATEMENT OF FACTS?\s*[-=]*\s*\n(.*?)(?=(?:##\s*)?(?:DISCUSSION|CONCLUSION|$))',
        r'(?:##\s*)?DISCUSSION\s*[-=]*\s*\n(.*?)(?=(?:##\s*)?CONCLUSION|$)',
        r'(?:##\s*)?CONCLUSION\s*[-=]*\s*\n(.*?)$',
    ]

    section_names = [
        'Question Presented',
        'Brief Answer',
        'Statement of Facts',
        'Discussion',
        'Conclusion'
    ]

    for pattern, name in zip(section_patterns, section_names):
        match = re.search(pattern, memo_text, re.IGNORECASE | re.DOTALL)
        if match:
            content = match.group(1).strip()
            # Clean up markdown formatting
            content = re.sub(r'^#+\s*', '', content, flags=re.MULTILINE)
            content = re.sub(r'^[-=]+\s*$', '', content, flags=re.MULTILINE)
            if content:
                sections[name] = content

    # If no sections found, treat entire text as content
    if not sections:
        sections['Content'] = memo_text

    return sections


def get_formatted_output(content: Dict[str, str], title: str = None) -> str:
    """
    Get formatted plain text output ready for copy-paste into Word.

    This is the primary function for MCP integration when the user
    wants formatted text rather than a file.

    Args:
        content: Section name -> content dictionary
        title: Optional title

    Returns:
        Formatted text with proper spacing
    """
    return format_as_plain_text(content, title)


# ============================================================================
# HELPER FOR CHECKING DOCX AVAILABILITY
# ============================================================================

def check_docx_support() -> Tuple[bool, str]:
    """
    Check if python-docx is available for Word document generation.

    Returns:
        Tuple of (is_available, message)
    """
    if DOCX_AVAILABLE:
        return True, "python-docx is available. Word documents can be generated."
    else:
        return False, (
            "python-docx is not installed. Install with: pip install python-docx\n"
            "Without it, formatted plain text will be provided instead."
        )
