"""Module to parse the output of the Aider model."""

import re
from dataclasses import dataclass
from typing import List


@dataclass
class AiderOutputBlock:
    """A single block of output from the Aider model."""

    content: str
    type: str  # 'text', 'code', 'prompt', 'separator', 'token_info'
    is_code_block: bool = False
    is_prompt: bool = False
    is_separator: bool = False
    is_token_info: bool = False


class AiderOutputParser:
    """Output parser for Aider model."""

    def __init__(self) -> None:
        """Initialize the output parser."""
        self.code_block_pattern = re.compile(
            r"<<<<<<< SEARCH\n(.*?)\n=======\n(.*?)\n>>>>>>> REPLACE", re.DOTALL
        )
        self.prompt_pattern = re.compile(r"^> (.*)$", re.MULTILINE)
        self.separator_pattern = re.compile(r"^─+$", re.MULTILINE)
        self.token_info_pattern = re.compile(
            r"Tokens: (.*?) sent, (.*?) received\. Cost: (.*?) message, (.*?) session\."
        )
        self.aider_pattern = re.compile(r"\baider\b", re.IGNORECASE)

    def parse_output(self, output: str) -> List[AiderOutputBlock]:
        """Parse the output from the Aider model into blocks."""
        blocks = []
        current_block = []

        for line in output.split("\n"):

            line = line.rstrip("\r")

            line = re.sub(self.aider_pattern, "Dev Assist", line)

            if not line.strip():
                if current_block:
                    blocks.append(self._create_text_block("".join(current_block)))
                    current_block = []
                continue

            if "<<<<<<< SEARCH" in line:
                if current_block:
                    blocks.append(self._create_text_block("".join(current_block)))
                    current_block = []
                blocks.append(self._create_code_block(line))
                continue

            if line.startswith("> "):
                if current_block:
                    blocks.append(self._create_text_block("".join(current_block)))
                    current_block = []
                prompt_content = line[2:].strip()
                blocks.append(self._create_prompt_block(prompt_content))
                continue

            if re.match(self.separator_pattern, line):
                if current_block:
                    blocks.append(self._create_text_block("".join(current_block)))
                    current_block = []
                blocks.append(self._create_separator_block(line))
                continue

            if re.match(self.token_info_pattern, line):
                if current_block:
                    blocks.append(self._create_text_block("".join(current_block)))
                    current_block = []
                blocks.append(self._create_token_info_block(line))
                continue

            current_block.append(line + "\n")

        if current_block:
            blocks.append(self._create_text_block("".join(current_block)))

        return blocks

    def _create_text_block(self, content: str) -> AiderOutputBlock:
        return AiderOutputBlock(content=content, type="text")

    def _create_code_block(self, content: str) -> AiderOutputBlock:
        return AiderOutputBlock(content=content, type="code", is_code_block=True)

    def _create_prompt_block(self, content: str) -> AiderOutputBlock:
        return AiderOutputBlock(content=content, type="prompt", is_prompt=True)

    def _create_separator_block(self, content: str) -> AiderOutputBlock:
        return AiderOutputBlock(content=content, type="separator", is_separator=True)

    def _create_token_info_block(self, content: str) -> AiderOutputBlock:
        return AiderOutputBlock(content=content, type="token_info", is_token_info=True)