student/code_generator.py

"""LLM-based code generator for creating web applications."""
import json
from pathlib import Path
from typing import Any

import anthropic
import openai

from shared.config import settings
from shared.logger import setup_logger
from shared.models import Attachment, TaskRequest
from shared.utils import decode_data_uri

logger = setup_logger(__name__)


class CodeGenerator:
    """Generate code using LLM based on task requirements."""

    def __init__(self) -> None:
        """Initialize code generator with LLM client."""
        self.provider = settings.llm_provider
        self.model = settings.llm_model

        if self.provider == "anthropic":
            self.client = anthropic.Anthropic(api_key=settings.anthropic_api_key)
        elif self.provider == "openai":
            self.client = openai.OpenAI(api_key=settings.openai_api_key)
        elif self.provider == "aipipe":
            # Use OpenAI client with AIPipe endpoints
            self.client = openai.OpenAI(
                api_key=settings.aipipe_token,
                base_url=settings.aipipe_base_url
            )
        else:
            raise ValueError(f"Unsupported LLM provider: {self.provider}")

        logger.info(f"Initialized CodeGenerator with {self.provider}/{self.model}")

    def generate_app(self, task: TaskRequest, output_dir: Path) -> dict[str, str]:
        """Generate application code based on task requirements.

        Args:
            task: Task request containing brief and requirements
            output_dir: Directory to save generated files

        Returns:
            Dictionary mapping filenames to their content
        """
        logger.info(f"Generating app for task {task.task}")

        # Prepare context with attachments
        attachment_info = self._prepare_attachments(task.attachments, output_dir)

        # Build prompt
        prompt = self._build_generation_prompt(task, attachment_info)

        # Generate code
        generated_files = self._generate_with_llm(prompt)

        # Save files
        self._save_files(generated_files, output_dir)

        logger.info(f"Generated {len(generated_files)} files for task {task.task}")
        return generated_files

    def _prepare_attachments(
        self, attachments: list[Attachment], output_dir: Path
    ) -> list[dict[str, str]]:
        """Decode and save attachments, return metadata.

        Args:
            attachments: List of attachments with data URIs
            output_dir: Directory to save attachments

        Returns:
            List of attachment metadata
        """
        attachment_info = []

        for att in attachments:
            try:
                mime_type, content = decode_data_uri(att.url)
                file_path = output_dir / att.name
                file_path.write_bytes(content)

                attachment_info.append(
                    {
                        "name": att.name,
                        "mime_type": mime_type,
                        "size": len(content),
                        "preview": content[:200].decode("utf-8", errors="ignore")
                        if mime_type.startswith("text/")
                        else "[binary data]",
                    }
                )
                logger.debug(f"Saved attachment {att.name} ({mime_type}, {len(content)} bytes)")
            except Exception as e:
                logger.error(f"Failed to process attachment {att.name}: {e}")
                attachment_info.append({"name": att.name, "error": str(e)})

        return attachment_info

    def _build_generation_prompt(
        self, task: TaskRequest, attachment_info: list[dict[str, str]]
    ) -> str:
        """Build prompt for LLM code generation.

        Args:
            task: Task request
            attachment_info: Attachment metadata

        Returns:
            Formatted prompt
        """
        attachments_section = ""
        if attachment_info:
            attachments_section = "\n\n**Attachments:**\n" + "\n".join(
                f"- {att['name']}: {att.get('mime_type', 'unknown')}"
                for att in attachment_info
            )

        checks_section = "\n\n**Requirements (will be tested):**\n" + "\n".join(
            f"- {check}" for check in task.checks
        )

        prompt = f"""You are an expert web developer. Create a complete, production-ready single-page web application based on the following requirements.

**Task:** {task.task}

**Brief:** {task.brief}{attachments_section}{checks_section}

**Instructions:**
1. Create a minimal, functional web application that meets ALL requirements
2. Use only vanilla HTML, CSS, and JavaScript (no build tools required)
3. Include all necessary CDN links for external libraries (Bootstrap, marked, highlight.js, etc.)
4. Ensure the app is self-contained in a single index.html file or minimal files
5. Follow best practices for code quality, accessibility, and user experience
6. Include helpful comments explaining key functionality
7. Make the UI clean and professional using Bootstrap 5 or similar

**Output Format:**
Provide the complete code for each file in JSON format:
```json
{{
  "index.html": "<!DOCTYPE html>...",
  "style.css": "/* optional styles */",
  "script.js": "// optional separate JS",
  "README.md": "# Project Title\\n\\n..."
}}
```

Generate ONLY the JSON output, no other text. Ensure all code is complete and functional.
"""
        return prompt

    def _generate_with_llm(self, prompt: str) -> dict[str, str]:
        """Call LLM API to generate code.

        Args:
            prompt: Generation prompt

        Returns:
            Dictionary of filename -> content
        """
        logger.info(f"Calling {self.provider} API for code generation")

        try:
            if self.provider == "anthropic":
                response = self.client.messages.create(
                    model=self.model,
                    max_tokens=4096,
                    temperature=0.3,
                    messages=[{"role": "user", "content": prompt}],
                )
                content = response.content[0].text

            elif self.provider in ["openai", "aipipe"]:
                # Both OpenAI and AIPipe use the same API format
                response = self.client.chat.completions.create(
                    model=self.model,
                    messages=[{"role": "user", "content": prompt}],
                    temperature=0.3,
                    max_tokens=4096,
                )
                content = response.choices[0].message.content

            else:
                raise ValueError(f"Unsupported provider: {self.provider}")

            # Extract JSON from response
            files = self._extract_json(content)
            return files

        except Exception as e:
            logger.error(f"LLM generation failed: {e}")
            # Fallback to minimal template
            return self._get_fallback_template()

    def _extract_json(self, content: str) -> dict[str, str]:
        """Extract JSON from LLM response.

        Args:
            content: LLM response text

        Returns:
            Parsed JSON dictionary
        """
        # Try to find JSON in markdown code block (greedy match for nested braces)
        import re

        json_match = re.search(r"```json\s*(\{.*\})\s*```", content, re.DOTALL)
        if json_match:
            result = json.loads(json_match.group(1))
            # Validate that all values are strings
            return {k: v for k, v in result.items() if v is not None}

        # Try to find JSON in plain code block
        json_match = re.search(r"```\s*(\{.*\})\s*```", content, re.DOTALL)
        if json_match:
            result = json.loads(json_match.group(1))
            # Validate that all values are strings
            return {k: v for k, v in result.items() if v is not None}

        # Try to parse the whole content as JSON
        try:
            result = json.loads(content)
            # Validate that all values are strings
            return {k: v for k, v in result.items() if v is not None}
        except json.JSONDecodeError:
            # Try to find any JSON object (greedy match)
            json_match = re.search(r"\{.*\}", content, re.DOTALL)
            if json_match:
                result = json.loads(json_match.group(0))
                # Validate that all values are strings
                return {k: v for k, v in result.items() if v is not None}

        logger.error(f"Could not extract JSON from LLM response: {content[:200]}")
        raise ValueError("Could not extract JSON from LLM response")

    def _get_fallback_template(self) -> dict[str, str]:
        """Get fallback template when LLM generation fails.

        Returns:
            Basic HTML template
        """
        return {
            "index.html": """<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Generated App</title>
    <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0/dist/css/bootstrap.min.css" rel="stylesheet">
</head>
<body>
    <div class="container mt-5">
        <h1>Application</h1>
        <p>This is a minimal fallback template.</p>
    </div>
    <script src="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0/dist/js/bootstrap.bundle.min.js"></script>
</body>
</html>""",
            "README.md": """# Generated Application

This is an automatically generated web application.

## Setup

Simply open `index.html` in a web browser.

## License

MIT License
""",
        }

    def _save_files(self, files: dict[str, str], output_dir: Path) -> None:
        """Save generated files to output directory.

        Args:
            files: Dictionary of filename -> content
            output_dir: Directory to save files
        """
        output_dir.mkdir(parents=True, exist_ok=True)

        for filename, content in files.items():
            # Skip None or empty content
            if content is None:
                logger.warning(f"Skipping {filename} - content is None")
                continue
            
            # Ensure content is a string
            if not isinstance(content, str):
                logger.warning(f"Converting {filename} content to string")
                content = str(content)
            
            file_path = output_dir / filename
            file_path.write_text(content, encoding="utf-8")
            logger.debug(f"Saved {filename} ({len(content)} bytes)")