import re
import shutil
import sys

from pygments import highlight
from pygments.formatters import Terminal256Formatter
from pygments.lexers import PythonLexer


class FinalUnifiedStreamer:
    ANSI = {
        "reset": "\033[0m",
        "bold": "\033[1m",
        "italic": "\033[3m",
        "code_in": "\033[36m",
        "block_bg": "\033[48;5;235m",  # Dark background
        "header": "\033[1;95m",
        "quote": "\033[1;34m",
        "hr": "\033[90m",
        "clear_line": "\r\033[2K",  # Clear whole line and move to start
        "fill_bg": "\033[K",  # Fill background to end of line
    }

    def __init__(self):
        self.buffer = ""
        self.in_code_block = False
        self.current_code_line = ""
        self.at_line_start = True
        self.active_styles = set()
        self.lexer = PythonLexer()
        self.formatter = Terminal256Formatter(style="monokai")

    def apply_pygments(self, text):
        # Generate highlighted text
        highlighted = highlight(text, self.lexer, self.formatter)
        # Combine: BG Color + Highlighted Text + Fill to end of line
        return f"{self.ANSI['block_bg']}{highlighted.rstrip()}{self.ANSI['fill_bg']}{self.ANSI['reset']}"

    def process_buffer(self, final=False):
        while len(self.buffer) >= (1 if final else 15):
            # --- 1. CODE BLOCK STATE (Full Line BG) ---
            if self.in_code_block:
                if self.buffer.startswith("```"):
                    sys.stdout.write(f"{self.ANSI['clear_line']}")
                    self.in_code_block = False
                    self.current_code_line = ""
                    self.buffer = self.buffer[3:].lstrip("\n")
                    self.at_line_start = True
                    sys.stdout.write(self.ANSI["reset"] + "\n")
                    continue

                char = self.buffer[0]
                self.buffer = self.buffer[1:]

                if char == "\n":
                    hl = self.apply_pygments(self.current_code_line)
                    sys.stdout.write(f"{self.ANSI['clear_line']}{hl}\n")
                    self.current_code_line = ""
                else:
                    self.current_code_line += char
                    # Apply background and fill to end for every character drafted
                    sys.stdout.write(
                        f"\r{self.ANSI['block_bg']}{self.current_code_line}{self.ANSI['fill_bg']}"
                    )

                sys.stdout.flush()
                continue

            # --- 2. BLOCK LEVEL (Headers, HR, Quotes) ---
            if self.at_line_start:
                if self.buffer.startswith("```"):
                    self.in_code_block = True
                    match = re.match(r"^```[a-zA-Z]*\n?", self.buffer)
                    self.buffer = self.buffer[match.end() :]
                    continue

                if self.buffer.startswith("---"):
                    width = shutil.get_terminal_size().columns
                    sys.stdout.write(
                        f"{self.ANSI['hr']}{'─' * width}{self.ANSI['reset']}\n"
                    )
                    self.buffer = self.buffer[3:].lstrip("\n")
                    continue

                if self.buffer.startswith("> "):
                    sys.stdout.write(f"{self.ANSI['quote']}┃ {self.ANSI['reset']}")
                    self.buffer = self.buffer[2:]
                    continue

                if self.buffer.startswith("#"):
                    match = re.match(r"^(#+) ", self.buffer)
                    if match:
                        h_text = match.group()
                        sys.stdout.write(f"{self.ANSI['header']}{h_text}")
                        self.buffer = self.buffer[len(h_text) :]
                        self.at_line_start = False
                        continue

            # --- 3. INLINE FORMATTING ---
            if self.buffer.startswith("\\"):
                sys.stdout.write(self.buffer[1])
                self.buffer = self.buffer[2:]
            elif self.buffer.startswith("**"):
                self._toggle("bold")
                self.buffer = self.buffer[2:]
            elif self.buffer.startswith(("*", "_")) and not self.buffer.startswith(" "):
                self._toggle("italic")
                self.buffer = self.buffer[1:]
            elif self.buffer.startswith("`"):
                self._toggle("code_in")
                self.buffer = self.buffer[1:]
            else:
                char = self.buffer[0]
                sys.stdout.write(char)
                self.at_line_start = char == "\n"
                if self.at_line_start:
                    sys.stdout.write(self.ANSI["reset"])
                    self.active_styles.clear()
                self.buffer = self.buffer[1:]

            sys.stdout.flush()

    def _toggle(self, style):
        if style in self.active_styles:
            self.active_styles.remove(style)
        else:
            self.active_styles.add(style)
        sys.stdout.write(
            self.ANSI["reset"] + "".join(self.ANSI[s] for s in self.active_styles)
        )

    def stream(self, text):
        for char in text:
            self.buffer += char
            self.process_buffer()
        self.process_buffer(final=True)


parser = FinalUnifiedStreamer()

complex_test = """# 🚀 Extreme Streamer Test
---
This is a test of **nested bolding** and `inline code`.
Here is an escaped star: \\* and a backslash: \\\\.

> "The best way to predict the future is to create it."
> — Even with **bolding** inside quotes.

---

## 🐍 Python Block 1 (Logic)
```python
import time

class Tester:
    def __init__(self, name="Gemini"):
        self.name = name # Comment check
        self.data = [1, 2, 3]

    def run(self):
        print(f"Hello {self.name}!")
        return True
```
---
```
# Testing a different lexer if supported, or just raw formatting
echo "Installing dependencies..."
pip install pygments
mkdir -p ./output/logs
```
"""

import re
import shutil
import sys
import time

from pygments import highlight
from pygments.formatters import Terminal256Formatter
from pygments.lexers import PythonLexer


class FinalUnifiedStreamer:
    ANSI = {
        "reset": "\033[0m",
        "bold": "\033[1m",
        "italic": "\033[3m",
        "code_in": "\033[36m",
        "block_bg": "\033[48;5;235m",
        "header": "\033[1;95m",
        "quote": "\033[1;34m",
        "hr": "\033[90m",
        "clear_line": "\r\033[2K",
        "fill_bg": "\033[K",
    }

    def __init__(self):
        self.buffer = ""
        self.in_code_block = False
        self.current_code_line = ""
        self.at_line_start = True
        self.active_styles = set()

        # Word Wrapping State
        self.word_buffer = ""
        self.line_pos = 0
        self.term_width = shutil.get_terminal_size().columns

        # Pygments
        self.lexer = PythonLexer()
        self.formatter = Terminal256Formatter(style="monokai")

    def apply_pygments(self, text):
        highlighted = highlight(text, self.lexer, self.formatter)
        return f"{self.ANSI['block_bg']}{highlighted.rstrip()}{self.ANSI['fill_bg']}{self.ANSI['reset']}"

    def _flush_word(self, prefix=""):
        """Decides if the word fits or needs a newline before printing."""
        if not self.word_buffer:
            return

        # Check if word length exceeds remaining space
        if self.line_pos + len(self.word_buffer) >= self.term_width - 1:
            sys.stdout.write("\n" + prefix)
            self.line_pos = len(
                re.sub(r"\033\[[0-9;]*m", "", prefix)
            )  # Count only visible chars

        sys.stdout.write(self.word_buffer)
        self.line_pos += len(re.sub(r"\033\[[0-9;]*m", "", self.word_buffer))
        self.word_buffer = ""

    def process_buffer(self, final=False):
        # Window of 15 to handle markdown triggers
        while len(self.buffer) >= (1 if final else 15):
            # --- 1. CODE BLOCK (No word-wrap, uses truncation/scroll logic) ---
            if self.in_code_block:
                if self.buffer.startswith("```"):
                    sys.stdout.write(f"{self.ANSI['clear_line']}")
                    self.in_code_block = False
                    self.current_code_line = ""
                    self.buffer = self.buffer[3:].lstrip("\n")
                    self.at_line_start = True
                    sys.stdout.write(self.ANSI["reset"] + "\n")
                    self.line_pos = 0
                    continue

                char = self.buffer[0]
                self.buffer = self.buffer[1:]
                if char == "\n":
                    hl = self.apply_pygments(self.current_code_line)
                    sys.stdout.write(f"{self.ANSI['clear_line']}{hl}\n")
                    self.current_code_line = ""
                else:
                    self.current_code_line += char
                    # Fill background as we type
                    sys.stdout.write(
                        f"\r{self.ANSI['block_bg']}{self.current_code_line}{self.ANSI['fill_bg']}"
                    )
                sys.stdout.flush()
                continue

            # --- 2. BLOCK LEVEL ---
            if self.at_line_start:
                if self.buffer.startswith("```"):
                    self.in_code_block = True
                    match = re.match(r"^```[a-zA-Z]*\n?", self.buffer)
                    self.buffer = self.buffer[match.end() :]
                    continue

                if self.buffer.startswith("---"):
                    sys.stdout.write(
                        f"{self.ANSI['hr']}{'─' * self.term_width}{self.ANSI['reset']}\n"
                    )
                    self.buffer = self.buffer[3:].lstrip("\n")
                    self.line_pos = 0
                    continue

                if self.buffer.startswith("#"):
                    match = re.match(r"^(#+) ", self.buffer)
                    if match:
                        text = f"{self.ANSI['header']}{match.group()}"
                        self.word_buffer = text
                        self._flush_word()
                        self.buffer = self.buffer[len(match.group()) :]
                        self.at_line_start = False
                        continue

            # --- 3. INLINE & SMART WRAP ---
            prefix = (
                f"{self.ANSI['quote']}┃ {self.ANSI['reset']}"
                if "> " in self.buffer
                or self.at_line_start
                and self.buffer.startswith("> ")
                else ""
            )
            if self.at_line_start and self.buffer.startswith("> "):
                sys.stdout.write(prefix)
                self.line_pos = 2
                self.buffer = self.buffer[2:]

            char = self.buffer[0]

            if char == "\\":  # Escape
                self.word_buffer += self.buffer[1]
                self.buffer = self.buffer[2:]
            elif self.buffer.startswith("**"):
                self.word_buffer += self.ANSI["reset"] + (
                    self.ANSI["bold"] if "bold" not in self.active_styles else ""
                )
                self._toggle_state("bold")
                self.buffer = self.buffer[2:]
            elif self.buffer.startswith("`"):
                self.word_buffer += self.ANSI["reset"] + (
                    self.ANSI["code_in"] if "code_in" not in self.active_styles else ""
                )
                self._toggle_state("code_in")
                self.buffer = self.buffer[1:]
            elif char in (" ", "\n"):
                self._flush_word(prefix=prefix)
                if char == "\n":
                    sys.stdout.write("\n")
                    self.line_pos = 0
                    self.at_line_start = True
                    sys.stdout.write(self.ANSI["reset"])
                    self.active_styles.clear()
                else:
                    sys.stdout.write(" ")
                    self.line_pos += 1
                self.buffer = self.buffer[1:]
            else:
                self.word_buffer += char
                self.buffer = self.buffer[1:]

            sys.stdout.flush()

    def _toggle_state(self, style):
        if style in self.active_styles:
            self.active_styles.remove(style)
        else:
            self.active_styles.add(style)

    def stream(self, text):
        self.term_width = shutil.get_terminal_size().columns
        for char in text:
            self.buffer += char
            self.process_buffer()
            time.sleep(0.005)
        self.process_buffer(final=True)


parser = FinalUnifiedStreamer()
# parser.stream(complex_test)
