"""
Legal Writing Style Engine
==========================
Implements:
1. Chicago Manual of Style (17th ed.) for legal writing
2. Plain English for Lawyers principles (Richard Wydick)

Key principles applied:
- Active voice over passive
- Short sentences (average 20-25 words)
- Concrete subjects and verbs
- No nominalizations (turn nouns back into verbs)
- Front-load key information
- Eliminate throat-clearing phrases
- Use "the court held" not "it was held by the court"
"""

import re
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
from enum import Enum


class WritingIssue(Enum):
    """Types of writing issues to flag."""
    PASSIVE_VOICE = "passive_voice"
    NOMINALIZATION = "nominalization"
    WORDY_PHRASE = "wordy_phrase"
    THROAT_CLEARING = "throat_clearing"
    LONG_SENTENCE = "long_sentence"
    LEGALESE = "legalese"
    REDUNDANT_PAIR = "redundant_pair"
    BURIED_VERB = "buried_verb"


@dataclass
class WritingSuggestion:
    """A suggested improvement to text."""
    original: str
    suggested: str
    issue_type: WritingIssue
    explanation: str
    position: Tuple[int, int] = (0, 0)  # start, end


# ============================================================================
# PLAIN ENGLISH PATTERNS
# ============================================================================

# Passive voice indicators
PASSIVE_PATTERNS = [
    (r'\b(was|were|been|being|is|are|am)\s+([\w]+ed)\b', 'passive_voice'),
    (r'\b(was|were|been|being|is|are)\s+([\w]+en)\b', 'passive_voice'),
    (r'\bit\s+(?:was|is)\s+(?:held|found|determined|concluded|decided)\b', 'passive_voice'),
]

# Nominalizations - nouns that should be verbs
NOMINALIZATIONS = {
    'consideration': 'consider',
    'determination': 'determine',
    'authorization': 'authorize',
    'investigation': 'investigate',
    'implementation': 'implement',
    'utilization': 'use',
    'examination': 'examine',
    'termination': 'terminate/end',
    'violation': 'violate',
    'continuation': 'continue',
    'notification': 'notify',
    'recommendation': 'recommend',
    'allocation': 'allocate',
    'participation': 'participate',
    'calculation': 'calculate',
    'interpretation': 'interpret',
    'application': 'apply',
    'modification': 'modify',
    'compensation': 'compensate',
    'acquisition': 'acquire',
    'assumption': 'assume',
    'conclusion': 'conclude',
    'decision': 'decide',
    'agreement': 'agree',
    'requirement': 'require',
    'establishment': 'establish',
    'achievement': 'achieve',
    'assessment': 'assess',
    'commencement': 'begin/start',
    'performance': 'perform',
    'occurrence': 'occur',
    'assistance': 'assist/help',
    'compliance': 'comply',
    'reliance': 'rely',
    'existence': 'exist',
    'maintenance': 'maintain',
    'furtherance': 'further',
    'continuance': 'continue',
    'observance': 'observe',
}

# Wordy phrases to simplify
WORDY_PHRASES = {
    'at the present time': 'now',
    'at this point in time': 'now',
    'at that point in time': 'then',
    'in the event that': 'if',
    'in the event of': 'if',
    'in order to': 'to',
    'for the purpose of': 'to',
    'with respect to': 'about/regarding',
    'in regard to': 'about/regarding',
    'with regard to': 'about/regarding',
    'in connection with': 'about/regarding',
    'in relation to': 'about/regarding',
    'with reference to': 'about/regarding',
    'pertaining to': 'about/regarding',
    'relating to': 'about/regarding',
    'as to': 'about',
    'by means of': 'by/with',
    'by virtue of': 'by/under',
    'by reason of': 'because of',
    'on the grounds that': 'because',
    'due to the fact that': 'because',
    'owing to the fact that': 'because',
    'in light of the fact that': 'because',
    'given the fact that': 'because',
    'considering the fact that': 'because',
    'inasmuch as': 'because/since',
    'on the basis of': 'based on',
    'on the part of': 'by',
    'a number of': 'several/many',
    'a majority of': 'most',
    'a sufficient number of': 'enough',
    'an adequate number of': 'enough',
    'during the course of': 'during',
    'during the time that': 'while',
    'subsequent to': 'after',
    'prior to': 'before',
    'in advance of': 'before',
    'until such time as': 'until',
    'at such time as': 'when',
    'in the absence of': 'without',
    'in the nature of': 'like',
    'of a [adjective] nature': '[adjective]',
    'of a [adjective] character': '[adjective]',
    'there is no doubt that': '[omit]',
    'it is clear that': '[omit]',
    'it is obvious that': '[omit]',
    'it is apparent that': '[omit]',
    'it is important to note that': '[omit]',
    'it should be noted that': '[omit]',
    'it must be remembered that': '[omit]',
    'it goes without saying': '[omit]',
    'needless to say': '[omit]',
    'as a matter of fact': '[omit]',
    'the fact that': 'that',
    'despite the fact that': 'although/even though',
    'notwithstanding the fact that': 'although/even though',
    'regardless of the fact that': 'although/even though',
    'is able to': 'can',
    'is unable to': 'cannot',
    'has the ability to': 'can',
    'has the capacity to': 'can',
    'is in a position to': 'can',
    'make a decision': 'decide',
    'make a determination': 'determine',
    'make an argument': 'argue',
    'make an assumption': 'assume',
    'give consideration to': 'consider',
    'have knowledge of': 'know',
    'be in possession of': 'have',
    'enter into an agreement': 'agree',
    'file a complaint': 'sue/complain',
    'render a decision': 'decide',
    'reach a conclusion': 'conclude',
    'achieve a resolution': 'resolve',
}

# Throat-clearing phrases (usually can be deleted)
THROAT_CLEARING = [
    r'^[Ii]t is (?:important|significant|noteworthy|interesting|relevant) (?:to note |that )',
    r'^[Aa]s (?:previously |earlier )?(?:stated|mentioned|noted|discussed|indicated)',
    r'^[Ii]n (?:this|the present) case,?',
    r'^[Ii]t (?:should|must|is important to) be (?:noted|emphasized|stressed|remembered) that',
    r'^[Tt]here (?:is|are|can be) no (?:doubt|question|dispute) that',
    r'^[Cc]learly,',
    r'^[Oo]bviously,',
    r'^[Ii]ndeed,',
    r'^[Ii]n fact,',
    r'^[Aa]s a matter of fact,',
    r'^[Ii]t is (?:well[- ])?(?:settled|established) (?:law )?that',
]

# Redundant legal pairs (pick one, not both)
REDUNDANT_PAIRS = [
    ('null', 'void'),
    ('cease', 'desist'),
    ('aid', 'abet'),
    ('each', 'every'),
    ('full', 'complete'),
    ('true', 'correct'),
    ('alter', 'change'),
    ('make', 'enter into'),
    ('authorize', 'empower'),
    ('by', 'through'),
    ('final', 'conclusive'),
    ('from', 'after'),
    ('give', 'grant'),
    ('have', 'hold'),
    ('keep', 'maintain'),
    ('let', 'permit'),
    ('made', 'entered into'),
    ('new', 'novel'),
    ('null', 'void'),
    ('over', 'above'),
    ('save', 'except'),
    ('sole', 'exclusive'),
    ('suffer', 'permit'),
    ('then', 'in that event'),
    ('unless', 'except'),
    ('void', 'of no effect'),
]

# Legalese to replace
LEGALESE_REPLACEMENTS = {
    'aforementioned': 'this/that/the',
    'aforesaid': 'this/that/the',
    'herein': 'in this document/here',
    'hereinafter': 'below/later',
    'hereinbefore': 'above/earlier',
    'heretofore': 'until now',
    'hereunder': 'below/under this',
    'herewith': 'with this',
    'thereby': 'by that/so',
    'therefrom': 'from that',
    'therein': 'in that/there',
    'thereof': 'of that',
    'thereon': 'on that',
    'thereto': 'to that',
    'thereunder': 'under that',
    'therewith': 'with that',
    'whereby': 'by which',
    'wherein': 'in which/where',
    'whereof': 'of which',
    'whereon': 'on which',
    'whosoever': 'whoever',
    'whatsoever': 'whatever/any',
    'whensoever': 'whenever',
    'wheresoever': 'wherever',
    'howsoever': 'however',
    'said': 'the/this/that',
    'same': 'it/them/the [noun]',
    'such': 'the/this/that',
    'forthwith': 'immediately',
    'witnesseth': '[omit]',
    'whereas': '[omit or use "because"]',
    'now therefore': '[omit]',
    'know all men by these presents': '[omit]',
}


# ============================================================================
# CHICAGO MANUAL OF STYLE - LEGAL CITATIONS
# ============================================================================

# Chicago style uses footnotes for citations, but legal briefs use inline
# Key Chicago principles for legal writing:
# 1. Spell out numbers one through one hundred
# 2. Use serial comma (Oxford comma)
# 3. Titles in italics for case names
# 4. Specific date formats

CHICAGO_NUMBER_RULES = {
    'spell_out_under': 100,  # Spell out numbers under 100
    'use_numerals_for': ['dates', 'page numbers', 'citations', 'percentages', 'money'],
}


# ============================================================================
# ANALYSIS FUNCTIONS
# ============================================================================

def analyze_writing(text: str) -> List[WritingSuggestion]:
    """
    Analyze text for Plain English and style issues.

    Args:
        text: Text to analyze

    Returns:
        List of WritingSuggestion objects
    """
    suggestions = []

    # Check for passive voice
    suggestions.extend(find_passive_voice(text))

    # Check for nominalizations
    suggestions.extend(find_nominalizations(text))

    # Check for wordy phrases
    suggestions.extend(find_wordy_phrases(text))

    # Check for throat-clearing
    suggestions.extend(find_throat_clearing(text))

    # Check for long sentences
    suggestions.extend(find_long_sentences(text))

    # Check for legalese
    suggestions.extend(find_legalese(text))

    # Check for redundant pairs
    suggestions.extend(find_redundant_pairs(text))

    return suggestions


def find_passive_voice(text: str) -> List[WritingSuggestion]:
    """Find passive voice constructions."""
    suggestions = []

    for pattern, _ in PASSIVE_PATTERNS:
        for match in re.finditer(pattern, text, re.IGNORECASE):
            suggestions.append(WritingSuggestion(
                original=match.group(0),
                suggested="[Rewrite in active voice]",
                issue_type=WritingIssue.PASSIVE_VOICE,
                explanation="Use active voice: identify who performed the action and make them the subject.",
                position=(match.start(), match.end())
            ))

    return suggestions


def find_nominalizations(text: str) -> List[WritingSuggestion]:
    """Find nominalizations (buried verbs)."""
    suggestions = []

    for noun, verb in NOMINALIZATIONS.items():
        pattern = rf'\b{noun}\b'
        for match in re.finditer(pattern, text, re.IGNORECASE):
            suggestions.append(WritingSuggestion(
                original=match.group(0),
                suggested=f"Consider using '{verb}' instead",
                issue_type=WritingIssue.NOMINALIZATION,
                explanation=f"'{noun}' is a nominalization. Use the verb '{verb}' for stronger, clearer writing.",
                position=(match.start(), match.end())
            ))

    return suggestions


def find_wordy_phrases(text: str) -> List[WritingSuggestion]:
    """Find wordy phrases that can be simplified."""
    suggestions = []
    text_lower = text.lower()

    for wordy, simple in WORDY_PHRASES.items():
        if wordy in text_lower:
            # Find actual position in original text
            pattern = re.escape(wordy)
            for match in re.finditer(pattern, text, re.IGNORECASE):
                suggestions.append(WritingSuggestion(
                    original=match.group(0),
                    suggested=simple,
                    issue_type=WritingIssue.WORDY_PHRASE,
                    explanation=f"Replace '{wordy}' with '{simple}' for conciseness.",
                    position=(match.start(), match.end())
                ))

    return suggestions


def find_throat_clearing(text: str) -> List[WritingSuggestion]:
    """Find throat-clearing phrases at sentence starts."""
    suggestions = []

    # Split into sentences
    sentences = re.split(r'(?<=[.!?])\s+', text)

    for sentence in sentences:
        for pattern in THROAT_CLEARING:
            match = re.match(pattern, sentence)
            if match:
                suggestions.append(WritingSuggestion(
                    original=match.group(0),
                    suggested="[Delete or revise]",
                    issue_type=WritingIssue.THROAT_CLEARING,
                    explanation="This phrase delays the main point. Start with your key information.",
                    position=(0, 0)  # Position within sentence
                ))
                break

    return suggestions


def find_long_sentences(text: str, max_words: int = 35) -> List[WritingSuggestion]:
    """Find sentences that are too long."""
    suggestions = []

    # Split into sentences
    sentences = re.split(r'(?<=[.!?])\s+', text)

    for sentence in sentences:
        word_count = len(sentence.split())
        if word_count > max_words:
            suggestions.append(WritingSuggestion(
                original=sentence[:50] + "..." if len(sentence) > 50 else sentence,
                suggested=f"[Break into shorter sentences - currently {word_count} words]",
                issue_type=WritingIssue.LONG_SENTENCE,
                explanation=f"Sentence has {word_count} words. Aim for 20-25 words average. Consider breaking at natural pause points.",
                position=(0, 0)
            ))

    return suggestions


def find_legalese(text: str) -> List[WritingSuggestion]:
    """Find archaic legal terms."""
    suggestions = []

    for legalese, replacement in LEGALESE_REPLACEMENTS.items():
        pattern = rf'\b{legalese}\b'
        for match in re.finditer(pattern, text, re.IGNORECASE):
            suggestions.append(WritingSuggestion(
                original=match.group(0),
                suggested=replacement,
                issue_type=WritingIssue.LEGALESE,
                explanation=f"'{legalese}' is archaic legalese. Use '{replacement}' instead.",
                position=(match.start(), match.end())
            ))

    return suggestions


def find_redundant_pairs(text: str) -> List[WritingSuggestion]:
    """Find redundant word pairs."""
    suggestions = []
    text_lower = text.lower()

    for word1, word2 in REDUNDANT_PAIRS:
        # Check if both words appear near each other
        pattern = rf'\b{word1}\b.{{0,10}}\b{word2}\b|\b{word2}\b.{{0,10}}\b{word1}\b'
        for match in re.finditer(pattern, text_lower):
            suggestions.append(WritingSuggestion(
                original=match.group(0),
                suggested=f"Use just '{word1}' or '{word2}', not both",
                issue_type=WritingIssue.REDUNDANT_PAIR,
                explanation="This is a redundant pair. Choose one word.",
                position=(match.start(), match.end())
            ))

    return suggestions


# ============================================================================
# TEXT IMPROVEMENT FUNCTIONS
# ============================================================================

def improve_text(text: str, apply_fixes: bool = False) -> Tuple[str, List[WritingSuggestion]]:
    """
    Analyze and optionally improve text.

    Args:
        text: Text to analyze
        apply_fixes: If True, apply automatic fixes for simple issues

    Returns:
        Tuple of (improved_text, suggestions)
    """
    suggestions = analyze_writing(text)
    improved = text

    if apply_fixes:
        # Apply simple replacements (wordy phrases only - these are multi-word and unambiguous)
        for wordy, simple in WORDY_PHRASES.items():
            if simple != '[omit]':
                improved = re.sub(
                    re.escape(wordy),
                    simple.split('/')[0],  # Use first suggestion
                    improved,
                    flags=re.IGNORECASE
                )

        # For legalese, only auto-apply SAFE replacements
        # Skip context-dependent single words that could have normal usage
        # (e.g., "said" can mean past tense of "say", not just legal "the said document")
        UNSAFE_LEGALESE = {
            'said',  # "he said" vs "said defendant"
            'same',  # replacement "it/them/the [noun]" has placeholder
            'such',  # often legitimate: "no such evidence"
        }

        # Safe legalese: multi-word or archaic terms with clear replacements
        # Words that replace "the X" entirely (determiner-inclusive)
        DETERMINER_LEGALESE = {
            'aforementioned': 'this',
            'aforesaid': 'this',
        }

        # Words that are simple replacements (don't affect preceding article)
        SIMPLE_LEGALESE = {
            'herein': 'here',
            'hereinafter': 'below',
            'hereinbefore': 'above',
            'heretofore': 'until now',
            'hereunder': 'under this',
            'herewith': 'with this',
            'forthwith': 'immediately',
            'whosoever': 'whoever',
            'whatsoever': 'any',
            'whensoever': 'whenever',
            'wheresoever': 'wherever',
            'howsoever': 'however',
        }

        # Handle determiner + legalese patterns (e.g., "the aforementioned" → "this")
        for legalese, replacement in DETERMINER_LEGALESE.items():
            # Replace "the aforementioned" with "this" (remove redundant article)
            improved = re.sub(
                rf'\b(?:the|an?)\s+{legalese}\b',
                replacement,
                improved,
                flags=re.IGNORECASE
            )
            # Also replace standalone occurrences
            improved = re.sub(
                rf'\b{legalese}\b',
                replacement,
                improved,
                flags=re.IGNORECASE
            )

        # Handle simple replacements
        for legalese, replacement in SIMPLE_LEGALESE.items():
            improved = re.sub(
                rf'\b{legalese}\b',
                replacement,
                improved,
                flags=re.IGNORECASE
            )

    return improved, suggestions


def format_suggestions_report(suggestions: List[WritingSuggestion]) -> str:
    """Format suggestions as a readable report."""
    if not suggestions:
        return "No writing issues found. Text follows Plain English principles."

    output = []
    output.append("WRITING ANALYSIS REPORT")
    output.append("=" * 70)
    output.append(f"\nTotal issues found: {len(suggestions)}\n")

    # Group by type
    by_type: Dict[WritingIssue, List[WritingSuggestion]] = {}
    for s in suggestions:
        if s.issue_type not in by_type:
            by_type[s.issue_type] = []
        by_type[s.issue_type].append(s)

    # Report each type
    type_labels = {
        WritingIssue.PASSIVE_VOICE: "Passive Voice",
        WritingIssue.NOMINALIZATION: "Nominalizations (Buried Verbs)",
        WritingIssue.WORDY_PHRASE: "Wordy Phrases",
        WritingIssue.THROAT_CLEARING: "Throat-Clearing",
        WritingIssue.LONG_SENTENCE: "Long Sentences",
        WritingIssue.LEGALESE: "Archaic Legalese",
        WritingIssue.REDUNDANT_PAIR: "Redundant Pairs",
        WritingIssue.BURIED_VERB: "Buried Verbs",
    }

    for issue_type, items in by_type.items():
        output.append(f"\n{type_labels.get(issue_type, issue_type.value).upper()} ({len(items)} found)")
        output.append("-" * 70)

        for item in items[:10]:  # Limit to 10 per category
            output.append(f"\n  Found: \"{item.original}\"")
            output.append(f"  Suggest: {item.suggested}")
            output.append(f"  Why: {item.explanation}")

        if len(items) > 10:
            output.append(f"\n  ... and {len(items) - 10} more")

    return '\n'.join(output)


# ============================================================================
# LEGAL WRITING TEMPLATES
# ============================================================================

def generate_rule_statement_template(rule: str, cases: List[str]) -> str:
    """
    Generate a properly structured rule statement.

    Uses Plain English principles:
    - Lead with the rule
    - Short sentences
    - Active voice
    """
    template = []

    # Main rule first
    template.append(rule)

    # Supporting cases
    if cases:
        template.append(f"See {cases[0]}.")
        if len(cases) > 1:
            template.append(f"Accord {'; '.join(cases[1:])}.")

    return ' '.join(template)


def generate_application_paragraph(
    rule_element: str,
    facts: str,
    conclusion: str
) -> str:
    """
    Generate an application paragraph using Plain English structure.

    Structure:
    1. Topic sentence stating conclusion
    2. Rule element (brief)
    3. Facts that satisfy/don't satisfy
    4. Conclusion
    """
    template = f"""{conclusion} {rule_element} Here, {facts} Therefore, {conclusion.lower()}"""

    return template


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

def analyze_writing_style(text: str) -> str:
    """
    Main entry point for MCP server.

    Args:
        text: Text to analyze

    Returns:
        Formatted report of writing issues and suggestions
    """
    suggestions = analyze_writing(text)
    return format_suggestions_report(suggestions)


def improve_legal_writing(text: str) -> str:
    """
    Improve text by applying automatic fixes.

    Args:
        text: Text to improve

    Returns:
        Improved text with report
    """
    improved, suggestions = improve_text(text, apply_fixes=True)

    output = []
    output.append("IMPROVED TEXT")
    output.append("=" * 70)
    output.append(improved)
    output.append("\n" + "=" * 70)
    output.append(format_suggestions_report(suggestions))

    return '\n'.join(output)
