"""Dynamic planning system for scaffold branch command.

This module implements the Dynamic Planning System that generates human-readable
implementation plans before code execution, enables plan review and refinement,
and provides bidirectional communication with Claude for interactive sessions.
"""

import asyncio
from typing import TYPE_CHECKING, Any

import click
from dagster_dg_core.context import DgContext
from dagster_shared.record import record
from dagster_shared.utils.timing import format_duration

from dagster_dg_cli.cli.scaffold.branch.claude.diagnostics import ClaudeDiagnostics
from dagster_dg_cli.cli.scaffold.branch.constants import ALLOWED_COMMANDS_PLANNING, ModelType
from dagster_dg_cli.cli.scaffold.branch.version_utils import ensure_claude_sdk_python_version

if TYPE_CHECKING:
    from dagster_dg_cli.cli.scaffold.branch.claude.sdk_client import ClaudeSDKClient, OutputChannel


@record
class GeneratedPlan:
    """Container for a prompt-generated implementation plan.

    Attributes:
        markdown_content: The complete plan as markdown text generated by Claude
        metadata: Additional metadata about plan generation
    """

    markdown_content: str
    metadata: dict[str, Any]


@record
class PlanningContext:
    model: ModelType
    prompt_text: str
    dg_context: DgContext
    project_structure: dict[str, Any]
    verbose: bool


class PlanGenerator:
    """Main orchestrator for plan generation and refinement."""

    def __init__(self, claude_client: "ClaudeSDKClient", diagnostics: ClaudeDiagnostics):
        """Initialize the plan generator.

        Args:
            claude_client: Client for AI interactions
            diagnostics: Diagnostics service for logging
        """
        self.claude_client = claude_client
        self.diagnostics = diagnostics

    def generate_initial_plan(
        self,
        context: PlanningContext,
        output_channel: "OutputChannel",
    ) -> GeneratedPlan:
        """Generate initial implementation plan from user input.

        Args:
            context: Planning context with user input and project information

        Returns:
            Generated implementation plan as markdown content

        Raises:
            Exception: If plan generation fails
        """
        self.diagnostics.info(
            category="planning_generation_start",
            message="Starting initial plan generation",
            data={
                "user_input_length": len(context.prompt_text),
            },
        )

        prompt = self._create_initial_plan_prompt(context)

        allowed_tools = ALLOWED_COMMANDS_PLANNING.copy()

        # Use Claude to generate structured plan as markdown
        messages = asyncio.run(
            self.claude_client.scaffold_with_streaming(
                prompt=prompt,
                model=context.model,
                allowed_tools=allowed_tools,
                output_channel=output_channel,
                disallowed_tools=["Bash(python:*)", "WebSearch", "WebFetch"],
                verbose=context.verbose,
            )
        )

        # Extract the plan content from Claude's response
        plan_content = self._extract_plan_from_messages(messages, output_channel)

        plan = GeneratedPlan(
            markdown_content=plan_content,
            metadata={
                "generation_method": "prompt_driven_claude_generation",
                "messages_count": len(messages),
                "user_input": context.prompt_text,
            },
        )

        self.diagnostics.info(
            category="planning_generation_completed",
            message="Initial plan generation completed",
            data={
                "plan_content_length": len(plan_content),
                "messages_count": len(messages),
            },
        )

        return plan

    def refine_plan(
        self,
        context: PlanningContext,
        current_plan: GeneratedPlan,
        user_feedback: str,
        output_channel: "OutputChannel",
    ) -> GeneratedPlan:
        """Refine existing plan based on user feedback.

        Args:
            current_plan: The current plan to refine
            user_feedback: User feedback for plan improvements

        Returns:
            Refined implementation plan
        """
        self.diagnostics.info(
            category="planning_refinement_start",
            message="Starting plan refinement",
            data={
                "current_plan_length": len(current_plan.markdown_content),
                "feedback_length": len(user_feedback),
            },
        )

        prompt = self._create_refinement_prompt(current_plan, user_feedback)

        allowed_tools = ALLOWED_COMMANDS_PLANNING.copy()

        messages = asyncio.run(
            self.claude_client.scaffold_with_streaming(
                prompt=prompt,
                model=context.model,
                allowed_tools=allowed_tools,
                output_channel=output_channel,
                disallowed_tools=["Bash(python:*)", "WebSearch", "WebFetch"],
                verbose=context.verbose,
            )
        )

        # Extract refined plan content from Claude's response
        refined_content = self._extract_plan_from_messages(messages, output_channel)

        refined_plan = GeneratedPlan(
            markdown_content=refined_content,
            metadata={
                **current_plan.metadata,
                "refinement_method": "prompt_driven_claude_refinement",
                "refinement_messages": len(messages),
                "user_feedback": user_feedback,
            },
        )

        self.diagnostics.info(
            category="planning_refinement_completed",
            message="Plan refinement completed",
            data={
                "refined_content_length": len(refined_content),
            },
        )

        return refined_plan

    def _create_initial_plan_prompt(self, context: PlanningContext) -> str:
        """Create a natural prompt for initial plan generation by loading from template file.

        Args:
            context: Planning context with user input and project information

        Returns:
            Natural language prompt for Claude to generate a plan
        """
        from pathlib import Path

        # Load prompt template from external file
        prompt_path = Path(__file__).parent / "prompts" / "planning_prompt.md"
        template = prompt_path.read_text()

        context_info = ""
        if context.project_structure:
            context_info = f"\n\nProject Structure Overview: {context.project_structure}"

        # Format template with actual values
        return template.format(
            user_input=context.prompt_text,
            context_info=context_info,
        )

    def _create_refinement_prompt(self, current_plan: GeneratedPlan, user_feedback: str) -> str:
        """Create a natural prompt for plan refinement.

        Args:
            current_plan: The current plan to refine
            user_feedback: User's feedback for improvements

        Returns:
            Natural language prompt for Claude to refine the plan
        """
        return f"""You are refining an implementation plan based on user feedback.

Here is the current plan:

{current_plan.markdown_content}

---

User Feedback: {user_feedback}

---

Please generate an improved version of the implementation plan that addresses the user's feedback. Keep the same markdown structure and format, but incorporate the requested changes. Be specific about what you're modifying and why.

Provide the complete updated plan in the same markdown format as before."""

    def _extract_plan_from_messages(self, messages: list, output_channel: "OutputChannel") -> str:
        """Extract plan content from Claude's response messages.

        Args:
            messages: List of message objects from Claude client

        Returns:
            The plan content as markdown text
        """
        ensure_claude_sdk_python_version()

        from claude_code_sdk.types import ResultMessage

        self.diagnostics.debug(
            category="plan_extraction_start",
            message="Extracting plan from Claude messages",
            data={"message_count": len(messages)},
        )

        # Look specifically for ResultMessage and extract the text content
        plan_content = None
        success_result_found = False
        result_message = None

        for i, message in enumerate(messages):
            self.diagnostics.debug(
                category="plan_message_processing",
                message=f"Processing message {i}",
                data={
                    "message_class": type(message).__name__,
                },
            )

            if isinstance(message, ResultMessage):
                success_result_found = True
                result_message = message
                if message.result and message.result.strip():
                    plan_content = message.result.strip()
                    self.diagnostics.info(
                        category="plan_success_result_found",
                        message=f"Found ResultMessage with {len(plan_content)} chars of content",
                        data={
                            "message_index": i,
                        },
                    )
                    break  # Found what we need, stop processing
                else:
                    self.diagnostics.error(
                        category="plan_success_result_empty",
                        message="Found ResultMessage but result is empty",
                        data={"message_index": i},
                    )

        # If we didn't find AssistantMessage, error and exit
        if not success_result_found:
            message_types = [type(msg).__name__ for msg in messages]
            self.diagnostics.error(
                category="plan_no_success_result",
                message="No ResultMessage found in Claude response",
                data={
                    "message_count": len(messages),
                    "message_types": message_types,
                },
            )
            raise Exception(
                f"Expected ResultMessage from Claude SDK but got {len(messages)} message(s) "
                f"of types: {message_types}. This indicates a problem with Claude SDK communication "
                f"or model response. Check your Claude SDK installation and try again."
            )

        # If we found ResultMessage but no content, error and exit
        if not plan_content:
            raise Exception(
                "Found ResultMessage but the text content was empty. "
                "Claude generated a response but provided no plan content."
            )

        combined_content = plan_content

        # Display success summary to user
        if result_message:
            output_channel.write(
                f"✅ Plan generation completed ({len(plan_content):,} characters, ${result_message.total_cost_usd:.2f}, {format_duration(result_message.duration_ms)})."
            )

        self.diagnostics.debug(
            category="plan_extraction_result",
            message="Plan extraction completed",
            data={
                "content_length": len(combined_content),
                "has_implementation_plan": "# Implementation Plan" in combined_content,
            },
        )

        # Final validation - this should not happen since we validated above, but just in case
        if not combined_content or len(combined_content) < 100:
            raise Exception(
                f"AssistantMessage contained only {len(combined_content)} characters of content, "
                f"which is too short for a meaningful plan. Expected at least 100 characters."
            )

        return combined_content


def get_user_plan_approval(plan: GeneratedPlan) -> tuple[bool, str | None]:
    """Get user approval for a generated plan.

    Args:
        plan: The plan to review

    Returns:
        Tuple of (approved, feedback) where approved indicates if the user
        approved the plan, and feedback contains refinement suggestions
    """
    click.echo("\n" + "=" * 60)
    click.echo("IMPLEMENTATION PLAN REVIEW")
    click.echo("=" * 60)
    click.echo("")

    # Display the plan content directly
    click.echo(plan.markdown_content)

    click.echo("")
    click.echo("=" * 60)

    while True:
        choice = click.prompt(
            "Plan Review Options:\n"
            "  [a]pprove - Execute this plan as-is\n"
            "  [r]efine - Provide feedback to improve the plan\n"
            "  [c]ancel - Cancel the operation\n\n"
            "Your choice",
            type=click.Choice(["a", "r", "c", "approve", "refine", "cancel"], case_sensitive=False),
            default="approve",
        ).lower()

        if choice in ("a", "approve"):
            return True, None
        elif choice in ("r", "refine"):
            feedback = click.prompt(
                "\nWhat would you like to change about this plan?\n"
                "Be specific about steps, files, or approaches you'd like modified",
                type=str,
            )
            return False, feedback
        elif choice in ("c", "cancel"):
            raise click.ClickException("Operation cancelled by user")
