from __future__ import annotations

from enum import Enum
from openai_harmony import (
    load_harmony_encoding,
    HarmonyEncodingName,
    StreamableParser,
    Role
)    

class ChannelType(Enum):
    """Enumeration of harmony channel types."""

    ANALYSIS = "analysis"
    COMMENTARY = "commentary" 
    FINAL = "final"

class ToolParserState(Enum):
    """Enumeration of parser states."""
    NORMAL = "normal"
    FOUND_ARGUMENTS = "found_arguments"
    END_STREAM = "end_stream"

class HarmonyParser:
    """Parser for Harmony encoding."""

    def __init__(self):
        self.encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
        self.parser = StreamableParser(self.encoding, role=Role.ASSISTANT)

        self.end_tool_chunk = "<|call|>"
        self.state = ToolParserState.NORMAL
        self.arguments_buffer = []  # Changed to list for efficient concatenation
        self.function_name_buffer = ""

    def parse(self, text: str) -> dict[str, str] | None:
        """
        Parse the text and return the parsed content.
        """
        if self.end_tool_chunk in text:
            text = text.split(self.end_tool_chunk)[0]

        result = {
            "content": None,
            "tool_calls": [],
            "reasoning_content": None,
        }
        tokens = self.encoding.encode(text, allowed_special="all")
        parsed_messages = self.encoding.parse_messages_from_completion_tokens(tokens, role=Role.ASSISTANT)
        for message in parsed_messages:
            if message.channel == ChannelType.ANALYSIS.value:
                result["reasoning_content"] = message.content[0].text
            elif message.channel == ChannelType.COMMENTARY.value:
                result["tool_calls"].append({
                    "name": message.recipient.replace("functions.", ""),
                    "arguments": message.content[0].text
                })
            elif message.channel == ChannelType.FINAL.value:
                result["content"] = message.content[0].text
        return result
    
    def _build_result(
        self, 
        reasoning_contents: list[str], 
        tool_calls: list[dict[str, str]] | None, 
        contents: list[str]
    ) -> dict[str, str | list | None]:
        """Build the result dictionary from accumulated content."""
        return {
            "reasoning_content": "".join(reasoning_contents) or None,
            "tool_calls": tool_calls,
            "content": "".join(contents) or None,
        }
        
    def parse_streaming(self, chunk: str) -> tuple[dict[str, str | list | None] | None, bool]:
        """Parse the chunk and return the parsed content."""
        if self.state == ToolParserState.END_STREAM:
            return None, True

        reasoning_contents = []
        contents = []
        end_stream_state = False

        # Check for end marker and truncate
        if self.end_tool_chunk in chunk:
            chunk = chunk[:chunk.find(self.end_tool_chunk)]
            end_stream_state = True
        
        # Process chunk tokens
        chunk_tokens = self.encoding.encode(chunk, allowed_special="all")
        for chunk_token in chunk_tokens:
            stream_text = self.parser.process(chunk_token)
            content = stream_text.last_content_delta

            if not content:
                continue

            # Handle FOUND_ARGUMENTS state separately
            if self.state == ToolParserState.FOUND_ARGUMENTS:
                self.arguments_buffer.append(content)
                continue

            # Handle different channels
            current_channel = stream_text.current_channel
            if current_channel == ChannelType.ANALYSIS.value:
                reasoning_contents.append(content)
            elif current_channel == ChannelType.COMMENTARY.value:
                self.state = ToolParserState.FOUND_ARGUMENTS
                self.arguments_buffer.append(content)
                self.function_name_buffer = stream_text.current_recipient.replace("functions.", "")
            elif current_channel == ChannelType.FINAL.value:
                contents.append(content)

        # Handle end of stream
        if end_stream_state:
            tool_calls = [{
                "name": self.function_name_buffer,
                "arguments": "".join(self.arguments_buffer)
            }]
            self.arguments_buffer = []
            self.function_name_buffer = ""
            self.state = ToolParserState.END_STREAM
            return self._build_result(reasoning_contents, tool_calls, contents), True
        
        return self._build_result(reasoning_contents, None, contents), False