"""
PlanGenerationGraph: LangGraph workflow for plan generation.

This module implements a stateful workflow for generating code plans using
LangGraph's StateGraph. The workflow includes context loading, requirement
analysis, plan generation, and validation with retry logic.

Workflow Stages
---------------
1. **load_context**: Loads and prepares context information
2. **analyze_requirements**: Analyzes user prompt for requirements
3. **generate_plan**: Generates code changes based on requirements
4. **validate**: Validates generated changes (with retry logic)

Example Usage
-------------
::

    from cleveragents.agents import PlanGenerationGraph

    # Create the agent
    graph = PlanGenerationGraph(max_retries=3)

    # Execute the workflow
    result = graph.invoke(project, plan, contexts, thread_id="my-thread")

    # Or stream for real-time progress
    for event in graph.stream(project, plan, contexts):
        print(event)
"""

import re
from collections.abc import Iterator
from pathlib import Path
from typing import Any, TypedDict, cast
from uuid import uuid4

from langchain_core.documents import Document
from langchain_core.language_models import BaseLanguageModel
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate  # type: ignore[attr-defined]
from langgraph.checkpoint.memory import MemorySaver  # type: ignore[import-untyped]
from langgraph.graph import END, StateGraph  # type: ignore[import-untyped]

from cleveragents.agents.context_analysis import (
    ContextAnalysisAgent,
    ContextAnalysisState,
)
from cleveragents.application.services.prompt_sanitizer import PromptSanitizer
from cleveragents.domain.models.core import (
    Change,
    Context,
    OperationType,
    Plan,
    Project,
)
from cleveragents.domain.providers.ai_provider import ActorInvocationContext

# Module-level sanitizer for prompt boundary markers (mechanism 2)
_SANITIZER = PromptSanitizer()

PATH_HINT_RE = re.compile(r"@(?P<path>[\w./-]+)")


class BoundedMemorySaver(MemorySaver):
    """Memory saver that retains only the latest checkpoints per thread."""

    def __init__(self, *, max_checkpoints: int = 2) -> None:
        super().__init__()
        self.max_checkpoints = max(1, max_checkpoints)

    def _prune(self, thread_id: str, checkpoint_ns: str) -> None:
        storage = cast(Any, self).storage
        blobs = cast(Any, self).blobs
        writes = cast(Any, self).writes
        serde = cast(Any, self).serde

        checkpoints = storage[thread_id][checkpoint_ns]
        if len(checkpoints) <= self.max_checkpoints:
            return

        to_remove = sorted(checkpoints.keys())[: -self.max_checkpoints]
        for checkpoint_id in to_remove:
            checkpoint_bytes, _, _ = checkpoints.pop(checkpoint_id)
            try:
                checkpoint_data = serde.loads_typed(checkpoint_bytes)
                channel_versions = cast(
                    dict[str, Any], checkpoint_data.get("channel_versions", {}) or {}
                )
            except Exception:  # pragma: no cover - defensive cleanup
                channel_versions = {}

            for channel, version in channel_versions.items():
                blobs.pop((thread_id, checkpoint_ns, channel, version), None)
            writes.pop((thread_id, checkpoint_ns, checkpoint_id), None)

    def put(  # type: ignore[override]
        self,
        config: dict[str, Any],
        checkpoint: dict[str, Any],
        metadata: dict[str, Any],
        new_versions: dict[str, Any],
    ) -> dict[str, Any]:
        base_put = cast(Any, super())
        updated = cast(
            dict[str, Any],
            base_put.put(config, checkpoint, metadata, new_versions),
        )
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        self._prune(thread_id, checkpoint_ns)
        return updated


class PlanGenerationState(TypedDict):
    """State structure for plan generation workflow.

    Attributes:
        project: The project being modified
        plan: The plan containing instructions
        contexts: List of context files
        prompt: User's instruction prompt
        analyzed_requirements: Structured requirements from analysis
        generated_changes: List of generated code changes
        validation_result: Validation status and messages
        retry_count: Number of retry attempts
        error: Error message if any
    """

    project: Project
    plan: Plan
    contexts: list[Context]
    context_summary: str
    context_dependencies: dict[str, list[str]]
    context_relevance: dict[str, float]
    context_analysis_error: str | None
    actor_name: str | None
    actor_options: dict[str, Any]
    actor_graph_descriptor: dict[str, Any] | None
    actor_initial_context: dict[str, Any]
    prompt: str
    analyzed_requirements: dict[str, Any]
    generated_changes: list[Change]
    validation_result: dict[str, Any]
    retry_count: int
    error: str | None


class PlanGenerationGraph:
    """LangGraph workflow for generating code plans.

    This workflow orchestrates the plan generation process through multiple nodes:
    1. load_context: Loads and prepares context information
    2. analyze_requirements: Analyzes user prompt for requirements
    3. generate_plan: Generates code changes based on requirements
    4. validate: Validates generated changes

    The workflow includes conditional edges for retry logic and checkpointing
    for resumable execution.
    """

    def __init__(
        self,
        llm: BaseLanguageModel | None = None,
        max_retries: int = 3,
        context_llm: BaseLanguageModel | None = None,
        checkpoint_limit: int = 2,
    ):
        """Initialize the plan generation graph.

        Args:
            llm: Language model used for requirements, generation, and validation
            max_retries: Maximum number of retry attempts
            context_llm: Optional language model dedicated to context analysis
            checkpoint_limit: Maximum checkpoints to retain per thread
        """
        self.max_retries = max(1, max_retries)

        # Initialize LLMs - an LLM must be provided explicitly
        if llm is None:
            raise ValueError(
                "No LLM provider configured. "
                "Set provider credentials or enable core.mock_providers for testing."
            )
        self.llm = llm
        self.context_llm = context_llm or llm

        # Create prompts
        self._create_prompts()

        # Build the graph
        self.graph = self._build_graph()

        # Compile with checkpointing (bounded to reduce overhead)
        self.checkpointer = BoundedMemorySaver(max_checkpoints=checkpoint_limit)
        self.app = self.graph.compile(checkpointer=self.checkpointer)

    def _create_prompts(self) -> None:
        """Create prompt templates for each workflow node.

        System-level instructions are augmented with boundary marker
        recognition directives (mechanism 2 — prompt boundary markers).
        User content placeholders are wrapped with boundary markers.
        """
        _bi = _SANITIZER.BOUNDARY_INSTRUCTION
        _bs = _SANITIZER.BOUNDARY_START
        _be = _SANITIZER.BOUNDARY_END

        # Requirements analysis prompt
        self.analyze_prompt: Any = PromptTemplate(
            input_variables=["prompt", "context_summary"],
            template=(
                f"{_bi}\n\n"
                "You are a requirements analyst. "
                "Extract concise technical requirements from the request "
                "and context.\n\n"
                f"Request: {_bs}\n{{prompt}}\n{_be}\n"
                f"Context:\n{_bs}\n{{context_summary}}\n{_be}\n\n"
                "Return:\n"
                "- Key requirements\n"
                "- Files to modify or create\n"
                "- Needed dependencies\n"
                "- Risks or challenges\n"
            ),
        )

        # Code generation prompt
        self.generate_prompt: Any = PromptTemplate(
            input_variables=["requirements", "context_summary"],
            template=(
                f"{_bi}\n\n"
                "You are a code generator. "
                "Produce production-ready code that satisfies the requirements using "
                "the provided context.\n\n"
                f"Requirements:\n{_bs}\n{{requirements}}\n{_be}\n\n"
                f"Context:\n{_bs}\n{{context_summary}}\n{_be}\n\n"
                "Return concise, documented code that follows best practices.\n"
            ),
        )

        # Validation prompt
        self.validate_prompt: Any = PromptTemplate(
            input_variables=["generated_code"],
            template=(
                f"{_bi}\n\n"
                "You are a reviewer. "
                "Validate the generated code for correctness and safety.\n\n"
                f"Code:\n{_bs}\n{{generated_code}}\n{_be}\n\n"
                "Check syntax, logic, best practices, security, and performance. "
                "Respond with PASS or FAIL and list issues found.\n"
            ),
        )

    def _chain_with_retry(self, chain: Any) -> Any:
        """Wrap a LangChain runnable with the configured retry attempts."""

        attempts = max(1, self.max_retries)
        if hasattr(chain, "with_retry"):
            return chain.with_retry(stop_after_attempt=attempts)
        return chain

    def _build_graph(self) -> StateGraph:
        """Build the state graph for plan generation workflow.

        Returns:
            Configured state graph
        """
        # Create graph
        workflow = StateGraph(PlanGenerationState)

        # Add nodes
        workflow.add_node("load_context", self._load_context)
        workflow.add_node("analyze_requirements", self._analyze_requirements)
        workflow.add_node("generate_plan", self._generate_plan)
        workflow.add_node("validate", self._validate)

        # Set entry point
        workflow.set_entry_point("load_context")

        # Add edges
        workflow.add_edge("load_context", "analyze_requirements")
        workflow.add_edge("analyze_requirements", "generate_plan")
        workflow.add_edge("generate_plan", "validate")

        # Add conditional edge for retry logic
        workflow.add_conditional_edges(
            "validate",
            self._should_retry,
            {
                "retry": "analyze_requirements",
                "end": END,
            },
        )

        return workflow

    def _load_context(self, state: PlanGenerationState) -> dict[str, Any]:
        """Load and prepare context information.

        Args:
            state: Current workflow state

        Returns:
            Updated state with context loaded
        """
        contexts = state.get("contexts", [])
        analysis = self._analyze_contexts(contexts)

        return {
            "retry_count": 0,
            "error": None,
            **analysis,
        }

    def _analyze_requirements(self, state: PlanGenerationState) -> dict[str, Any]:
        """Analyze user prompt to extract requirements.

        Args:
            state: Current workflow state

        Returns:
            Updated state with analyzed requirements
        """
        # Prepare context summary
        contexts = state.get("contexts", [])
        context_summary = state.get("context_summary") or self._format_context_summary(
            contexts
        )

        # Create analysis chain
        chain = self._chain_with_retry(
            self.analyze_prompt | self.llm | StrOutputParser()
        )

        try:
            # Run analysis
            analysis = chain.invoke(
                {
                    "prompt": state["prompt"],
                    "context_summary": context_summary,
                }
            )

            # Parse requirements (simplified for now)
            requirements: dict[str, Any] = {
                "description": str(analysis),
                "files_to_modify": [ctx.path for ctx in contexts],
                "operation": "modify" if contexts else "create",
            }

            return {
                "analyzed_requirements": requirements,
                "error": None,
            }

        except Exception as e:
            return {
                "analyzed_requirements": {},
                "error": f"Requirements analysis failed: {e!s}",
            }

    def _generate_plan(self, state: PlanGenerationState) -> dict[str, Any]:
        """Generate code changes based on requirements.

        Args:
            state: Current workflow state

        Returns:
            Updated state with generated changes
        """
        requirements = state.get("analyzed_requirements", {})

        if not requirements:
            return {
                "generated_changes": [],
                "error": "No requirements to generate from",
            }

        # Prepare context summary
        contexts = state.get("contexts", [])
        context_summary = state.get("context_summary") or self._format_context_summary(
            contexts
        )

        # Create generation chain
        chain = self._chain_with_retry(
            self.generate_prompt | self.llm | StrOutputParser()
        )

        try:
            # Generate code
            result = chain.invoke(
                {
                    "requirements": requirements.get("description", ""),
                    "context_summary": context_summary,
                }
            )
            generated_code = str(result)

            # Determine file path and operation
            operation = requirements.get("operation", "create")
            contexts = state.get("contexts", [])
            explicit_path = self._extract_path_from_prompt(state.get("prompt", ""))

            if operation == "modify" and contexts:
                matched_context = self._find_matching_context(contexts, explicit_path)
                target_context = matched_context or contexts[0]
                file_path = target_context.path
                operation_type = OperationType.MODIFY
                original_content = target_context.content
            elif explicit_path:
                file_path = explicit_path
                operation_type = (
                    OperationType.MODIFY if contexts else OperationType.CREATE
                )
                original_content = None

                existing_path = Path(file_path)
                if existing_path.exists():
                    if existing_path.is_dir():
                        file_path = str(existing_path / "generated.py")
                        operation_type = OperationType.CREATE
                        original_content = None
                    else:
                        operation_type = OperationType.MODIFY
                        try:
                            original_content = existing_path.read_text()
                        except Exception:  # pragma: no cover - best effort
                            original_content = None
                elif existing_path.suffix == "":
                    file_path = str(existing_path / "generated.py")
                    operation_type = OperationType.CREATE
                    original_content = None
            else:
                prompt_lower = state["prompt"].lower()
                if "test" in prompt_lower:
                    file_path = "test_generated.py"
                elif "error" in prompt_lower or "exception" in prompt_lower:
                    file_path = "error_handler.py"
                else:
                    file_path = "generated.py"
                operation_type = OperationType.CREATE
                original_content = None

            # Create change object
            plan_id = state["plan"].id if state["plan"].id else 0

            changes = [
                Change(
                    id=None,
                    plan_id=plan_id,
                    file_path=file_path,
                    operation=operation_type,
                    original_content=original_content,
                    new_content=generated_code,
                    applied=False,
                    applied_at=None,
                    new_path=None,
                )
            ]

            return {
                "generated_changes": changes,
                "error": None,
            }

        except Exception as e:
            return {
                "generated_changes": [],
                "error": f"Code generation failed: {e!s}",
            }

    def _extract_path_from_prompt(self, prompt: str) -> str | None:
        """Return a file or directory path only when explicitly prefixed with @."""

        match = PATH_HINT_RE.search(prompt or "")
        return match.group("path") if match else None

    def _find_matching_context(
        self, contexts: list[Context], file_path: str | None
    ) -> Context | None:
        """Return a context whose path matches the hinted file."""

        if not file_path:
            return None

        for ctx in contexts:
            ctx_path = ctx.path
            if ctx_path and (ctx_path == file_path or ctx_path.endswith(file_path)):
                return ctx
        return None

    def _validate(self, state: PlanGenerationState) -> dict[str, Any]:
        """Validate generated code changes.

        Args:
            state: Current workflow state

        Returns:
            Updated state with validation results
        """
        changes = state.get("generated_changes", [])

        if not changes:
            return {
                "validation_result": {
                    "status": "FAIL",
                    "message": "No changes to validate",
                },
            }

        # Create validation chain
        chain = self._chain_with_retry(
            self.validate_prompt | self.llm | StrOutputParser()
        )

        try:
            # Validate all changes
            all_code = "\n\n".join(
                f"File: {change.file_path}\n{change.new_content}" for change in changes
            )

            result = chain.invoke(
                {
                    "generated_code": all_code,
                }
            )
            validation = str(result)

            # Simple validation check (in real implementation, parse the LLM response)
            is_valid = "PASS" in validation.upper() or len(all_code) > 10

            return {
                "validation_result": {
                    "status": "PASS" if is_valid else "FAIL",
                    "message": validation,
                },
            }

        except Exception as e:
            return {
                "validation_result": {
                    "status": "FAIL",
                    "message": f"Validation failed: {e!s}",
                },
            }

    def _should_retry(self, state: PlanGenerationState) -> str:
        """Determine if workflow should retry based on validation.

        Args:
            state: Current workflow state

        Returns:
            "retry" or "end" based on validation and retry count
        """
        validation = state.get("validation_result", {})
        retry_count = state.get("retry_count", 0)

        # Check if validation failed and retries available
        if validation.get("status") == "FAIL" and retry_count < self.max_retries:
            # Increment retry count
            state["retry_count"] = retry_count + 1
            return "retry"

        return "end"

    def _format_context_summary(self, contexts: list[Context]) -> str:
        """Format context files into a summary string.

        Args:
            contexts: List of context files

        Returns:
            Formatted context summary
        """
        if not contexts:
            return "No context files provided"

        summary_parts: list[str] = []
        for ctx in contexts[:5]:  # Limit to first 5 files
            content_preview = ctx.content[:300] if ctx.content else ""
            summary_parts.append(f"File: {ctx.path}\nPreview: {content_preview}...\n")

        if len(contexts) > 5:
            summary_parts.append(f"... and {len(contexts) - 5} more files")

        return "\n".join(summary_parts)

    def _analyze_contexts(self, contexts: list[Context]) -> dict[str, Any]:
        """Run ContextAnalysisAgent to enrich the workflow state."""

        fallback_summary = self._format_context_summary(contexts)
        if not contexts:
            return {
                "context_summary": fallback_summary,
                "context_dependencies": {},
                "context_relevance": {},
                "context_analysis_error": None,
            }

        documents: list[Document] = []
        for ctx in contexts:
            if ctx.content:
                documents.append(
                    Document(
                        page_content=ctx.content,
                        metadata={"source": ctx.path},
                    )
                )

        initial_state: ContextAnalysisState = {
            "file_paths": [ctx.path for ctx in contexts if ctx.path],
            "documents": documents,
            "dependencies": {},
            "summary": "",
            "relevance_scores": {},
            "chunks": [],
            "error": None,
        }

        agent = ContextAnalysisAgent(llm=self.context_llm)
        config = {"configurable": {"thread_id": f"context-analysis-{uuid4()}"}}

        try:
            analysis = agent.invoke(initial_state, config)
            summary = analysis.get("summary") or fallback_summary
            dependencies = analysis.get("dependencies", {})
            relevance = analysis.get("relevance_scores", {})
            analysis_error = analysis.get("error")
        except Exception as exc:  # pragma: no cover - defensive fallback
            summary = fallback_summary
            dependencies = {}
            relevance = {}
            analysis_error = f"Context analysis failed: {exc!s}"

        return {
            "context_summary": summary,
            "context_dependencies": dependencies,
            "context_relevance": relevance,
            "context_analysis_error": analysis_error,
        }

    def invoke(
        self,
        project: Project,
        plan: Plan,
        contexts: list[Context],
        thread_id: str = "default",
        actor_context: ActorInvocationContext | None = None,
    ) -> dict[str, Any]:
        """Execute the plan generation workflow synchronously.

        Args:
            project: The project being modified
            plan: The plan containing instructions
            contexts: List of context files
            thread_id: Thread ID for checkpointing
            actor_context: Optional actor metadata/options to include in state

        Returns:
            Final workflow state
        """
        initial_state: dict[str, Any] = {
            "project": project,
            "plan": plan,
            "contexts": contexts,
            "context_summary": "",
            "context_dependencies": {},
            "context_relevance": {},
            "context_analysis_error": None,
            "actor_name": actor_context.name if actor_context else None,
            "actor_options": actor_context.options if actor_context else {},
            "actor_graph_descriptor": actor_context.graph_descriptor
            if actor_context
            else None,
            "actor_initial_context": actor_context.initial_context
            if actor_context
            else {},
            "prompt": plan.prompt or "",
            "analyzed_requirements": {},
            "generated_changes": [],
            "validation_result": {},
            "retry_count": 0,
            "error": None,
        }

        config = {"configurable": {"thread_id": thread_id}}

        result = self.app.invoke(initial_state, config)
        return result

    async def ainvoke(
        self,
        project: Project,
        plan: Plan,
        contexts: list[Context],
        thread_id: str = "default",
        actor_context: ActorInvocationContext | None = None,
    ) -> dict[str, Any]:
        """Execute the plan generation workflow asynchronously.

        Args:
            project: The project being modified
            plan: The plan containing instructions
            contexts: List of context files
            thread_id: Thread ID for checkpointing
            actor_context: Optional actor metadata/options to include in state

        Returns:
            Final workflow state
        """
        initial_state: dict[str, Any] = {
            "project": project,
            "plan": plan,
            "contexts": contexts,
            "context_summary": "",
            "context_dependencies": {},
            "context_relevance": {},
            "context_analysis_error": None,
            "actor_name": actor_context.name if actor_context else None,
            "actor_options": actor_context.options if actor_context else {},
            "actor_graph_descriptor": actor_context.graph_descriptor
            if actor_context
            else None,
            "actor_initial_context": actor_context.initial_context
            if actor_context
            else {},
            "prompt": plan.prompt or "",
            "analyzed_requirements": {},
            "generated_changes": [],
            "validation_result": {},
            "retry_count": 0,
            "error": None,
        }

        config = {"configurable": {"thread_id": thread_id}}

        result = await self.app.ainvoke(initial_state, config)
        return result

    def stream(
        self,
        project: Project,
        plan: Plan,
        contexts: list[Context],
        thread_id: str = "default",
        actor_context: ActorInvocationContext | None = None,
    ) -> Iterator[dict[str, Any]]:
        """Stream the plan generation workflow execution.

        Args:
            project: The project being modified
            plan: The plan containing instructions
            contexts: List of context files
            thread_id: Thread ID for checkpointing

        Yields:
            Streaming workflow events
        """
        initial_state: dict[str, Any] = {
            "project": project,
            "plan": plan,
            "contexts": contexts,
            "context_summary": "",
            "context_dependencies": {},
            "context_relevance": {},
            "context_analysis_error": None,
            "actor_name": actor_context.name if actor_context else None,
            "actor_options": actor_context.options if actor_context else {},
            "actor_graph_descriptor": actor_context.graph_descriptor
            if actor_context
            else None,
            "actor_initial_context": actor_context.initial_context
            if actor_context
            else {},
            "prompt": plan.prompt or "",
            "analyzed_requirements": {},
            "generated_changes": [],
            "validation_result": {},
            "retry_count": 0,
            "error": None,
        }

        config = {"configurable": {"thread_id": thread_id}}

        yield from self.app.stream(initial_state, config)
