file_manager.py

import os
from datetime import datetime
from typing import Dict
import subprocess

class FileManager:
    def __init__(self):
        self.output_base = "output"
        
    def create_output_dir(self, project_name: str, timestamp: str) -> str:
        """Create output directory structure for this run."""
        dir_name = f"{project_name}_{timestamp}"
        output_dir = os.path.join(self.output_base, dir_name)
        
        # Create main directories
        os.makedirs(output_dir, exist_ok=True)
        os.makedirs(os.path.join(output_dir, "code"), exist_ok=True)
        os.makedirs(os.path.join(output_dir, "analysis"), exist_ok=True)
        os.makedirs(os.path.join(output_dir, "images"), exist_ok=True)
        
        return output_dir
    
    def save_step_result(self, output_dir: str, step_name: str, content: str):
        """Save intermediate step result with clean PlantUML code."""
        ext = ".pu" if "diagram" in step_name else ".md"
        
        # Determine target directory based on file type
        if ext == ".pu":
            target_dir = os.path.join(output_dir, "code")
            # Render PlantUML diagram
            self._render_plantuml(content, output_dir, step_name)
        else:
            target_dir = os.path.join(output_dir, "analysis")
            
        filename = f"{step_name}{ext}"
        filepath = os.path.join(target_dir, filename)
        
        print(f"Saving {filename} to {target_dir}")  # Add this line for debugging
        
        # Clean PlantUML code by removing markdown code block markers
        if ext == ".pu":
            content = self._clean_plantuml_content(content)
        
        with open(filepath, "w") as f:
            f.write(content)

    def _render_plantuml(self, content: str, output_dir: str, step_name: str):
        """Render PlantUML diagram to PNG."""
        # Clean the content first
        content = self._clean_plantuml_content(content)
        
        # Create temporary .pu file
        temp_pu = os.path.join(output_dir, "code", f"{step_name}.pu")
        with open(temp_pu, "w") as f:
            f.write(content)
        
        # Render to PNG
        output_png = os.path.join(output_dir, "images", f"{step_name}.png")
        try:
            subprocess.run([
                "plantuml", 
                "-tpng",
                temp_pu,
                "-o",  "images"
            ], check=True)
        except subprocess.CalledProcessError as e:
            print(f"Error rendering PlantUML diagram: {e}")
        except FileNotFoundError:
            print("PlantUML not found. Please install PlantUML to generate images.")

    def _clean_plantuml_content(self, content: str) -> str:
        """Remove triple quotes and plantuml markers from the content."""
        # Remove '''plantuml and ''' markers
        content = content.replace("'''plantuml", "").replace("'''", "")
        content = content.replace("```plantuml", "").replace("```", "")
        
        
        # Remove any remaining leading/trailing whitespace
        content = content.strip()
        
        return content
    
    def save_final_results(self, state: Dict):
        """Save all final results."""
        output_dir = state["output_dir"]
        
        # Save all analysis results and diagram codes
        for key, value in state.items():
            # Skip metadata fields
            if key in ["project_name", "project_description", "output_dir", "selected_diagrams"]:
                continue
            
            if value:  # Only save if there's content
                self.save_step_result(output_dir, key, value)
        
        # Save metadata
        metadata = {
            "project_name": state["project_name"],
            "project_description": state["project_description"],
            "generated_diagrams": state["selected_diagrams"],
            "timestamp": datetime.now().isoformat()
        }
        
        import json
        with open(os.path.join(output_dir, "metadata.json"), "w") as f:
            json.dump(metadata, f, indent=2)