import re
from dataclasses import dataclass
from typing import List


@dataclass
class ParsedCommand:
    keystrokes: str
    duration: float


@dataclass
class ParseResult:
    commands: List[ParsedCommand]
    is_task_complete: bool
    error: str
    warning: str
    analysis: str = ""
    plan: str = ""


class TerminusXMLPlainParser:
    """Parser for terminus XML plain response format."""

    def __init__(self):
        self.required_sections = ["analysis", "plan", "commands"]

    def parse_response(self, response: str) -> ParseResult:
        """
        Parse a terminus XML plain response and extract commands.

        Args:
            response: The full LLM response string

        Returns:
            ParseResult with commands, completion status, errors and warnings
        """
        # Try normal parsing first
        result = self._try_parse_response(response)

        if result.error:
            # Try auto-fixes in order until one works
            for fix_name, fix_function in self._get_auto_fixes():
                corrected_response, was_fixed = fix_function(response, result.error)
                if was_fixed:
                    corrected_result = self._try_parse_response(corrected_response)

                    if corrected_result.error == "":
                        # Success! Add auto-correction warning
                        auto_warning = (
                            f"AUTO-CORRECTED: {fix_name} - "
                            "please fix this in future responses"
                        )
                        corrected_result.warning = self._combine_warnings(
                            auto_warning, corrected_result.warning
                        )
                        return corrected_result

        # Return original result if no fix worked
        return result

    def _try_parse_response(self, response: str) -> ParseResult:
        """
        Try to parse a terminus XML plain response.

        Args:
            response: The full LLM response string

        Returns:
            ParseResult with commands, completion status, errors and warnings
        """
        warnings = []

        # Check for extra text before/after <response> tags
        self._check_extra_text(response, warnings)

        # Extract <response> content
        response_content = self._extract_response_content(response)
        if not response_content:
            return ParseResult(
                [],
                False,
                "No <response> tag found",
                "- " + "\n- ".join(warnings) if warnings else "",
                "",
                "",
            )

        # Check if task is complete first
        is_complete = self._check_task_complete(response_content)

        # Check for required sections and extract content
        sections = self._extract_sections(response_content, warnings)

        # Extract analysis and plan for reasoning content
        analysis = sections.get("analysis", "")
        plan = sections.get("plan", "")

        # Extract commands section
        commands_content = sections.get("commands", "")
        if not commands_content:
            if "commands" in sections:
                # Commands section exists but is empty
                if not is_complete:
                    warnings.append(
                        "Commands section is empty; not taking any action. "
                        "If you want to wait a specific amount of time please use "
                        "`sleep`, but if you're waiting for a command to finish then "
                        "continue to wait."
                    )
                return ParseResult(
                    [],
                    is_complete,
                    "",
                    "- " + "\n- ".join(warnings) if warnings else "",
                    analysis,
                    plan,
                )
            else:
                # Commands section is missing entirely
                if is_complete:
                    return ParseResult(
                        [],
                        True,
                        "",
                        "- " + "\n- ".join(warnings) if warnings else "",
                        analysis,
                        plan,
                    )
                return ParseResult(
                    [],
                    False,
                    "Missing <commands> section",
                    "- " + "\n- ".join(warnings) if warnings else "",
                    analysis,
                    plan,
                )

        # Parse commands directly from XML (no code blocks in plain format)
        commands, parse_error = self._parse_xml_commands(commands_content, warnings)
        if parse_error:
            # If task is complete, parse errors are just warnings
            if is_complete:
                warnings.append(parse_error)
                return ParseResult(
                    [],
                    True,
                    "",
                    "- " + "\n- ".join(warnings) if warnings else "",
                    analysis,
                    plan,
                )
            return ParseResult(
                [],
                False,
                parse_error,
                "- " + "\n- ".join(warnings) if warnings else "",
                analysis,
                plan,
            )

        return ParseResult(
            commands,
            is_complete,
            "",
            "- " + "\n- ".join(warnings) if warnings else "",
            analysis,
            plan,
        )

    def _get_auto_fixes(self):
        """Return list of auto-fix functions to try in order."""
        return [
            (
                "Missing </response> tag was automatically inserted",
                self._fix_missing_response_tag,
            ),
        ]

    def _combine_warnings(self, auto_warning: str, existing_warning: str) -> str:
        """Combine auto-correction warning with existing warnings."""
        if existing_warning:
            return f"- {auto_warning}\n{existing_warning}"
        else:
            return f"- {auto_warning}"

    def _fix_missing_response_tag(self, response: str, error: str) -> tuple[str, bool]:
        """Fix missing </response> closing tag by appending it."""
        if "Missing </response> closing tag" not in error:
            return response, False

        # Simply append </response> at the end
        corrected = response.rstrip() + "\n</response>"
        return corrected, True

    def _check_extra_text(self, response: str, warnings: List[str]) -> None:
        """Check for extra text before/after <response> tags."""
        # Find response tag positions
        start_pos = response.find("<response>")
        end_pos = response.find("</response>")

        if start_pos == -1:
            return  # Will be handled as error later

        # Check text before <response>
        before_text = response[:start_pos].strip()
        if before_text:
            warnings.append("Extra text detected before <response> tag")

        # Check text after </response> if closing tag exists
        if end_pos != -1:
            after_text = response[end_pos + len("</response>") :].strip()
            if after_text:
                warnings.append("Extra text detected after </response> tag")

                # Count total <response> tags
                total_response_count = response.count("<response>")
                if total_response_count > 1:
                    warnings.append(
                        f"IMPORTANT: Only issue one <response> block at a time. "
                        f"You issued {total_response_count} and only the first "
                        f"was executed."
                    )

    def _extract_response_content(self, response: str) -> str:
        """Extract content from <response> tags."""
        start_pos = response.find("<response>")
        if start_pos == -1:
            return ""

        end_pos = response.find("</response>", start_pos)
        if end_pos == -1:
            # Missing closing tag - return content from opening tag to end
            return response[start_pos + len("<response>") :].strip()

        return response[start_pos + len("<response>") : end_pos].strip()

    def _extract_sections(self, content: str, warnings: List[str]) -> dict:
        """Extract analysis, plan, commands, and task_complete sections."""
        sections = {}
        found_sections = set()

        # Define patterns for each section
        section_patterns = {
            "analysis": (
                r"<analysis>(.*?)</analysis>",
                r"<analysis\s*/>",
                r"<analysis></analysis>",
            ),
            "plan": (r"<plan>(.*?)</plan>", r"<plan\s*/>", r"<plan></plan>"),
            "commands": (
                r"<commands>(.*?)</commands>",
                r"<commands\s*/>",
                r"<commands></commands>",
            ),
            "task_complete": (
                r"<task_complete>(.*?)</task_complete>",
                r"<task_complete\s*/>",
                r"<task_complete></task_complete>",
            ),
        }

        for section_name, patterns in section_patterns.items():
            full_pattern, self_closing_pattern, empty_pattern = patterns

            # Try full pattern first
            match = re.search(full_pattern, content, re.DOTALL)
            if match:
                sections[section_name] = match.group(1).strip()
                found_sections.add(section_name)
                continue

            # Try self-closing pattern
            if re.search(self_closing_pattern, content):
                sections[section_name] = ""  # Self-closing = empty content
                found_sections.add(section_name)
                continue

            # Try empty pattern
            if re.search(empty_pattern, content):
                sections[section_name] = ""  # Empty = empty content
                found_sections.add(section_name)
                continue

        # Check for missing required sections
        required = set(self.required_sections)
        missing = required - found_sections
        for section in missing:
            if section != "task_complete":  # task_complete is optional
                warnings.append(f"Missing <{section}> section")

        # Check for unexpected tags at the direct child level of <response> only
        # Find all top-level tags (not nested inside other tags)
        top_level_tags = self._find_top_level_tags(content)
        expected_tags = set(self.required_sections + ["task_complete"])
        unexpected = set(top_level_tags) - expected_tags
        for tag in unexpected:
            warnings.append(
                f"Unknown tag found: <{tag}>, expected "
                f"analysis/plan/commands/task_complete"
            )

        # Check for multiple instances of same tag
        for section_name in self.required_sections + ["task_complete"]:
            tag_count = len(re.findall(f"<{section_name}(?:\\s|>|/>)", content))
            if tag_count > 1:
                if section_name == "commands":
                    warnings.append(
                        f"IMPORTANT: Only issue one <commands> block at a time. "
                        f"You issued {tag_count} and only the first was executed."
                    )
                else:
                    warnings.append(f"Multiple <{section_name}> sections found")

        # Check for correct order of sections (analysis, plan, commands)
        self._check_section_order(content, warnings)

        return sections

    def _parse_xml_commands(
        self, xml_content: str, warnings: List[str]
    ) -> tuple[List[ParsedCommand], str]:
        """Parse XML content and extract command objects manually."""

        # Find all keystrokes elements manually for better error reporting
        commands = []
        keystrokes_pattern = re.compile(
            r"<keystrokes([^>]*)>(.*?)</keystrokes>", re.DOTALL
        )

        matches = keystrokes_pattern.findall(xml_content)
        for i, (attributes_str, keystrokes_content) in enumerate(matches):
            # Check for attribute issues
            self._check_attribute_issues(attributes_str, i + 1, warnings)

            # Parse attributes
            duration = 1.0

            # Parse duration attribute
            duration_match = re.search(
                r'duration\s*=\s*["\']([^"\']*)["\']', attributes_str
            )
            if duration_match:
                try:
                    duration = float(duration_match.group(1))
                except ValueError:
                    warnings.append(
                        f"Command {i + 1}: Invalid duration value "
                        f"'{duration_match.group(1)}', using default 1.0"
                    )
            else:
                warnings.append(
                    f"Command {i + 1}: Missing duration attribute, using default 1.0"
                )

            # Check for newline at end of keystrokes
            if i < len(matches) - 1 and not keystrokes_content.endswith("\n"):
                warnings.append(
                    f"Command {i + 1} should end with newline when followed "
                    f"by another command. Otherwise the two commands will be "
                    f"concatenated together on the same line."
                )

            commands.append(
                ParsedCommand(keystrokes=keystrokes_content, duration=duration)
            )

        # Check for XML entities and warn
        entities = {
            "&lt;": "<",
            "&gt;": ">",
            "&amp;": "&",
            "&quot;": '"',
            "&apos;": "'",
        }
        for entity, char in entities.items():
            if entity in xml_content:
                warnings.append(
                    f"Warning: {entity} is read verbatim and not converted to {char}. "
                    f"NEVER USE {entity}, unless you want these exact characters to "
                    f"appear directly in the output."
                )

        # Check for \r\n line endings and warn
        if "\\r\\n" in xml_content:
            warnings.append(
                "Warning: \\r\\n line endings are not necessary - use \\n "
                "instead for simpler output"
            )

        return commands, ""

    def _find_top_level_tags(self, content: str) -> List[str]:
        """Find all top-level XML tags (direct children of response), not
        nested tags."""
        top_level_tags = []
        depth = 0
        i = 0

        while i < len(content):
            if content[i] == "<":
                # Find the end of this tag
                tag_end = content.find(">", i)
                if tag_end == -1:
                    break

                tag_content = content[i + 1 : tag_end]

                # Skip comments, CDATA, etc.
                if tag_content.startswith("!") or tag_content.startswith("?"):
                    i = tag_end + 1
                    continue

                # Check if this is a closing tag
                if tag_content.startswith("/"):
                    depth -= 1
                    i = tag_end + 1
                    continue

                # Check if this is a self-closing tag
                is_self_closing = tag_content.endswith("/")

                # Extract tag name (first word)
                tag_name = tag_content.split()[0] if " " in tag_content else tag_content
                if tag_name.endswith("/"):
                    tag_name = tag_name[:-1]

                # If we're at depth 0, this is a top-level tag
                if depth == 0:
                    top_level_tags.append(tag_name)

                # Adjust depth for opening tags (but not self-closing)
                if not is_self_closing:
                    depth += 1

                i = tag_end + 1
            else:
                i += 1

        return top_level_tags

    def _check_section_order(self, content: str, warnings: List[str]) -> None:
        """Check if sections appear in the correct order: analysis, plan, commands."""
        # Find positions of each section
        positions = {}
        for section in ["analysis", "plan", "commands"]:
            # Look for opening tags
            match = re.search(f"<{section}(?:\\s|>|/>)", content)
            if match:
                positions[section] = match.start()

        # Check if we have at least 2 sections to compare order
        if len(positions) < 2:
            return

        # Expected order
        expected_order = ["analysis", "plan", "commands"]

        # Get sections that are present, in the order they appear
        present_sections = []
        for section in expected_order:
            if section in positions:
                present_sections.append((section, positions[section]))

        # Sort by position to get actual order
        actual_order = [
            section for section, pos in sorted(present_sections, key=lambda x: x[1])
        ]

        # Get expected order for present sections only
        expected_present = [s for s in expected_order if s in positions]

        # Compare orders
        if actual_order != expected_present:
            actual_str = " → ".join(actual_order)
            expected_str = " → ".join(expected_present)
            warnings.append(
                f"Sections appear in wrong order. Found: {actual_str}, "
                f"expected: {expected_str}"
            )

    def _check_attribute_issues(
        self, attributes_str: str, command_num: int, warnings: List[str]
    ) -> None:
        """Check for attribute-related issues."""
        # Check for missing quotes
        unquoted_pattern = re.compile(r'(\w+)\s*=\s*([^"\'\s>]+)')
        unquoted_matches = unquoted_pattern.findall(attributes_str)
        for attr_name, attr_value in unquoted_matches:
            warnings.append(
                f"Command {command_num}: Attribute '{attr_name}' value should be "
                f'quoted: {attr_name}="{attr_value}"'
            )

        # Check for single quotes (should use double quotes for consistency)
        single_quote_pattern = re.compile(r"(\w+)\s*=\s*'([^']*)'")
        single_quote_matches = single_quote_pattern.findall(attributes_str)
        for attr_name, attr_value in single_quote_matches:
            warnings.append(
                f"Command {command_num}: Use double quotes for attribute "
                f"'{attr_name}': {attr_name}=\"{attr_value}\""
            )

        # Check for unknown attribute names
        known_attributes = {"duration"}
        all_attributes = re.findall(r"(\w+)\s*=", attributes_str)
        for attr_name in all_attributes:
            if attr_name not in known_attributes:
                warnings.append(
                    f"Command {command_num}: Unknown attribute '{attr_name}' - "
                    f"known attributes are: {', '.join(sorted(known_attributes))}"
                )

    def _check_task_complete(self, response_content: str) -> bool:
        """Check if the response indicates the task is complete."""
        # Check for <task_complete>true</task_complete>
        true_match = re.search(
            r"<task_complete>\s*true\s*</task_complete>",
            response_content,
            re.IGNORECASE,
        )
        if true_match:
            return True

        # All other cases (false, empty, self-closing, missing) = not complete
        return False

    def salvage_truncated_response(
        self, truncated_response: str
    ) -> tuple[str | None, bool]:
        """
        Try to salvage a valid response from truncated output.

        Args:
            truncated_response: The truncated response from the LLM

        Returns:
            Tuple of (salvaged_response, has_multiple_blocks)
            - salvaged_response: Clean response up to </response> if salvageable, "
                "None otherwise
            - has_multiple_blocks: True if multiple response/commands blocks "
                "were detected
        """
        # Check if we can find a complete response structure
        commands_end = truncated_response.find("</commands>")
        if commands_end == -1:
            return None, False

        # Find the </response> tag after </commands>
        response_end = truncated_response.find("</response>", commands_end)
        if response_end == -1:
            return None, False

        # We have a complete response up to at least </commands>
        # Truncate cleanly at </response>
        clean_response = truncated_response[: response_end + len("</response>")]

        # Check if this is a valid response with no critical issues
        try:
            parse_result = self.parse_response(clean_response)

            # Check if there are no errors and no "multiple blocks" warnings
            has_multiple_blocks = False
            if parse_result.warning:
                warning_lower = parse_result.warning.lower()
                has_multiple_blocks = (
                    "only issue one" in warning_lower
                    and "block at a time" in warning_lower
                )

            if not parse_result.error and not has_multiple_blocks:
                # Valid response! Return the clean truncated version
                return clean_response, False
            else:
                # Has errors or multiple blocks
                return None, has_multiple_blocks

        except Exception:
            # If parsing fails, not salvageable
            return None, False
