import nltk
from rich.console import Console, ConsoleOptions, RenderResult
from rich.markdown import Markdown, Paragraph
from rich.text import Text

POS_COLORS = {
    "NN": "cyan bold",   
    "NNS": "cyan bold",  
    "NNP": "blue bold",  
    "NNPS": "blue bold", 
    "VB": "green",       
    "VBD": "green",      
    "VBG": "green",      
    "VBN": "green",      
    "VBP": "green",      
    "VBZ": "green",      
    "JJ": "magenta",     
    "JJR": "magenta",    
    "JJS": "magenta",    
    "RB": "yellow italic",      
    "RBR": "yellow italic",     
    "RBS": "yellow italic",     
    "IN": "dim",         
    "DT": "dim",         
}

class POSParagraph(Paragraph):
    def __rich_console__(self, console: Console, options: ConsoleOptions) -> RenderResult:
        text_content = self.text.plain
        
        try:
            # defensive programming: ensure nltk is available and works
            # basic split fallback if text is too long or complex could be added here
            tokens = nltk.word_tokenize(text_content)
            tagged = nltk.pos_tag(tokens)
            
            styled_text = Text()
            for word, tag in tagged:
                style = POS_COLORS.get(tag, None)
                styled_text.append(word, style=style)
                styled_text.append(" ") # simple spacing reconstruction
            
            yield styled_text

        except (LookupError, OSError, UnicodeDecodeError) as e:
            # if nltk fails (missing data, file system error), return plain text
            yield Text(text_content)
        except Exception:
            # catch-all for any other nlp crashes
            yield Text(text_content)

class POSMarkdown(Markdown):
    elements = Markdown.elements.copy()
    elements["paragraph"] = POSParagraph
