# Copyright(C) 2024-2025 Advanced Micro Devices, Inc. All rights reserved.
# SPDX-License-Identifier: MIT
"""Checklist Generator for LLM-Driven Code Generation.

This module uses an LLM to generate a checklist of template invocations
based on the user's request and the available template catalog.

The generator:
1. Receives user request and project context
2. Sends prompt to LLM with template catalog
3. Parses LLM response into structured checklist
4. Validates checklist items against template definitions

The resulting checklist is then executed deterministically by
ChecklistExecutor.
"""

import json
import logging
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Protocol

from .steps.base import UserContext
from .template_catalog import get_catalog_prompt, validate_checklist_item

logger = logging.getLogger(__name__)


class ChatSDK(Protocol):
    """Protocol for chat SDK interface."""

    def send(self, message: str, timeout: int = 600, no_history: bool = False) -> Any:
        """Send a message and get response."""
        ...


@dataclass
class ChecklistItem:
    """Single item in the generated checklist.

    Represents a template invocation with its parameters and
    the LLM's reasoning for including it.
    """

    template: str
    params: Dict[str, Any]
    description: str

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary representation."""
        return {
            "template": self.template,
            "params": self.params,
            "description": self.description,
        }


@dataclass
class GeneratedChecklist:
    """Complete checklist generated by LLM.

    Contains the list of template invocations and the LLM's
    overall reasoning for the chosen approach.
    """

    items: List[ChecklistItem]
    reasoning: str
    raw_response: str = ""
    validation_errors: List[str] = field(default_factory=list)

    @property
    def is_valid(self) -> bool:
        """Check if checklist passed validation."""
        return len(self.validation_errors) == 0

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary representation."""
        return {
            "reasoning": self.reasoning,
            "checklist": [item.to_dict() for item in self.items],
            "is_valid": self.is_valid,
            "validation_errors": self.validation_errors,
        }


@dataclass
class ProjectState:
    """Current state of the project for context."""

    exists: bool = False
    has_package_json: bool = False
    has_prisma: bool = False
    has_next_config: bool = False
    existing_models: List[str] = field(default_factory=list)
    existing_routes: List[str] = field(default_factory=list)
    existing_pages: List[str] = field(default_factory=list)

    def to_prompt(self) -> str:
        """Generate prompt-friendly description of project state."""
        if not self.exists:
            return "Project does not exist yet - will be created fresh."

        lines = ["Current project state:"]

        if self.has_package_json:
            lines.append("- ✓ package.json exists (Node.js project)")
        if self.has_next_config:
            lines.append("- ✓ next.config.ts exists (Next.js configured)")
        if self.has_prisma:
            lines.append("- ✓ Prisma configured")
            if self.existing_models:
                lines.append(f"  - Models: {', '.join(self.existing_models)}")

        if self.existing_routes:
            lines.append(f"- Existing API routes: {', '.join(self.existing_routes)}")

        if self.existing_pages:
            lines.append(f"- Existing pages: {', '.join(self.existing_pages)}")

        return "\n".join(lines)


CHECKLIST_SYSTEM_PROMPT = """You are a code generation planner. Your task is to analyze the user's request and generate a checklist of template invocations that will create the requested application.

{catalog_prompt}

## Instructions

1. Analyze the user's request carefully
2. Consider what the user ACTUALLY wants (semantic understanding)
3. Select templates that will fulfill the request
4. Add semantic enhancements based on the request type:
   - For "todo" apps: add checkboxes for completion status
   - For "blog" apps: add date fields for posts
   - For "e-commerce": add price, inventory fields
5. Ensure dependencies are satisfied (run setup before data, data before API, API before UI)
6. Generate a complete checklist that creates a working application
7. When follow-up fixes are requested, use `fix_code` to repair the specific files called out by validation/test logs. Extract the precise file paths and line numbers from the errors (see the Raw Validation Logs) and pass those line numbers inside the error description so the fixer knows exactly where to focus. Always reference the latest findings to decide which fixes to schedule before running validations again.

This workflow repeats until all validations pass, so each checklist should either advance new functionality or explicitly repair the failures reported in the latest validation logs.

## IMPORTANT: Complete CRUD Applications

For any app that manages resources (todos, posts, users, etc.), you MUST generate ALL of these UI components:

1. **List page** (variant: "list") - Main page showing all items at /resources
2. **Form component** (variant: "form") - Reusable form for create
3. **New page** (variant: "new") - Create new item at /resources/new
4. **Edit page** (variant: "detail") - Edit single item at /resources/[id] with pre-populated form

Missing any of these will result in broken navigation! Always include all 4 variants.

## REQUIRED: Setup and Validation Commands

**CRITICAL**: The following commands are REQUIRED for a valid plan:

1. **setup_app_styling** MUST be included after creating the application (after `create_next_app`). This configures app-wide styling with modern dark theme design system.

2. **setup_testing** MUST be included after `setup_app_styling`. This sets up the testing infrastructure.

3. **generate_style_tests** MUST be included after `setup_testing`. This generates CSS integrity tests that validate the design system.

4. **The final 2 commands MUST be in this exact order:**
   - Second-to-last: `run_typescript_check` (validates TypeScript compilation)
   - Last: `validate_styles` (validates CSS files and design system)

These setup and validation commands are mandatory - a plan without them is INVALID.

## Output Format

Respond with ONLY a JSON object (no markdown code blocks):
{{
  "reasoning": "Brief explanation of your approach and any semantic enhancements",
  "checklist": [
    {{"template": "template_name", "params": {{}}, "description": "Why this step is needed"}}
  ]
}}

Important:
- Use exact template names from the catalog
- Provide all required parameters
- Order items by dependency (setup first, then data, then API, then UI)
- Add semantic enhancements that make the app intuitive (e.g., checkboxes for todos)
- For CRUD apps, ALWAYS include all 4 UI variants: list, form, new, detail
- REQUIRED: Include `setup_app_styling` after `create_next_app`
- REQUIRED: Include `setup_testing` after `setup_app_styling`
- REQUIRED: Include `generate_style_tests` after `setup_testing`
- REQUIRED: End with `run_typescript_check`, then `validate_styles` as the last 2 commands
- When converting a raw validation error into `fix_code`, copy the exact snippet (file, line, column, and message). For example:

  Raw Validation Logs (example):
  ```
  {{"template": "run_typescript_check", "output": {{"errors": "path/to/File.tsx(10,5): error TS1234: <error text>\\n"}}}}
  ```

  Corresponding checklist item:
  ```
  {{
    "template": "fix_code",
    "params": {{
      "file_path": "path/to/File.tsx",
      "error_description": "path/to/File.tsx(10,5): error TS1234: <error text>"
    }},
    "description": "Fix the TypeScript compiler error reported for File.tsx."
  }}
  ```
  Always keep the error text verbatim so the fixer knows exactly where to edit."""


class ChecklistGenerator:
    """Generate execution checklist using LLM.

    The generator sends the user request, project state, and template
    catalog to an LLM, which returns a structured checklist of template
    invocations.
    """

    def __init__(self, chat_sdk: ChatSDK):
        """Initialize the checklist generator.

        Args:
            chat_sdk: Chat SDK instance for LLM communication
        """
        self.chat = chat_sdk

    def generate_initial_checklist(
        self,
        context: UserContext,
        project_state: Optional[ProjectState] = None,
    ) -> GeneratedChecklist:
        """Generate the initial project-scaffolding checklist."""
        if project_state is None:
            project_state = ProjectState()

        system_prompt = CHECKLIST_SYSTEM_PROMPT.format(
            catalog_prompt=get_catalog_prompt()
        )
        user_prompt = self._build_initial_prompt(context, project_state)
        full_prompt = f"{system_prompt}\n\n## User Request\n\n{user_prompt}"
        return self._generate_from_prompt(full_prompt)

    def generate_debug_checklist(
        self,
        context: UserContext,
        project_state: Optional[ProjectState],
        prior_errors: Optional[List[str]],
        validation_logs: Optional[List[Any]],
    ) -> GeneratedChecklist:
        """Generate a remediation checklist to fix outstanding errors."""
        if project_state is None:
            project_state = ProjectState()

        debug_prompt = self._build_debug_prompt(
            context=context,
            project_state=project_state,
            prior_errors=prior_errors or [],
            validation_logs=validation_logs or [],
        )
        system_prompt = CHECKLIST_SYSTEM_PROMPT.format(
            catalog_prompt=get_catalog_prompt()
        )
        full_prompt = f"{system_prompt}\n\n## Remediation Context\n\n{debug_prompt}"
        return self._generate_from_prompt(full_prompt)

    def _generate_from_prompt(self, full_prompt: str) -> GeneratedChecklist:
        """Common checklist generation logic with retries."""
        logger.debug("Generating checklist with LLM...")
        logger.debug(f"Checklist prompt: {full_prompt}")

        max_attempts = 3
        last_failure_reason = "unknown error"

        for attempt in range(1, max_attempts + 1):
            try:
                response = self.chat.send(full_prompt, timeout=1200)

                response_text = self._extract_response_text(response)

                logger.debug(f"LLM response (attempt {attempt}): {response_text}")

                checklist = self._parse_checklist(response_text)
            except Exception as exc:  # pylint: disable=broad-exception-caught
                last_failure_reason = str(exc)
                logger.warning(
                    "Checklist generation attempt %d/%d failed: %s",
                    attempt,
                    max_attempts,
                    exc,
                )
                continue

            if not checklist.items:
                last_failure_reason = "LLM returned an empty checklist"
                logger.warning(
                    "Checklist generation attempt %d/%d returned no items, retrying...",
                    attempt,
                    max_attempts,
                )
                continue

            self._validate_checklist(checklist)
            if checklist.validation_errors:
                last_failure_reason = "; ".join(checklist.validation_errors)
                logger.warning(
                    "Checklist generation attempt %d/%d failed validation: %s",
                    attempt,
                    max_attempts,
                    checklist.validation_errors,
                )
                continue

            logger.debug(
                "Generated checklist with %d items on attempt %d",
                len(checklist.items),
                attempt,
            )
            return checklist

        raise RuntimeError(
            f"Failed to generate a valid checklist after {max_attempts} attempts: "
            f"{last_failure_reason}"
        )

    def _build_initial_prompt(
        self,
        context: UserContext,
        project_state: ProjectState,
    ) -> str:
        """Build the user prompt with all context.

        Args:
            context: User context
            project_state: Current project state

        Returns:
            Formatted user prompt string
        """
        lines = [f"**User Request**: {context.user_request}"]

        if context.entity_name:
            lines.append(f"\n**Inferred Entity**: {context.entity_name}")

        if context.schema_fields:
            lines.append(f"\n**Inferred Fields**: {json.dumps(context.schema_fields)}")

        lines.append(f"\n**Project Directory**: {context.project_dir}")
        lines.append(f"\n**Language**: {context.language}")
        lines.append(f"\n**Project Type**: {context.project_type}")

        lines.append(f"\n{project_state.to_prompt()}")

        if context.fix_feedback:
            lines.append("\n**Outstanding Fix Requests**:")
            for note in context.fix_feedback[-5:]:
                lines.append(f"- {note}")

        if context.validation_reports:
            lines.append("\n**Recent Validation/Test Findings**:")
            for log in context.validation_reports[-5:]:
                status = "PASS" if log.get("success", True) else "FAIL"
                template = log.get("template", "validation_step")
                description = log.get("description", "")
                lines.append(f"- [{status}] {template}: {description}")
                if log.get("error"):
                    lines.append(f"  Error: {log['error']}")

                output = log.get("output", {})
                snippet = ""
                if isinstance(output, dict):
                    for key in ("stdout", "stderr", "message", "details"):
                        if output.get(key):
                            snippet = str(output[key])[:200]
                            break
                    if not snippet and output:
                        snippet = json.dumps(output)[:200]
                elif output:
                    snippet = str(output)[:200]
                if snippet:
                    lines.append(f"  Output: {snippet}")

        lines.append("\nGenerate a checklist to fulfill this request.")

        return "\n".join(lines)

    def _build_debug_prompt(
        self,
        context: UserContext,
        project_state: ProjectState,
        prior_errors: List[str],
        validation_logs: List[Any],
    ) -> str:
        """Build prompt for remediation/debug checklists."""
        lines = [
            "You are a remediation planner for the GAIA web development agent. "
            "The project has already been scaffolded; focus exclusively on fixing outstanding issues."
        ]
        lines.append(f"\n**User Request**: {context.user_request}")
        lines.append(f"\n**Project Directory**: {context.project_dir}")

        if context.entity_name:
            lines.append(f"\n**Entity**: {context.entity_name}")
        if context.schema_fields:
            lines.append(f"\n**Schema Fields**: {json.dumps(context.schema_fields)}")

        lines.append(f"\n{project_state.to_prompt()}")

        if prior_errors:
            lines.append("\n**Execution Errors From Last Attempt:**")
            for err in prior_errors:
                lines.append(f"- {err}")

        if validation_logs:
            lines.append("\n**Recent Validation/Test Results:**")
            raw_entries = []
            for log in validation_logs[-10:]:
                entry = log.to_dict() if hasattr(log, "to_dict") else log
                template = entry.get("template", "unknown_step")
                success = entry.get("success", True)
                desc = entry.get("description", "")
                status = "PASS" if success else "FAIL"
                lines.append(f"- [{status}] {template}: {desc}")
                if entry.get("error"):
                    lines.append(f"  Error: {entry['error']}")
                output = entry.get("output") or {}
                for key in ("stdout", "stderr", "details", "message"):
                    if output.get(key):
                        snippet = str(output[key])[:200]
                        lines.append(f"  Output: {snippet}")
                        break
                raw_entries.append(entry)

            if raw_entries:
                lines.append(
                    "\n**Raw Validation Logs (exact text for follow-up fixes):**"
                )
                for entry in raw_entries:
                    lines.append(json.dumps(entry, ensure_ascii=False))

        if context.fix_feedback:
            lines.append("\n**Outstanding Fix Instructions:**")
            for note in context.fix_feedback[-10:]:
                lines.append(f"- {note}")

        lines.append(
            "\nYour job: draft a concise checklist that repairs the errors above, "
            "regenerates any broken code, and re-runs critical validations."
        )
        lines.append(
            "\n**Critical Requirements for Debug Checklists:**\n"
            "1. Use `fix_code` to repair the specific files referenced in the failures above.\n"
            "2. Re-run any validations or tests that previously failed once fixes are applied.\n"
            "3. Always include `run_typescript_check` as the second-to-last command to capture current compiler errors.\n"
            "4. Always include `validate_styles` as the final command to capture CSS/design regressions."
        )

        return "\n".join(lines)

    def _extract_response_text(self, response: Any) -> str:
        """Extract text from LLM response.

        Handles different response formats from various SDKs.

        Args:
            response: Response from chat SDK

        Returns:
            Response text string
        """
        if isinstance(response, str):
            return response

        # Handle response objects with text attribute
        if hasattr(response, "text"):
            return response.text

        # Handle response objects with content attribute
        if hasattr(response, "content"):
            return response.content

        # Handle dict-like responses
        if isinstance(response, dict):
            return response.get("text", response.get("content", str(response)))

        return str(response)

    def _parse_checklist(self, response_text: str) -> GeneratedChecklist:
        """Parse LLM response into GeneratedChecklist.

        Args:
            response_text: Raw LLM response text

        Returns:
            Parsed GeneratedChecklist
        """
        try:
            # Try to extract JSON from the response
            json_str = self._extract_json(response_text)

            data = json.loads(json_str)

            # Parse items
            items = []
            for item_data in data.get("checklist", []):
                item = ChecklistItem(
                    template=item_data.get("template", ""),
                    params=item_data.get("params", {}),
                    description=item_data.get("description", ""),
                )
                items.append(item)

            return GeneratedChecklist(
                items=items,
                reasoning=data.get("reasoning", ""),
                raw_response=response_text,
            )

        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse checklist JSON: {e}")
            return GeneratedChecklist(
                items=[],
                reasoning="",
                raw_response=response_text,
                validation_errors=[f"Failed to parse JSON: {str(e)}"],
            )

    def _extract_json(self, text: str) -> str:
        """Extract JSON from text that might contain markdown or other content.

        Args:
            text: Text that may contain JSON

        Returns:
            Extracted JSON string
        """
        # Try to find JSON in markdown code block
        code_block_match = re.search(r"```(?:json)?\s*\n?(.*?)\n?```", text, re.DOTALL)
        if code_block_match:
            return code_block_match.group(1).strip()

        # Try to find raw JSON object
        json_match = re.search(r"\{.*\}", text, re.DOTALL)
        if json_match:
            return json_match.group(0)

        # Return as-is and let JSON parser handle it
        return text.strip()

    def _validate_checklist(self, checklist: GeneratedChecklist) -> None:
        """Validate checklist items against template definitions.

        Adds validation errors to the checklist if any are found.

        Args:
            checklist: Checklist to validate (modified in place)
        """
        for item in checklist.items:
            errors = validate_checklist_item(item.template, item.params)
            checklist.validation_errors.extend(errors)

        # Check for duplicate templates (some are ok, like multiple API routes)
        seen_templates = {}
        for item in checklist.items:
            key = f"{item.template}:{json.dumps(item.params, sort_keys=True)}"
            if key in seen_templates:
                checklist.validation_errors.append(
                    f"Duplicate checklist item: {item.template} with same params"
                )
            seen_templates[key] = True

        # Validate required setup: setup_app_styling must come after create_next_app
        create_app_index = None
        setup_styling_index = None
        setup_testing_index = None
        for i, item in enumerate(checklist.items):
            if item.template == "create_next_app":
                create_app_index = i
            if item.template == "setup_app_styling":
                setup_styling_index = i
            if item.template == "setup_testing":
                setup_testing_index = i

        if create_app_index is not None:
            if setup_styling_index is None:
                checklist.validation_errors.append(
                    "REQUIRED: 'setup_app_styling' must be included after 'create_next_app'"
                )
            elif setup_styling_index <= create_app_index:
                checklist.validation_errors.append(
                    "REQUIRED: 'setup_app_styling' must come after 'create_next_app' in the checklist"
                )

        # Validate required testing setup: setup_testing must come after setup_app_styling
        if setup_styling_index is not None:
            if setup_testing_index is None:
                checklist.validation_errors.append(
                    "REQUIRED: 'setup_testing' must be included after 'setup_app_styling'"
                )
            elif setup_testing_index <= setup_styling_index:
                checklist.validation_errors.append(
                    "REQUIRED: 'setup_testing' must come after 'setup_app_styling' in the checklist"
                )

        # Validate required final validation commands: run_typescript_check, validate_styles
        if len(checklist.items) < 2:
            checklist.validation_errors.append(
                "REQUIRED: Checklist must end with 'run_typescript_check', "
                "'validate_styles' as the last two commands"
            )
        else:
            last_item = checklist.items[-1]
            second_last_item = checklist.items[-2]

            if last_item.template != "validate_styles":
                checklist.validation_errors.append(
                    "REQUIRED: The last command must be 'validate_styles'"
                )
            if second_last_item.template != "run_typescript_check":
                checklist.validation_errors.append(
                    "REQUIRED: The second-to-last command must be 'run_typescript_check'"
                )

        # Validate generate_style_tests is included (after setup_testing)
        generate_style_tests_index = None
        for i, item in enumerate(checklist.items):
            if item.template == "generate_style_tests":
                generate_style_tests_index = i

        if setup_testing_index is not None and generate_style_tests_index is None:
            checklist.validation_errors.append(
                "REQUIRED: 'generate_style_tests' must be included after 'setup_testing'"
            )
        elif (
            generate_style_tests_index is not None
            and setup_testing_index is not None
            and generate_style_tests_index <= setup_testing_index
        ):
            checklist.validation_errors.append(
                "REQUIRED: 'generate_style_tests' must come after 'setup_testing'"
            )

        if checklist.validation_errors:
            logger.warning(
                f"Checklist validation errors: {checklist.validation_errors}"
            )


def create_checklist_from_workflow(
    workflow_phases: List[Any],
    context: UserContext,
) -> GeneratedChecklist:
    """Create a checklist from existing workflow phases (for comparison/testing).

    This converts the old step-based workflow into the new checklist format,
    useful for testing and migration.

    Args:
        workflow_phases: List of WorkflowPhase objects from factory
        context: User context

    Returns:
        GeneratedChecklist representing the workflow
    """
    items = []

    for phase in workflow_phases:
        for step in phase.steps:
            # Map step names to template names
            template_map = {
                "create_next_app": "create_next_app",
                "setup_styling": "setup_app_styling",
                "install_deps": "setup_prisma",
                "setup_testing": "setup_testing",
                "prisma_init": "setup_prisma",
                "setup_prisma": "setup_prisma",
                "manage_data_model": "generate_prisma_model",
                "manage_api_endpoint": "generate_api_route",
                "manage_api_endpoint_dynamic": "generate_api_route",
                "manage_react_component": "generate_react_component",
                "update_landing_page": "update_landing_page",
                "validate_typescript": "run_typescript_check",
                "validate_crud_structure": "run_typescript_check",
                "test_crud_api": "run_typescript_check",
            }

            template_name = template_map.get(step.name, step.name)

            # Extract params from step
            params = {}
            if hasattr(step, "get_tool_invocation"):
                invocation = step.get_tool_invocation(context)
                if invocation:
                    _, step_params = invocation
                    params = {
                        k: v for k, v in step_params.items() if k != "project_dir"
                    }

            items.append(
                ChecklistItem(
                    template=template_name,
                    params=params,
                    description=step.description,
                )
            )

    return GeneratedChecklist(
        items=items,
        reasoning="Converted from existing workflow",
    )
