import re
import logging

from dataclasses import dataclass
from itertools import chain

from parsimonious.nodes import (
    Node,
    NodeVisitor,
)
from parsimonious.grammar import Grammar

from maccarone.openai import CachedChatAPI

logger = logging.getLogger(__name__)

class Piece:
    pass

@dataclass
class PresentPiece(Piece):
    text: str

@dataclass
class MissingPiece(Piece):
    indent: str
    guidance: str
    inlined: str | None = None

grammar = Grammar(
    r"""
    maccarone = human_source maccarone_chunk*
    maccarone_chunk = snippet human_source?

    snippet = snippet_open (ai_source snippet_close)?
    snippet_open = snippet_open_single / snippet_open_multi
    snippet_open_single = guidance_open guidance_inner ">>" nl
    snippet_open_multi = guidance_open nl guidance_lines guidance_close
    snippet_close = ws "#<</>>" nl

    guidance_open = ws "#<<"
    guidance_close = ws "#>>" nl
    guidance_line = ws "#" guidance_inner nl
    guidance_lines = guidance_line+
    guidance_inner = ~"((?!>>).)*"

    human_source = source_line*
    ai_source = source_line*
    source_line = !(guidance_open / guidance_close / snippet_close) ws ~".*" nl?

    ws = ~"[ \t]*"
    nl = ws ~"[\r\n]"
    """
)

@dataclass
class GuidanceOpen:
    indent: str

@dataclass
class Guidance:
    text: str

@dataclass
class SnippetOpen:
    indent: str
    guidance: str

class RawSourceVisitor(NodeVisitor):
    def generic_visit(self, node: Node, visited_children: list[Node]):
        return visited_children or node

    def visit_maccarone(self, node: Node, visited_children: list):
        (first_source, chunks) = visited_children

        return [first_source] + list(chain(*chunks))

    def visit_maccarone_chunk(self, node: Node, visited_children: list):
        (snippet, source) = visited_children

        if isinstance(source, list):
            source_list = source
        else:
            source_list = []

        return [snippet] + source_list

    def visit_snippet(self, node: Node, visited_children: list):
        (snippet_open, quantified_source) = visited_children

        if isinstance(quantified_source, list):
            ((source, _),) = quantified_source
        else:
            source = None

        return MissingPiece(
            indent=snippet_open.indent,
            guidance=snippet_open.guidance,
            inlined=source,
        )

    def visit_snippet_open(self, node: Node, visited_children: list):
        (single_or_multi,) = visited_children

        return single_or_multi

    def visit_snippet_open_single(self, node: Node, visited_children: list):
        (guidance_open, guidance, _, _) = visited_children

        return SnippetOpen(
            indent=guidance_open.indent,
            guidance=guidance.text,
        )

    def visit_snippet_open_multi(self, node: Node, visited_children: list):
        (guidance_open, _, guidance, _) = visited_children

        return SnippetOpen(
            indent=guidance_open.indent,
            guidance=guidance.text,
        )

    def visit_guidance_open(self, node: Node, visited_children: list):
        (ws, _) = visited_children

        return GuidanceOpen(indent=ws.text)

    def visit_guidance_line(self, node: Node, visited_children: list):
        (_, _, guidance_inner, _) = visited_children

        return guidance_inner

    def visit_guidance_lines(self, node: Node, visited_children: list):
        return Guidance(
            text="\n".join(g.text for g in visited_children)
        )

    def visit_guidance_inner(self, node: Node, visited_children: list):
        return Guidance(text=node.text)

    def visit_human_source(self, node: Node, visited_children: list):
        return PresentPiece(text=node.text)

    def visit_ai_source(self, node: Node, visited_children: list):
        return node.text

def raw_source_to_pieces(input: str) -> list[Piece]:
    tree = grammar.parse(input)
    visitor = RawSourceVisitor()
    pieces = visitor.visit(tree)

    return pieces

def raw_pieces_to_tagged_input(raw_pieces: list[Piece]) -> str:
    tag_source = ""
    id = 0

    for piece in raw_pieces:
        match piece:
            case PresentPiece(text):
                tag_source += text
            case MissingPiece(indent, guidance):
                tag_source += f'# <write_this id="{id}">\n{indent}# {guidance}\n{indent}# </>'
                id += 1
            case _:
                raise TypeError("unknown piece type", piece)

    logger.debug("tagged input ↓\n%s", tag_source)

    return tag_source

def tagged_input_to_tagged_output(tagged_input: str, chat_api: CachedChatAPI) -> str:
    system_prompt = """
You are an expert programmer working on contract. Your client has written a partial program, but left pieces for you to complete. They have marked those with `<write_this>` tags inside Python comments, e.g.:

```
def add_two_numbers(x, y):
    # <write_this id="0">
    # add the two numbers
    # </>

# <write_this id="1">
# add two numbers from command line args, using argparse
# </>
```

You should produce a document that provides a `<completed>` tag for each missing piece, e.g.:

```
<completed id="0">
return x + y
</>
<completed id="1">
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("x", type=int)
parser.add_argument("y", type=int)
args = parser.parse_args()
return add_two_numbers(args.x, args.y)
</>
```

This formatting is very important. The client uses a custom tool to process your work product, and their tool requires this format. Follow this format exactly and do not copy anything outside a `<write_this>` tag.
"""
    chat_messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": tagged_input},
    ]
    tagged_output = chat_api.complete_chat("tagged_input_to_tagged_output", chat_messages)

    logger.debug("tagged output ↓\n%s", tagged_output)

    return tagged_output

def tagged_output_to_completed_pieces(tagged_output: str) -> dict[int, str]:
    pattern = re.compile(r'<completed id="(?P<id>\d+)">\n(?P<content>.+?)</>', re.DOTALL)
    matches = pattern.finditer(tagged_output)
    completed = {int(m.group("id")): m.group("content") for m in matches}

    return completed

def pieces_to_final_source(
        raw_pieces: list[Piece],
        completed_pieces: dict[int, str],
    ) -> str:
    id = 0
    final_source = ""

    for raw in raw_pieces:
        match raw:
            case PresentPiece(text):
                final_source += text

            case MissingPiece(indent, guidance):
                final_source += f"{indent}#<<{guidance}>>\n"
                completed = completed_pieces[id]
                final_source += indent + indent.join(completed.splitlines(True))
                final_source += f"{indent}#<</>>\n"
                id += 1

            case _:
                raise TypeError("unknown piece type", raw)

    logger.debug("final source ↓\n%s", final_source)

    return final_source

def preprocess_maccarone(
        raw_source: str,
        chat_api: CachedChatAPI,
    ) -> str:
    raw_pieces = raw_source_to_pieces(raw_source)
    tagged_input = raw_pieces_to_tagged_input(raw_pieces)
    tagged_output = tagged_input_to_tagged_output(tagged_input, chat_api)
    completed_pieces = tagged_output_to_completed_pieces(tagged_output)
    final_source = pieces_to_final_source(
        raw_pieces,
        completed_pieces,
    )

    return final_source
