src/nodes/system_design/lld_code_generator.py

import os
import logging
from typing import Dict, List, Optional
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field
from langchain_core.prompts import ChatPromptTemplate

# Setup Logger
logger = logging.getLogger(__name__)

# --- Pydantic Models for Structured Output ---
class CodeFile(BaseModel):
    filename: str = Field(description="Name of the file, e.g., 'NotificationFactory.java'")
    content: str = Field(description="Complete source code for the file")
    
class LLDCodeResponse(BaseModel):
    title: str = Field(description="Title of the design pattern being demonstrated")
    subtitle: str = Field(description="Subtitle explaining the specific implementation")
    files: List[CodeFile] = Field(description="List of code files. Max 3-4 files.")
    execution_output: str = Field(description="Console output for the Main execution file.")
    typing_speed: int = Field(default=25, description="Typing animation speed in ms per char")
    auto_run: bool = Field(default=True, description="Whether to auto-run the animation")

# Pattern-specific guidance
PATTERN_GUIDANCE = {
    "singleton": {
        "description": "Ensures a class has only one instance and provides global access",
        "files_needed": ["Singleton class", "Main/Demo class"],
        "key_concepts": ["Private constructor", "Static instance", "Thread safety"],
        "example_domain": "Database connection, Logger, Configuration manager"
    },
    "factory": {
        "description": "Creates objects without specifying exact class to create",
        "files_needed": ["Product interface", "Concrete products", "Factory class", "Main/Demo"],
        "key_concepts": ["Product interface", "Concrete implementations", "Factory method"],
        "example_domain": "Notification system (Email, SMS, Push), Vehicle factory, Payment processors"
    },
    "observer": {
        "description": "Defines one-to-many dependency between objects",
        "files_needed": ["Subject interface", "Observer interface", "Concrete subject", "Concrete observers", "Main/Demo"],
        "key_concepts": ["Subject", "Observer", "Notify mechanism", "Subscription"],
        "example_domain": "Event system, Stock price updates, News feed, Weather station"
    },
    "strategy": {
        "description": "Defines family of algorithms, encapsulates each, makes them interchangeable",
        "files_needed": ["Strategy interface", "Concrete strategies", "Context class", "Main/Demo"],
        "key_concepts": ["Strategy interface", "Concrete strategies", "Context switching"],
        "example_domain": "Payment methods, Sorting algorithms, Compression methods, Route calculation"
    },
    "builder": {
        "description": "Separates construction of complex object from representation",
        "files_needed": ["Product class", "Builder interface/class", "Concrete builder", "Director (optional)", "Main/Demo"],
        "key_concepts": ["Step-by-step construction", "Fluent interface", "Immutable product"],
        "example_domain": "Building complex objects (House, Pizza, Document), Query builder"
    },
    "decorator": {
        "description": "Attaches additional responsibilities to object dynamically",
        "files_needed": ["Component interface", "Concrete component", "Decorator base", "Concrete decorators", "Main/Demo"],
        "key_concepts": ["Component interface", "Wrapping", "Composition over inheritance"],
        "example_domain": "Coffee with addons, Text formatting, Stream wrappers"
    },
    "adapter": {
        "description": "Converts interface of class into another interface clients expect",
        "files_needed": ["Target interface", "Adaptee class", "Adapter class", "Main/Demo"],
        "key_concepts": ["Target interface", "Adaptee", "Adapter wrapping"],
        "example_domain": "Legacy system integration, Third-party API wrapper"
    },
    "command": {
        "description": "Encapsulates request as object, allowing parameterization",
        "files_needed": ["Command interface", "Concrete commands", "Receiver", "Invoker", "Main/Demo"],
        "key_concepts": ["Command interface", "Execute method", "Receiver", "Undo capability"],
        "example_domain": "Text editor operations, Remote control, Transaction system"
    },
    "design_pattern": {
        "description": "Generic design pattern implementation",
        "files_needed": ["Interface", "Implementation", "Main/Demo"],
        "key_concepts": ["Clean code", "SOLID principles", "Best practices"],
        "example_domain": "Based on topic context"
    }
}

class LLDCodeGenerator:
    """
    Generator for Low-Level Design code examples using LangChain for structured multi-file output.
    """
    
    def __init__(self, client=None):
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key and client and hasattr(client, "api_key"):
             api_key = client.api_key
             
        self.llm = ChatOpenAI(
            model="gpt-4o",
            temperature=0.2,
            api_key=api_key
        )
        self.logger = logging.getLogger(__name__)
    
    def generate_lld_code_content(
        self, 
        slide_config: Dict, 
        topic: str, 
        architecture_description: str
    ) -> Dict:
        """
        Generate multi-file LLD code examples with pattern-specific guidance.
        """
        focus = slide_config.get('focus', 'Low-Level Design')
        language = slide_config.get('language', 'java')
        lld_pattern = slide_config.get('lld_pattern', 'design_pattern')
        
        self.logger.info(f"Generating LLD code for: {focus} in {language} using pattern: {lld_pattern}")
        
        # Get pattern-specific guidance
        pattern_info = PATTERN_GUIDANCE.get(lld_pattern.lower(), PATTERN_GUIDANCE['design_pattern'])
        
        # Define input model for type safety
        class PromptInput(BaseModel):
            language: str
            topic: str
            focus: str
            lld_pattern: str
            pattern_description: str
            files_needed: str
            key_concepts: str
            example_domain: str
            context: str
        
        # Prepare Prompt
        system_prompt = """You are an expert software architect creating clean, educational code examples for Low-Level Design patterns.

Generate complete, working multi-file code that demonstrates design patterns and SOLID principles.

CRITICAL REQUIREMENTS:

1. **Pattern Implementation**:
   - Implement the EXACT pattern specified: {lld_pattern}
   - Pattern description: {pattern_description}
   - Required files: {files_needed}
   - Key concepts to demonstrate: {key_concepts}
   - Use domain context: {example_domain}

2. **Multiple Files** (3-4 files):
   - Interface/Abstract class (if applicable)
   - Concrete implementation classes (2-3 variants)
   - Main/Demo class showing usage
   - Each file should be focused and follow Single Responsibility Principle

3. **REAL-WORLD NAMING**:
   - **NEVER** use generic names like `ConcreteProductA`, `Component1`, `Observer1`, `StrategyA`
   - **ALWAYS** use domain-specific names based on the TOPIC ({topic}) and FOCUS ({focus})
   - Example for Notification System + Factory pattern:
     ✓ GOOD: `NotificationFactory`, `EmailNotification`, `SMSNotification`, `PushNotification`
     ✗ BAD: `Factory`, `ProductA`, `ProductB`, `ConcreteProduct1`
   - Example for Payment System + Strategy pattern:
     ✓ GOOD: `PaymentStrategy`, `CreditCardPayment`, `PayPalPayment`, `CryptoPayment`
     ✗ BAD: `Strategy`, `ConcreteStrategyA`, `ConcreteStrategyB`

4. **Clean Code**:
   - Use proper naming conventions for {language}
   - Add helpful comments explaining pattern-specific parts
   - COMPACT formatting: Use MINIMAL blank lines (max 1 blank line between methods/classes)
   - NO excessive spacing or empty lines at start of files
   - Keep code visually compact and professional

5. **Pattern-Specific Implementation**:
   - For **Factory**: Show polymorphic object creation without exposing instantiation logic
   - For **Singleton**: Include thread-safe implementation, private constructor, static instance
   - For **Observer**: Demonstrate subject-observer relationship, notify mechanism
   - For **Strategy**: Show algorithm family, context switching between strategies
   - For **Builder**: Implement step-by-step construction, fluent interface
   - For **Decorator**: Show component wrapping, responsibility addition
   - For **Adapter**: Demonstrate interface conversion between incompatible interfaces
   - For **Command**: Encapsulate requests as objects, show execute/undo

6. **Execution Output**:
   - Provide realistic console output showing the pattern in action
   - Output should demonstrate key pattern benefits
   - Show object creation, method calls, state changes

7. **Language**: {language}

OUTPUT STRUCTURE:
Return a JSON object with:
- title: Pattern name + brief description
- subtitle: Specific implementation focus
- files: list of objects {{filename, content}}
- execution_output: string containing the console output
- typing_speed: 25-35 (ms per character)
- auto_run: true
"""

        user_prompt = """
TOPIC: {topic}
FOCUS: {focus}
PATTERN: {lld_pattern}
PATTERN DESCRIPTION: {pattern_description}
KEY CONCEPTS: {key_concepts}
EXAMPLE DOMAIN: {example_domain}
CONTEXT: {context}

Generate the complete pattern implementation now with proper domain-specific names.
"""
        
        try:
            structured_llm = self.llm.with_structured_output(LLDCodeResponse)
            prompt = ChatPromptTemplate.from_messages([
                ("system", system_prompt),
                ("user", user_prompt)
            ])
            
            chain = prompt | structured_llm
            
            # Create input with Pydantic model
            prompt_input = PromptInput(
                language=language,
                topic=topic,
                focus=focus,
                lld_pattern=lld_pattern,
                pattern_description=pattern_info['description'],
                files_needed=", ".join(pattern_info['files_needed']),
                key_concepts=", ".join(pattern_info['key_concepts']),
                example_domain=pattern_info['example_domain'],
                context=architecture_description[:800]
            )
            
            result: LLDCodeResponse = chain.invoke(prompt_input.model_dump())
            
            # Convert List[CodeFile] back to Dict[str, str] for frontend compatibility
            code_files_dict = {f.filename: f.content for f in result.files}
            
            # Construct execution output dict (Main file gets the output, others empty)
            exec_output_dict = {name: "" for name in code_files_dict}
            
            # Heuristic: Find Main file or use the last one
            main_file = next(
                (f for f in code_files_dict if "Main" in f or "Demo" in f or "App" in f or "Test" in f), 
                list(code_files_dict.keys())[-1]
            )
            if main_file:
                exec_output_dict[main_file] = result.execution_output
            
            self.logger.info(f"Generated {len(code_files_dict)} files for {lld_pattern} pattern")
            
            return {
                "title": result.title,
                "subtitle": result.subtitle,
                "code_files": code_files_dict,
                "execution_output": exec_output_dict,
                "typing_speed": result.typing_speed,
                "auto_run": result.auto_run,
                "pattern": lld_pattern
            }
            
        except Exception as e:
            self.logger.error(f"Error generating LLD code with LangChain: {e}")
            return self._get_fallback_lld_code(focus, language, lld_pattern)

    def _get_fallback_lld_code(self, focus: str, language: str, pattern: str) -> Dict:
        """Fallback LLD code if generation fails - pattern-specific"""
        self.logger.warning(f"Using fallback LLD code for {pattern}")
        
        if pattern.lower() == "singleton":
            return {
                "title": "Singleton Pattern",
                "subtitle": "Thread-Safe Database Connection",
                "code_files": {
                    "DatabaseConnection.java": """public class DatabaseConnection {
    private static volatile DatabaseConnection instance;
    private String connectionString;
    
    private DatabaseConnection() {
        // Private constructor
        this.connectionString = "jdbc:mysql://localhost:3306/mydb";
        System.out.println("Database connection initialized");
    }
    
    public static DatabaseConnection getInstance() {
        if (instance == null) {
            synchronized (DatabaseConnection.class) {
                if (instance == null) {
                    instance = new DatabaseConnection();
                }
            }
        }
        return instance;
    }
    
    public void query(String sql) {
        System.out.println("Executing: " + sql);
    }
}""",
                    "Main.java": """public class Main {
    public static void main(String[] args) {
        // Only one instance created
        DatabaseConnection db1 = DatabaseConnection.getInstance();
        DatabaseConnection db2 = DatabaseConnection.getInstance();
        
        System.out.println("Same instance? " + (db1 == db2));
        
        db1.query("SELECT * FROM users");
    }
}"""
                },
                "execution_output": {
                    "Main.java": """Database connection initialized
Same instance? true
Executing: SELECT * FROM users"""
                },
                "typing_speed": 25,
                "auto_run": True,
                "pattern": "singleton"
            }
        
        # Generic fallback
        return {
            "title": focus,
            "subtitle": "Fallback Example",
            "code_files": {
                "Main.java": f"""public class Main {{
    public static void main(String[] args) {{
        System.out.println("Error generating {pattern} pattern code.");
        System.out.println("Focus: {focus}");
    }}
}}"""
            },
            "execution_output": {
                "Main.java": f"Error generating {pattern} pattern code.\nFocus: {focus}"
            },
            "typing_speed": 25,
            "auto_run": True,
            "pattern": pattern
        }