storage.py

import json
import os
from pathlib import Path
from typing import Dict, List, Optional
from datetime import datetime


class StorageManager:
    """Manage run history and artifacts using JSON files."""
    
    def __init__(self, base_dir: str = "runs"):
        self.base_dir = Path(base_dir)
        self.base_dir.mkdir(exist_ok=True)
        self.index_file = self.base_dir / "index.json"
        
        # Initialize index if it doesn't exist
        if not self.index_file.exists():
            self._save_index([])
    
    def save_run(self, run_id: str, run_data: Dict):
        """Save run data to disk."""
        run_dir = self.base_dir / run_id
        run_dir.mkdir(exist_ok=True)
        
        # Save main run data
        run_file = run_dir / "run.json"
        with open(run_file, 'w') as f:
            json.dump(run_data, f, indent=2)
        
        # Save individual artifacts
        if 'diagnosis' in run_data:
            with open(run_dir / "diagnosis.json", 'w') as f:
                json.dump(run_data['diagnosis'], f, indent=2)
        
        if 'plan' in run_data:
            with open(run_dir / "plan.json", 'w') as f:
                json.dump(run_data['plan'], f, indent=2)
        
        if 'patches' in run_data:
            with open(run_dir / "patches.json", 'w') as f:
                json.dump(run_data['patches'], f, indent=2)
        
        if 'validation_results' in run_data:
            with open(run_dir / "validation.json", 'w') as f:
                json.dump(run_data['validation_results'], f, indent=2)
        
        # Update index
        self._update_index(run_id, run_data)
    
    def load_run(self, run_id: str) -> Optional[Dict]:
        """Load run data from disk."""
        run_file = self.base_dir / run_id / "run.json"
        
        if not run_file.exists():
            return None
        
        with open(run_file, 'r') as f:
            return json.load(f)
    
    def get_recent_runs(self, limit: int = 10) -> List[Dict]:
        """Get recent runs from index."""
        index = self._load_index()
        
        # Sort by timestamp (most recent first)
        index.sort(key=lambda x: x.get('timestamp', ''), reverse=True)
        
        return index[:limit]
    
    def _update_index(self, run_id: str, run_data: Dict):
        """Update the index with a new run."""
        index = self._load_index()
        
        # Remove existing entry if present
        index = [entry for entry in index if entry.get('run_id') != run_id]
        
        # Add new entry
        index_entry = {
            'run_id': run_id,
            'timestamp': run_data.get('timestamp'),
            'repo_url': run_data.get('repo_url'),
            'status': run_data.get('status'),
            'pr_url': run_data.get('pr_url'),
            'issue_summary': run_data.get('diagnosis', {}).get('summary', 'No summary')
        }
        index.append(index_entry)
        
        self._save_index(index)
    
    def _load_index(self) -> List[Dict]:
        """Load the index file."""
        if not self.index_file.exists():
            return []
        
        try:
            with open(self.index_file, 'r') as f:
                return json.load(f)
        except:
            return []
    
    def _save_index(self, index: List[Dict]):
        """Save the index file."""
        with open(self.index_file, 'w') as f:
            json.dump(index, f, indent=2)
    
    def save_pr_ranking(self, repo_name: str, ranking_data: Dict):
        """Save PR ranking results."""
        rankings_dir = self.base_dir / "rankings"
        rankings_dir.mkdir(exist_ok=True)
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        safe_repo_name = repo_name.replace("/", "_")
        ranking_file = rankings_dir / f"{safe_repo_name}_{timestamp}.json"
        
        with open(ranking_file, 'w') as f:
            json.dump(ranking_data, f, indent=2)
    
    def get_latest_ranking(self, repo_name: str) -> Optional[Dict]:
        """Get the latest PR ranking for a repository."""
        rankings_dir = self.base_dir / "rankings"
        
        if not rankings_dir.exists():
            return None
        
        safe_repo_name = repo_name.replace("/", "_")
        ranking_files = sorted(rankings_dir.glob(f"{safe_repo_name}_*.json"), reverse=True)
        
        if not ranking_files:
            return None
        
        with open(ranking_files[0], 'r') as f:
            return json.load(f)