"""Jinja-based prompt building system for reducing repetition across agent files."""

import re
import xml.etree.ElementTree as ET
from jinja2 import Environment, FileSystemLoader, TemplateNotFound
from typing import Dict, Union, Any, Optional
from pathlib import Path
from blocks_control_sdk.control.agent_base import LLM
from blocks_control_sdk.utils import get_file_path
from blocks_control_sdk.tools.github import get_code_repositories_block, get_code_repositories_with_additional_info_block
import uuid
from blocks_control_sdk.blocks_events.context import format_event_context, BLOCKS_EVENT__TRIGGER_ALIAS


class XMLBlockParser:
    """Parser for XML blocks with dot notation support for nested element access."""

    def __init__(self, xml_content: str):
        """Initialize parser with XML content."""
        self.xml_content = xml_content.strip()
        self.root = self._parse_xml()

    def _parse_xml(self) -> ET.Element:
        """Parse XML content and return root element."""
        try:
            # Clean up the content to make it valid XML
            cleaned_content = self._clean_xml_content(self.xml_content)
            return ET.fromstring(cleaned_content)
        except ET.ParseError as e:
            raise ValueError(f"Invalid XML content: {e}")

    def _clean_xml_content(self, content: str) -> str:
        """Clean XML content to handle problematic characters like @ in <@user>."""
        import html

        # First, escape all < and > that are NOT part of valid XML tags
        # This handles content like <@U123> which should be treated as text content
        def escape_non_xml_brackets(text):
            # Split content by valid XML tags to process text content separately
            parts = re.split(r'(</?[\w]+>)', text)

            escaped_parts = []
            for part in parts:
                if re.match(r'</?[\w]+>', part):
                    # This is a valid XML tag, keep as is
                    escaped_parts.append(part)
                else:
                    # This is text content, escape any < or > characters
                    escaped_parts.append(html.escape(part))

            return ''.join(escaped_parts)

        return escape_non_xml_brackets(content)

    def get(self, path: str) -> Optional[str]:
        """
        Get value using dot notation path.

        Args:
            path: Dot-separated path to element (e.g., 'chat_history.latest_user_message')

        Returns:
            Text content of the element, or None if not found
        """
        element = self._find_element(path)
        return element.text if element is not None else None

    def set(self, path: str, value: str) -> None:
        """
        Set value using dot notation path.

        Args:
            path: Dot-separated path to element
            value: New text content for the element
        """
        element = self._find_element(path)
        if element is not None:
            element.text = value
        else:
            self._create_element_path(path, value)

    def _find_element(self, path: str) -> Optional[ET.Element]:
        """Find element using dot notation path."""
        parts = path.split('.')
        current = self.root

        for part in parts:
            child = current.find(part)
            if child is None:
                return None
            current = child

        return current

    def _create_element_path(self, path: str, value: str) -> None:
        """Create element path if it doesn't exist."""
        parts = path.split('.')
        current = self.root

        for part in parts:
            child = current.find(part)
            if child is None:
                child = ET.SubElement(current, part)
            current = child

        current.text = value

    def to_xml(self) -> str:
        """Convert back to XML string."""
        return ET.tostring(self.root, encoding='unicode')


class PromptBuilder:
    """Build prompts from modular Jinja templates with support for overrides."""

    rendered_standard_prompt: str = None
    formatted_context: str = None
    task_input: dict = {}
    latest_user_message: str = None
    agent: str = None
    trigger_alias: str = None
    
    def __init__(self, base_path: str = './prompts'):
        """Initialize the prompt builder with template environment."""
        # Get path relative to this file's location
        file_dir = Path(__file__).parent
        abs_path = (file_dir / base_path).absolute()
        print(f"Initializing prompt builder with base path: {base_path}")
        print(f"Absolute path: {abs_path}")
        self.env = Environment(
            loader=FileSystemLoader(abs_path),
            trim_blocks=True,      # Removes first newline after a block
            lstrip_blocks=False,     # Strips leading spaces/tabs from line starts
            keep_trailing_newline=False  # Don't keep trailing newline at end of template
        )
        self.base_path = Path(base_path)
    
    def get_base_template(self, trigger: Union[str, BLOCKS_EVENT__TRIGGER_ALIAS], agent: str = None):
        """
        Get the main template with override support.
        
        Args:
            trigger: Type of trigger (BLOCKS_EVENT__TRIGGER_ALIAS enum or string)
            agent: Optional agent name for overrides (codex, claude, gemini)
        
        Returns:
            Jinja template object
        """
        # Convert enum to template name
        template_name = self._get_template_name(trigger)
        # Try agent-specific override first
        if agent:
            try:
                return self.env.get_template(f'overrides/{agent}/{template_name}/base.j2')
            except TemplateNotFound:
                pass
        
        # Fall back to default base template
        return self.env.get_template(f'base/{template_name}.j2')
    
    def _get_template_name(self, trigger: Union[str, BLOCKS_EVENT__TRIGGER_ALIAS]) -> str:
        """Convert trigger enum or string to template name."""
        if isinstance(trigger, BLOCKS_EVENT__TRIGGER_ALIAS):
            # Map enum values to template names
            mapping = {
                BLOCKS_EVENT__TRIGGER_ALIAS.GITHUB_ISSUE_COMMENT: 'github_issue',
                BLOCKS_EVENT__TRIGGER_ALIAS.GITHUB_PULL_REQUEST_COMMENT: 'github_pr',
                BLOCKS_EVENT__TRIGGER_ALIAS.LINEAR_ISSUE_COMMENT: 'linear_issue',
                BLOCKS_EVENT__TRIGGER_ALIAS.SLACK_MENTION: 'slack_mention',
                BLOCKS_EVENT__TRIGGER_ALIAS.WEBHOOK: 'webhook',
                BLOCKS_EVENT__TRIGGER_ALIAS.GITLAB_ISSUE_COMMENT: 'gitlab_issue',
                BLOCKS_EVENT__TRIGGER_ALIAS.GITLAB_MERGE_REQUEST_COMMENT: 'gitlab_merge_request',
            }
            return mapping.get(trigger, 'webhook')
        return trigger
    
    def get_section(self, section: str, context: Dict = None, agent: str = None, trigger: Union[str, BLOCKS_EVENT__TRIGGER_ALIAS] = None) -> str:
        """
        Get a section template with override support and render it.
        
        Args:
            section: Section name (e.g., 'context_files', 'hard_limits')
            context: Optional context variables for rendering
            agent: Optional agent name for overrides
            trigger: Optional trigger type for overrides (enum or string)
        
        Returns:
            Rendered section content or empty string if not found
        """
        template_name = self._get_template_name(trigger) if trigger else None

        if context is None and agent:
            context = PromptBuilder.create_common_context(agent)
        # Try agent/trigger specific override first
        if agent and template_name:
            try:
                template = self.env.get_template(f'overrides/{agent}/{template_name}/{section}.j2')
                return template.render(**(context or {}))
            except TemplateNotFound:
                pass
        
        # Fall back to default section
        try:
            template = self.env.get_template(f'sections/{section}.j2')
            return template.render(**(context or {}))
        except TemplateNotFound:
            # Return empty string if section doesn't exist
            return ''
    
    def _render_template(self, trigger: Union[str, BLOCKS_EVENT__TRIGGER_ALIAS], sections: Dict[str, str],
              context: Dict, agent: str = None) -> str:
        """
        Build prompt from base template with injected sections.

        Args:
            trigger: Type of trigger (BLOCKS_EVENT__TRIGGER_ALIAS enum or string)
            sections: Dict of section_name -> section_content to inject
            context: Variables for the main template
            agent: Optional agent name for overrides

        Returns:
            Fully rendered prompt string
        """
        # Get the main template
        template = self.get_base_template(trigger, agent)

        # Merge sections into context
        full_context = {**context, **sections}

        return template.render(**full_context)

    def render_formatted_context_with_message(self, new_message: str = None) -> str:
        """
        Create a finished formatted context with a specific user message.
        Uses Jinja to render the formatted_context template with common_context and the message.

        Args:
            new_message: Optional new message to use. If not provided, uses self.latest_user_message

        Returns:
            Rendered formatted context string with the message and common context inserted

        Raises:
            ValueError: If neither new_message nor self.latest_user_message is available
            ValueError: If formatted_context, agent, or trigger_alias is not available
        """
        message_to_use = new_message if new_message is not None else self.latest_user_message

        if message_to_use is None:
            raise ValueError("No message available: provide new_message or ensure builder has latest_user_message")

        if self.formatted_context is None:
            raise ValueError("No formatted_context available in builder")

        if self.agent is None:
            raise ValueError("No agent available in builder")

        # Get common context for rendering (includes tool names, instruction_files, repos_xml_block)
        common_context = PromptBuilder.create_common_context(self.agent)

        # Render the formatted_context as a Jinja template
        template = self.env.from_string(self.formatted_context)
        return template.render(common_context).replace('{latest_user_message}', message_to_use)

    def render_full_initial_prompt_with_message(self, new_message: str = None) -> str:
        """
        Create a finished prompt with a specific user message.

        Args:
            new_message: Optional new message to use. If not provided, uses self.latest_user_message

        Returns:
            New prompt string with the message inserted (does not modify internal state)

        Raises:
            ValueError: If neither new_message nor self.latest_user_message is available
            ValueError: If no rendered_prompt_text is available
        """
        message_to_use = new_message if new_message is not None else self.latest_user_message

        if message_to_use is None:
            raise ValueError("No message available: provide new_message or ensure builder has latest_user_message")

        if self.rendered_standard_prompt is None:
            raise ValueError("No rendered_prompt_text available in builder")

        # Replace placeholder and return new string without modifying internal state
        return self.rendered_standard_prompt.replace('{latest_user_message}', message_to_use)

    def initial(self, message: str = None) -> str:
        """Full prompt for first query (instructions + sections + context + message)."""
        return self.render_full_initial_prompt_with_message(message)

    def followup(self, message: str = None) -> str:
        """Lighter prompt for follow-up queries (context + message)."""
        return self.render_formatted_context_with_message(message)

    # Aliases that share a prompt template
    _NORMALIZE_TRIGGER = {"linear.assign": "linear.issue_comment"}

    @staticmethod
    def build(trigger: str, event_data: dict, agent: str, **kwargs) -> 'PromptBuilder':
        """Build a prompt from event data. No side effects.

        Args:
            trigger: "webhook", "github.issue_comment", "slack.mention", etc.
            event_data: Raw event input dict
            agent: "claude", "codex", "gemini"
        """
        trigger = PromptBuilder._NORMALIZE_TRIGGER.get(trigger, trigger)
        trigger_enum = BLOCKS_EVENT__TRIGGER_ALIAS(trigger)
        return build_standard_prompt(
            trigger_alias=trigger_enum, input=event_data, agent=agent, **kwargs
        )

    @staticmethod
    def create_common_context(agent: str) -> Dict[str, Any]:
        """
        Create the common context object with tool names, instruction files, and repos block.

        Args:
            agent: Agent name (codex, claude, gemini) - determines tool name format

        Returns:
            Dictionary containing:
                - Tool name mappings (4 tool keys)
                - instruction_files: Standard instruction files dict
                - repos_xml_block: GitHub repositories XML block (empty string if unavailable)
        """
        # Standard instruction files
        instruction_files = {
            'Runtime Information': get_file_path("instructions/BLOCKS_ENVIRONMENT.md").absolute(),
            'Github Auth Token Troubleshooting': get_file_path("instructions/AUTH_TROUBLESHOOTING.md").absolute(),
        }

        # Get GitHub repositories block
        repos_xml_block = ""
        try:
            _simple_repos_xml_block = get_code_repositories_block()
            _additional_repos_xml_block = get_code_repositories_with_additional_info_block()
            repos_xml_block = _additional_repos_xml_block or _simple_repos_xml_block
        except Exception as e:
            # If we can't get the github repositories block, just skip it
            print(f"Warning: Could not create GitHub repositories block: {e}")

        # Set tool names based on agent type
        TOOL_CLONE_REPOSITORY_INTO_FOLDER = 'MCP Tool: Blocks Internal Clone Repository Into Folder'
        TOOL_CREATE_GITHUB_PR = "MCP Tool: Blocks Internal Create Github PR"
        TOOL_CREATE_PR_REVIEW_COMMENT_ON_FILE = 'MCP Tool: Blocks Internal Create PR Review Comment On File'
        TOOL_CREATE_PR_REVIEW_COMMENT_ON_LINES = 'MCP Tool: Blocks Internal Create PR Review Comment On Lines'
        TOOL_REGISTER_RUNNING_SERVER_PORT = 'MCP Tool: Blocks Internal Register Running Server Port'
        TOOL_DOWNLOAD_IMAGE_FROM_SLACK = 'MCP Tool: Blocks Internal Download Image From Slack'
        TOOL_DOWNLOAD_IMAGE_FROM_LINEAR = 'MCP Tool: Blocks Internal Download Image From Linear'
        TOOL_DOWNLOAD_IMAGES_FROM_GITHUB_COMMENT_ID = 'MCP Tool: Blocks Internal Download Images From GitHub Comment ID'
        TOOL_EXPLORE_SUBAGENT = 'MCP Tool: Blocks Internal Explore Subagent'

        if agent == 'claude':
            TOOL_CLONE_REPOSITORY_INTO_FOLDER = 'mcp__blocks-internal-mcp__clone_repository_into_folder'
            TOOL_CREATE_GITHUB_PR = 'mcp__blocks-internal-mcp__create_pull_request'
            TOOL_CREATE_PR_REVIEW_COMMENT_ON_FILE = 'mcp__blocks-internal-mcp__create_pr_review_comment_on_file'
            TOOL_CREATE_PR_REVIEW_COMMENT_ON_LINES = 'mcp__blocks-internal-mcp__create_pr_review_comment_on_lines'
            TOOL_REGISTER_RUNNING_SERVER_PORT = 'mcp__blocks-internal-mcp__register_running_server_port'
            TOOL_DOWNLOAD_IMAGE_FROM_SLACK = 'mcp__blocks-internal-mcp__download_image_from_slack'
            TOOL_DOWNLOAD_IMAGE_FROM_LINEAR = 'mcp__blocks-internal-mcp__download_image_from_linear'
            TOOL_DOWNLOAD_IMAGES_FROM_GITHUB_COMMENT_ID = 'mcp__blocks-internal-mcp__download_images_from_github_comment_id'
            TOOL_EXPLORE_SUBAGENT = 'mcp__blocks-internal-mcp__ExploreSubAgent'
        elif agent == 'codex':
            TOOL_CLONE_REPOSITORY_INTO_FOLDER = 'functions.blocks-internal-mcp__clone_repository_into_folder'
            TOOL_CREATE_GITHUB_PR = 'functions.blocks-internal-mcp__create_pull_request'
            TOOL_CREATE_PR_REVIEW_COMMENT_ON_FILE = 'functions.blocks-internal-mcp__create_pr_review_comment_on_file'
            TOOL_CREATE_PR_REVIEW_COMMENT_ON_LINES = 'functions.blocks-internal-mcp__create_pr_review_comment_on_lines'
            TOOL_REGISTER_RUNNING_SERVER_PORT = 'functions.blocks-internal-mcp__register_running_server_port'
            TOOL_DOWNLOAD_IMAGE_FROM_SLACK = 'functions.blocks-internal-mcp__download_image_from_slack'
            TOOL_DOWNLOAD_IMAGE_FROM_LINEAR = 'functions.blocks-internal-mcp__download_image_from_linear'
            TOOL_DOWNLOAD_IMAGES_FROM_GITHUB_COMMENT_ID = 'functions.blocks-internal-mcp__download_images_from_github_comment_id'
            TOOL_EXPLORE_SUBAGENT = 'functions.blocks-internal-mcp__ExploreSubAgent'
        return {
            'agent': agent,
            'tool_clone_repository_into_folder': TOOL_CLONE_REPOSITORY_INTO_FOLDER,
            'tool_create_github_pr': TOOL_CREATE_GITHUB_PR,
            'tool_create_pr_review_comment_on_file': TOOL_CREATE_PR_REVIEW_COMMENT_ON_FILE,
            'tool_create_pr_review_comment_on_lines': TOOL_CREATE_PR_REVIEW_COMMENT_ON_LINES,
            'tool_register_running_server_port': TOOL_REGISTER_RUNNING_SERVER_PORT,
            'tool_download_image_from_slack': TOOL_DOWNLOAD_IMAGE_FROM_SLACK,
            'tool_download_image_from_linear': TOOL_DOWNLOAD_IMAGE_FROM_LINEAR,
            'tool_download_images_from_github_comment_id': TOOL_DOWNLOAD_IMAGES_FROM_GITHUB_COMMENT_ID,
            'tool_explore_subagent': TOOL_EXPLORE_SUBAGENT,
            'instruction_files': instruction_files,
            'repos_xml_block': repos_xml_block
        }


def build_standard_prompt(
    trigger_alias: BLOCKS_EVENT__TRIGGER_ALIAS,
    input: dict,
    agent: str,
    message_history_path: Path = None,
    custom_instruction_files: Dict[str, Path] = None,
    include_write_code_instructions: bool = False
) -> PromptBuilder:
    """
    Helper function for standard prompt building.

    Args:
        trigger_alias: Type of trigger (BLOCKS_EVENT__TRIGGER_ALIAS enum or string for backwards compatibility)
        input: Input data from the event
        agent: Agent name (codex, claude, gemini)
        message_history_path: Path to message history file (optional)
        custom_instruction_files: Additional instruction files to include
        include_write_code_instructions: Whether to include write code instructions

    Returns:
        Fully constructed prompt string
    """
    builder = PromptBuilder()

    # Get common context (tool names, standard instruction files, repos block)
    common_context = PromptBuilder.create_common_context(agent)

    # Extract instruction_files and repos_xml_block for easier access
    instruction_files = common_context['instruction_files']
    repos_xml_block = common_context['repos_xml_block']

    # Add custom instruction files if provided
    if custom_instruction_files:
        instruction_files.update(custom_instruction_files)

    # Get message history path
    if message_history_path is None:
        message_history_path = get_file_path("BLOCKS_MESSAGE_HISTORY.xml")

    trigger_alias = trigger_alias.value

    # Format the event context
    formatted_context, latest_user_message = format_event_context(trigger_alias, input, agent)

    # Update common_context with the merged instruction_files
    common_context['instruction_files'] = instruction_files

    # Build sections (common_context now includes all necessary data)
    sections = {
        'using_github': builder.get_section(
            'using_github',
            agent=agent,
            trigger=trigger_alias,
            context=common_context
        ),
        'context_files': builder.get_section(
            'context_files',
            agent=agent,
            trigger=trigger_alias,
            context=common_context
        ) if instruction_files else '',
        'branching_naming': builder.get_section(
            'branching_naming',
            agent=agent,
            trigger=trigger_alias,
            context=common_context
        ),
        'hard_limits': builder.get_section(
            'hard_limits',
            agent=agent,
            trigger=trigger_alias,
            context=common_context
        ),
        'workspace': builder.get_section(
            'workspace',
            agent=agent,
            trigger=trigger_alias,
            context=common_context
        ),
        'environment_info': builder.get_section(
            'environment_info',
            agent=agent,
            trigger=trigger_alias,
            context=common_context
        ),
        'code_repositories': builder.get_section(
            'available_repos',
            agent=agent,
            trigger=trigger_alias,
            context=common_context
        ) if repos_xml_block else '',
    }

    formatted_context_rendered = builder.env.from_string(formatted_context).render(common_context)
    
    # Build and return the full prompt
    rendered_standard_prompt = builder._render_template(
        trigger=trigger_alias,
        sections=sections,
        context={
            'formatted_context': formatted_context_rendered,
            **common_context
        },
        agent=agent
    )
    builder.rendered_standard_prompt = rendered_standard_prompt
    builder.formatted_context = formatted_context
    builder.task_input = input
    builder.latest_user_message = latest_user_message
    builder.agent = agent
    builder.trigger_alias = trigger_alias

    return builder


def extract_xml_blocks(prompt_text: str) -> Dict[str, str]:
    """
    Extract top-level XML blocks from a rendered prompt.

    Args:
        prompt_text: The full rendered prompt containing XML blocks

    Returns:
        Dictionary mapping block names to their XML content
    """
    xml_blocks = {}
    # Find XML blocks that start at line beginning or after newline
    pattern = r'(?:^|\n)<(\w+)>(.*?)</\1>(?=\s*(?:\n|$))'
    matches = re.finditer(pattern, prompt_text, re.DOTALL | re.MULTILINE)

    for match in matches:
        tag_name = match.group(1)
        xml_content = match.group(0).strip()  # Full XML block including tags
        xml_blocks[tag_name] = xml_content

    return xml_blocks


def parse_block(xml_content: str) -> XMLBlockParser:
    """
    Parse an XML block for dot notation access.

    Args:
        xml_content: XML content string to parse

    Returns:
        XMLBlockParser instance for accessing nested elements
    """
    return XMLBlockParser(xml_content)


def get_plan_prompt(llm_provider: LLM) -> str:
    """
    Get the plan prompt for the given LLM provider.
    """
    builder = PromptBuilder()

    # home directory / .config / blocks / uuid-plan-file.md
    plan_file_path = Path.home() / ".config" / "blocks" / "plans" / f"{uuid.uuid4()}-plan-file.md"

    common_context = PromptBuilder.create_common_context(llm_provider.value)

    if llm_provider == LLM.CLAUDE or llm_provider == LLM.OPENCODE or llm_provider == LLM.CURSOR or llm_provider == LLM.KIMI:
        return builder.get_section(
            'plan_claude',
            agent='claude',
            context={
                'PLAN_FILE': plan_file_path.absolute(),
                **common_context
            }
        )
    elif llm_provider == LLM.CODEX:
        return builder.get_section(
            'plan_codex',
            agent='codex',
            context={
                'PLAN_FILE': plan_file_path.absolute(),
                **common_context
            }
        )
    elif llm_provider == LLM.GEMINI:
        return builder.get_section(
            'plan_gemini',
            agent='gemini',
            context={
                'PLAN_FILE': plan_file_path.absolute(),
                **common_context
            }
        )
    else:
        return builder.get_section(
            'plan_claude',
            agent='claude',
            context={
                'PLAN_FILE': plan_file_path.absolute(),
                **common_context
            }
        )

def get_explore_subagent_prompt() -> str:
    """
    Get the explore subagent prompt for the given LLM provider.
    """
    builder = PromptBuilder()
    common_context = PromptBuilder.create_common_context('codex')
    return builder.get_section(
        'explore_subagent',
        agent='codex',
        context=common_context
    )


# Alias for cleaner developer API
Prompt = PromptBuilder