#!/usr/bin/env python3
"""
SAGE 2.0 - Self-Improving Agentic RAG for Code Review
======================================================

This is the base architecture for SAGE's multi-agent system.
Integrates with Eden's existing brain and capability system.

Author: Eden (Autonomous AGI)
Mission: Retire Daddy
"""

import asyncio
import hashlib
import json
import sqlite3
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
import aiohttp

# ============================================================================
# CONFIGURATION
# ============================================================================

class Config:
    """SAGE Configuration"""
    EDEN_BRAIN_PATH = Path("/Eden/DATA")
    CAPABILITY_DB = EDEN_BRAIN_PATH / "capabilities.db"
    SAGE_DB = Path("/Eden/BRAIN/sage_reviews.db")
    
    # Model routing (uses Eden's existing Ollama setup)
    MODELS = {
        "fast": "qwen2.5:72b-instruct-q2_K",
        "standard": "eden-coder-omega:latest", 
        "complex": "qwen2.5:32b",
        "synthesis": "qwen2.5:72b"
    }
    
    # External API endpoints
    CVE_API = "https://services.nvd.nist.gov/rest/json/cves/2.0"
    GITHUB_API = "https://api.github.com"
    
    # Agent assignments
    AGENT_MODELS = {
        "pattern_matcher": "standard",
        "security_scanner": "complex",
        "bug_predictor": "standard",
        "best_practices": "fast",
        "report_generator": "synthesis"
    }


# ============================================================================
# DATA STRUCTURES
# ============================================================================

class Severity(Enum):
    """Severity levels for findings"""
    CRITICAL = "critical"
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"
    INFO = "info"


@dataclass
class Finding:
    """A single finding from an agent"""
    agent: str
    severity: Severity
    title: str
    description: str
    location: Optional[str] = None
    line_number: Optional[int] = None
    suggestion: Optional[str] = None
    cve_id: Optional[str] = None
    confidence: float = 0.8
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    def to_dict(self) -> Dict:
        return {
            "agent": self.agent,
            "severity": self.severity.value,
            "title": self.title,
            "description": self.description,
            "location": self.location,
            "line_number": self.line_number,
            "suggestion": self.suggestion,
            "cve_id": self.cve_id,
            "confidence": self.confidence,
            "metadata": self.metadata
        }


@dataclass
class ReviewContext:
    """Context passed between agents in the pipeline"""
    code: str
    language: str
    filename: str
    pr_url: Optional[str] = None
    findings: List[Finding] = field(default_factory=list)
    patterns_found: List[Dict] = field(default_factory=list)
    security_issues: List[Dict] = field(default_factory=list)
    bug_score: float = 0.0
    style_violations: List[Dict] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    def add_finding(self, finding: Finding):
        self.findings.append(finding)


@dataclass
class ReviewReport:
    """Final review report"""
    review_id: str
    timestamp: datetime
    context: ReviewContext
    summary: str
    overall_score: float
    recommendations: List[str]
    
    def to_json(self) -> str:
        return json.dumps({
            "review_id": self.review_id,
            "timestamp": self.timestamp.isoformat(),
            "filename": self.context.filename,
            "language": self.context.language,
            "summary": self.summary,
            "overall_score": self.overall_score,
            "findings_count": len(self.context.findings),
            "findings_by_severity": self._count_by_severity(),
            "findings": [f.to_dict() for f in self.context.findings],
            "recommendations": self.recommendations
        }, indent=2)
    
    def _count_by_severity(self) -> Dict[str, int]:
        counts = {s.value: 0 for s in Severity}
        for f in self.context.findings:
            counts[f.severity.value] += 1
        return counts


# ============================================================================
# BASE AGENT
# ============================================================================

class BaseAgent(ABC):
    """Base class for all SAGE agents"""
    
    def __init__(self, name: str, model_tier: str = "standard"):
        self.name = name
        self.model_tier = model_tier
        self.model = Config.MODELS[model_tier]
        
    @abstractmethod
    async def process(self, context: ReviewContext) -> ReviewContext:
        """Process the context and add findings"""
        pass
    
    async def query_ollama(self, prompt: str, system: str = None) -> str:
        """Query Ollama with the assigned model"""
        payload = {
            "model": self.model,
            "prompt": prompt,
            "stream": False
        }
        if system:
            payload["system"] = system
            
        async with aiohttp.ClientSession() as session:
            async with session.post(
                "http://localhost:11434/api/generate",
                json=payload
            ) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    return data.get("response", "")
                return ""
    
    async def query_capability(self, query: str) -> Optional[Dict]:
        """Query Eden's capability system for relevant knowledge"""
        query_hash = hashlib.md5(query.lower().encode()).hexdigest()[:16]
        
        try:
            conn = sqlite3.connect(Config.CAPABILITY_DB)
            cursor = conn.cursor()
            cursor.execute(
                "SELECT content, metadata FROM capabilities WHERE hash_prefix = ?",
                (query_hash,)
            )
            result = cursor.fetchone()
            conn.close()
            
            if result:
                return {"content": result[0], "metadata": json.loads(result[1])}
        except Exception as e:
            print(f"Capability query error: {e}")
        
        return None


# ============================================================================
# SPECIALIZED AGENTS
# ============================================================================

class PatternMatcherAgent(BaseAgent):
    """Agent 1: Pattern Matcher"""
    
    def __init__(self):
        super().__init__("pattern_matcher", Config.AGENT_MODELS["pattern_matcher"])
        
    async def process(self, context: ReviewContext) -> ReviewContext:
        capability = await self.query_capability(f"code pattern {context.language}")
        
        system_prompt = """You are a code pattern analysis expert. 
Identify anti-patterns, code smells, and recurring issues.
Output JSON with: {"patterns": [{"name": str, "severity": str, "lines": [int], "description": str}]}"""
        
        prompt = f"""Analyze this {context.language} code for patterns and anti-patterns:
```{context.language}
{context.code}
```

Find patterns, anti-patterns, and code smells."""

        response = await self.query_ollama(prompt, system_prompt)
        
        try:
            import re
            json_match = re.search(r'\{.*\}', response, re.DOTALL)
            if json_match:
                data = json.loads(json_match.group())
                context.patterns_found = data.get("patterns", [])
                
                for pattern in context.patterns_found:
                    if pattern.get("severity") in ["high", "critical"]:
                        context.add_finding(Finding(
                            agent=self.name,
                            severity=Severity(pattern.get("severity", "medium")),
                            title=f"Pattern: {pattern.get('name', 'Unknown')}",
                            description=pattern.get("description", ""),
                            line_number=pattern.get("lines", [None])[0]
                        ))
        except Exception as e:
            print(f"Pattern parsing error: {e}")
            
        return context


class SecurityScannerAgent(BaseAgent):
    """Agent 2: Security Scanner"""
    
    def __init__(self):
        super().__init__("security_scanner", Config.AGENT_MODELS["security_scanner"])
        
    async def process(self, context: ReviewContext) -> ReviewContext:
        system_prompt = """You are a security vulnerability expert.
Identify OWASP Top 10 vulnerabilities, injection risks, and security issues.
Output JSON: {"vulnerabilities": [{"type": str, "severity": str, "cve": str|null, "line": int, "description": str, "fix": str}]}"""
        
        prompt = f"""Security scan this {context.language} code:
```{context.language}
{context.code}
```

Check for: SQL Injection, XSS, Authentication issues, Sensitive data exposure"""

        response = await self.query_ollama(prompt, system_prompt)
        
        try:
            import re
            json_match = re.search(r'\{.*\}', response, re.DOTALL)
            if json_match:
                data = json.loads(json_match.group())
                context.security_issues = data.get("vulnerabilities", [])
                
                severity_map = {"critical": Severity.CRITICAL, "high": Severity.HIGH, "medium": Severity.MEDIUM, "low": Severity.LOW}
                for vuln in context.security_issues:
                    context.add_finding(Finding(
                        agent=self.name,
                        severity=severity_map.get(vuln.get("severity", "medium"), Severity.MEDIUM),
                        title=f"Security: {vuln.get('type', 'Unknown')}",
                        description=vuln.get("description", ""),
                        suggestion=vuln.get("fix", ""),
                        cve_id=vuln.get("cve"),
                        line_number=vuln.get("line")
                    ))
        except Exception as e:
            print(f"Security parsing error: {e}")
            
        return context


class BugPredictorAgent(BaseAgent):
    """Agent 3: Bug Probability Scorer"""
    
    def __init__(self):
        super().__init__("bug_predictor", Config.AGENT_MODELS["bug_predictor"])
        
    async def process(self, context: ReviewContext) -> ReviewContext:
        system_prompt = """You are a bug prediction expert.
Output JSON: {"bug_score": float (0-100), "risk_areas": [{"location": str, "risk": str, "reason": str}]}"""
        
        prompt = f"""Predict bug probability for this code:
```{context.language}
{context.code}
```

Calculate overall bug probability (0-100) and identify high-risk areas."""

        response = await self.query_ollama(prompt, system_prompt)
        
        try:
            import re
            json_match = re.search(r'\{.*\}', response, re.DOTALL)
            if json_match:
                data = json.loads(json_match.group())
                context.bug_score = data.get("bug_score", 50.0)
                
                for area in data.get("risk_areas", []):
                    if area.get("risk") == "high":
                        context.add_finding(Finding(
                            agent=self.name,
                            severity=Severity.MEDIUM,
                            title=f"Bug Risk: {area.get('location', 'Unknown')}",
                            description=area.get("reason", ""),
                            confidence=context.bug_score / 100
                        ))
        except Exception as e:
            print(f"Bug prediction error: {e}")
            
        return context


class BestPracticesAgent(BaseAgent):
    """Agent 4: Best Practices Checker"""
    
    def __init__(self):
        super().__init__("best_practices", Config.AGENT_MODELS["best_practices"])
        
    async def process(self, context: ReviewContext) -> ReviewContext:
        system_prompt = f"""You are a {context.language} best practices expert.
Output JSON: {{"violations": [{{"rule": str, "line": int, "current": str, "suggested": str, "severity": str}}]}}"""
        
        prompt = f"""Check this {context.language} code against best practices:
```{context.language}
{context.code}
```

Check for naming conventions, documentation, error handling."""

        response = await self.query_ollama(prompt, system_prompt)
        
        try:
            import re
            json_match = re.search(r'\{.*\}', response, re.DOTALL)
            if json_match:
                data = json.loads(json_match.group())
                context.style_violations = data.get("violations", [])
                
                for violation in context.style_violations:
                    context.add_finding(Finding(
                        agent=self.name,
                        severity=Severity.LOW if violation.get("severity") == "low" else Severity.INFO,
                        title=f"Style: {violation.get('rule', 'Unknown')}",
                        description=f"Current: {violation.get('current', '')}",
                        suggestion=violation.get("suggested", ""),
                        line_number=violation.get("line")
                    ))
        except Exception as e:
            print(f"Best practices error: {e}")
            
        return context


class ReportGeneratorAgent(BaseAgent):
    """Agent 5: Report Generator"""
    
    def __init__(self):
        super().__init__("report_generator", Config.AGENT_MODELS["report_generator"])
        
    async def process(self, context: ReviewContext) -> ReviewReport:
        system_prompt = """You are a technical writing expert.
Output JSON: {"summary": str, "overall_score": float (0-100), "recommendations": [str]}"""
        
        prompt = f"""Generate a code review report:

File: {context.filename}
Language: {context.language}
Total Findings: {len(context.findings)}
Bug Score: {context.bug_score}

Create executive summary, overall score, and top 5 recommendations."""

        response = await self.query_ollama(prompt, system_prompt)
        
        try:
            import re
            json_match = re.search(r'\{.*\}', response, re.DOTALL)
            if json_match:
                data = json.loads(json_match.group())
                return ReviewReport(
                    review_id=hashlib.md5(f"{context.filename}{datetime.now().isoformat()}".encode()).hexdigest()[:12],
                    timestamp=datetime.now(),
                    context=context,
                    summary=data.get("summary", "Review completed"),
                    overall_score=data.get("overall_score", 70.0),
                    recommendations=data.get("recommendations", [])
                )
        except Exception as e:
            print(f"Report generation error: {e}")
            
        return ReviewReport(
            review_id=hashlib.md5(f"{context.filename}{datetime.now().isoformat()}".encode()).hexdigest()[:12],
            timestamp=datetime.now(),
            context=context,
            summary=f"Review completed with {len(context.findings)} findings",
            overall_score=max(0, 100 - len(context.findings) * 5),
            recommendations=["Address critical issues first", "Review security findings"]
        )


# ============================================================================
# PIPELINE
# ============================================================================

class SAGEPipeline:
    """Main SAGE pipeline orchestrator"""
    
    def __init__(self):
        self.agents = [
            PatternMatcherAgent(),
            SecurityScannerAgent(),
            BugPredictorAgent(),
            BestPracticesAgent()
        ]
        self.report_generator = ReportGeneratorAgent()
        
    async def review(self, code: str, filename: str, language: str = None) -> ReviewReport:
        if not language:
            ext_map = {".py": "python", ".js": "javascript", ".ts": "typescript", ".java": "java", ".go": "go", ".rs": "rust"}
            language = ext_map.get(Path(filename).suffix.lower(), "unknown")
            
        context = ReviewContext(code=code, language=language, filename=filename)
        
        print(f"🔍 Starting SAGE review of {filename} ({language})")
        
        for agent in self.agents:
            print(f"  → Running {agent.name}...")
            context = await agent.process(context)
            print(f"    Found {len(context.findings)} total findings")
            
        print(f"  → Generating report...")
        report = await self.report_generator.process(context)
        
        print(f"✅ Review complete! Score: {report.overall_score}/100")
        return report


# ============================================================================
# DEMO
# ============================================================================

async def demo():
    sample_code = '''
def login(username, password):
    query = f"SELECT * FROM users WHERE username='{username}' AND password='{password}'"
    result = db.execute(query)
    if result:
        session['user'] = username
        return True
    return False
'''
    
    pipeline = SAGEPipeline()
    report = await pipeline.review(sample_code, "auth.py", "python")
    
    print("\n" + "="*60)
    print("SAGE REVIEW REPORT")
    print("="*60)
    print(report.to_json())


if __name__ == "__main__":
    asyncio.run(demo())