"""
RAG-based Incident Learning System
Captures deployment failures and network incidents for root cause analysis
Generates regression tests and updates LLM knowledge to prevent recurrence
"""
import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional
from dataclasses import dataclass, asdict

logger = logging.getLogger(__name__)


@dataclass
class Incident:
    """Network incident or deployment failure"""
    id: str
    timestamp: str
    severity: str  # critical, high, medium, low
    category: str  # deployment_failure, config_error, routing_issue, outage
    description: str
    root_cause: Optional[str] = None
    affected_devices: List[str] = None
    
    # Context for RAG
    network_model: Optional[Dict[str, Any]] = None
    config_changes: Optional[List[Dict[str, str]]] = None
    validation_errors: Optional[List[Dict[str, Any]]] = None
    
    # Resolution
    resolution: Optional[str] = None
    resolved_at: Optional[str] = None
    
    # Learning outcomes
    regression_test: Optional[str] = None
    llm_prompt_update: Optional[str] = None
    
    def __post_init__(self):
        if self.affected_devices is None:
            self.affected_devices = []
    
    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


class IncidentDatabase:
    """
    Local incident database with vector search
    Stores incidents for RAG retrieval
    """
    
    def __init__(self, db_path: Optional[Path] = None):
        """
        Initialize incident database
        
        Args:
            db_path: Path to incident database directory
        """
        self.db_path = db_path or Path.home() / ".overgrowth" / "incidents"
        self.db_path.mkdir(parents=True, exist_ok=True)
        
        self.incidents_file = self.db_path / "incidents.json"
        self.embeddings_file = self.db_path / "embeddings.json"
        
        self.use_chromadb = False
        self.chroma_client = None
        
        # Try to use ChromaDB for vector search
        try:
            import chromadb
            
            self.chroma_client = chromadb.PersistentClient(path=str(self.db_path / "chroma"))
            self.collection = self.chroma_client.get_or_create_collection(
                name="network_incidents",
                metadata={"description": "Network incidents and deployment failures"}
            )
            self.use_chromadb = True
            logger.info("ChromaDB initialized for incident vector search")
            
        except ImportError:
            logger.warning("chromadb not installed - using simple search")
            logger.info("Install with: pip install chromadb")
        except Exception as e:
            logger.warning(f"Failed to initialize ChromaDB: {e}")
    
    def add_incident(self, incident: Incident) -> str:
        """
        Add incident to database
        
        Args:
            incident: Incident to store
            
        Returns:
            Incident ID
        """
        # Load existing incidents
        incidents = self._load_incidents()
        
        # Add new incident
        incidents[incident.id] = incident.to_dict()
        
        # Save to file
        with open(self.incidents_file, 'w') as f:
            json.dump(incidents, f, indent=2)
        
        # Add to vector database
        if self.use_chromadb:
            self._add_to_vector_db(incident)
        
        logger.info(f"Added incident {incident.id} to database")
        return incident.id
    
    def _add_to_vector_db(self, incident: Incident):
        """Add incident to ChromaDB for vector search"""
        try:
            # Create searchable text from incident
            text = f"""
            Severity: {incident.severity}
            Category: {incident.category}
            Description: {incident.description}
            Root Cause: {incident.root_cause or 'Unknown'}
            Resolution: {incident.resolution or 'Unresolved'}
            Affected Devices: {', '.join(incident.affected_devices)}
            """
            
            # Add to collection
            self.collection.add(
                documents=[text],
                metadatas=[{
                    'incident_id': incident.id,
                    'severity': incident.severity,
                    'category': incident.category,
                    'timestamp': incident.timestamp
                }],
                ids=[incident.id]
            )
            
        except Exception as e:
            logger.error(f"Failed to add incident to vector DB: {e}")
    
    def search_similar(self, query: str, n_results: int = 5) -> List[Incident]:
        """
        Search for similar incidents using vector similarity
        
        Args:
            query: Search query (e.g., error message, symptoms)
            n_results: Number of results to return
            
        Returns:
            List of similar incidents
        """
        if self.use_chromadb:
            try:
                results = self.collection.query(
                    query_texts=[query],
                    n_results=n_results
                )
                
                incident_ids = results['ids'][0]
                incidents_dict = self._load_incidents()
                
                similar = []
                for incident_id in incident_ids:
                    if incident_id in incidents_dict:
                        incident_data = incidents_dict[incident_id]
                        similar.append(Incident(**incident_data))
                
                return similar
                
            except Exception as e:
                logger.error(f"Vector search failed: {e}")
        
        # Fallback: Simple keyword search
        return self._keyword_search(query, n_results)
    
    def _keyword_search(self, query: str, n_results: int) -> List[Incident]:
        """Fallback keyword-based search"""
        incidents_dict = self._load_incidents()
        query_lower = query.lower()
        
        matches = []
        for incident_data in incidents_dict.values():
            # Search in description, root_cause, resolution
            text = f"{incident_data.get('description', '')} {incident_data.get('root_cause', '')} {incident_data.get('resolution', '')}"
            
            if query_lower in text.lower():
                matches.append(Incident(**incident_data))
        
        return matches[:n_results]
    
    def get_incident(self, incident_id: str) -> Optional[Incident]:
        """Get incident by ID"""
        incidents = self._load_incidents()
        
        if incident_id in incidents:
            return Incident(**incidents[incident_id])
        
        return None
    
    def update_incident(self, incident_id: str, updates: Dict[str, Any]):
        """Update incident fields"""
        incidents = self._load_incidents()
        
        if incident_id in incidents:
            incidents[incident_id].update(updates)
            
            with open(self.incidents_file, 'w') as f:
                json.dump(incidents, f, indent=2)
            
            logger.info(f"Updated incident {incident_id}")
    
    def _load_incidents(self) -> Dict[str, Dict[str, Any]]:
        """Load incidents from file"""
        if not self.incidents_file.exists():
            return {}
        
        try:
            with open(self.incidents_file, 'r') as f:
                return json.load(f)
        except Exception as e:
            logger.error(f"Failed to load incidents: {e}")
            return {}
    
    def get_all_incidents(
        self,
        severity: Optional[str] = None,
        category: Optional[str] = None,
        limit: int = 100
    ) -> List[Incident]:
        """
        Get all incidents with optional filtering
        
        Args:
            severity: Filter by severity
            category: Filter by category
            limit: Max number to return
            
        Returns:
            List of incidents
        """
        incidents_dict = self._load_incidents()
        incidents = []
        
        for incident_data in incidents_dict.values():
            # Apply filters
            if severity and incident_data.get('severity') != severity:
                continue
            if category and incident_data.get('category') != category:
                continue
            
            incidents.append(Incident(**incident_data))
        
        # Sort by timestamp (newest first)
        incidents.sort(key=lambda x: x.timestamp, reverse=True)
        
        return incidents[:limit]


class RootCauseAnalyzer:
    """
    Analyzes incidents using RAG to find root causes
    Searches historical incidents for similar patterns
    """
    
    def __init__(self, incident_db: IncidentDatabase, llm_client=None):
        """
        Initialize analyzer
        
        Args:
            incident_db: Incident database for RAG
            llm_client: LLM client for analysis (Claude, GPT-4, etc.)
        """
        self.incident_db = incident_db
        self.llm_client = llm_client
    
    def analyze(self, incident: Incident) -> Dict[str, Any]:
        """
        Perform root cause analysis using RAG
        
        Args:
            incident: Incident to analyze
            
        Returns:
            Analysis results with suggested root cause
        """
        logger.info(f"Analyzing incident {incident.id}...")
        
        # 1. Search for similar historical incidents
        query = f"{incident.category} {incident.description}"
        similar = self.incident_db.search_similar(query, n_results=5)
        
        logger.info(f"Found {len(similar)} similar historical incidents")
        
        # 2. Extract patterns from similar incidents
        patterns = self._extract_patterns(similar)
        
        # 3. Generate root cause hypothesis
        if self.llm_client:
            root_cause = self._llm_analysis(incident, similar, patterns)
        else:
            root_cause = self._heuristic_analysis(incident, patterns)
        
        return {
            'suggested_root_cause': root_cause,
            'similar_incidents': [s.id for s in similar],
            'patterns_found': patterns,
            'confidence': self._calculate_confidence(similar, patterns)
        }
    
    def _extract_patterns(self, incidents: List[Incident]) -> List[str]:
        """Extract common patterns from incidents"""
        patterns = []
        
        # Look for common root causes
        root_causes = [i.root_cause for i in incidents if i.root_cause]
        if root_causes:
            # Find most common
            from collections import Counter
            common = Counter(root_causes).most_common(3)
            patterns.extend([f"Common root cause: {rc}" for rc, count in common])
        
        # Look for common device types
        all_devices = []
        for i in incidents:
            all_devices.extend(i.affected_devices)
        
        if all_devices:
            from collections import Counter
            common_devices = Counter(all_devices).most_common(3)
            patterns.append(f"Commonly affected devices: {', '.join([d for d, _ in common_devices])}")
        
        # Look for common categories
        categories = [i.category for i in incidents]
        if categories:
            from collections import Counter
            common_cat = Counter(categories).most_common(1)[0]
            patterns.append(f"Common category: {common_cat[0]}")
        
        return patterns
    
    def _heuristic_analysis(self, incident: Incident, patterns: List[str]) -> str:
        """Simple rule-based root cause analysis"""
        # Default analysis based on category
        root_causes = {
            'deployment_failure': 'Configuration syntax error or validation failure',
            'config_error': 'Invalid configuration parameters or conflicts',
            'routing_issue': 'Routing loop, missing route, or protocol misconfiguration',
            'outage': 'Link failure, device failure, or cascading failure'
        }
        
        base_cause = root_causes.get(incident.category, 'Unknown cause')
        
        # Enhance with patterns
        if patterns:
            return f"{base_cause}. Patterns from similar incidents: {patterns[0]}"
        
        return base_cause
    
    def _llm_analysis(
        self,
        incident: Incident,
        similar: List[Incident],
        patterns: List[str]
    ) -> str:
        """Use LLM for advanced root cause analysis"""
        # Build prompt with context from similar incidents
        context = "Similar historical incidents:\n"
        for s in similar[:3]:
            context += f"- {s.description} → Root cause: {s.root_cause}\n"
        
        prompt = f"""
        Analyze this network incident and suggest the root cause:
        
        Current Incident:
        Severity: {incident.severity}
        Category: {incident.category}
        Description: {incident.description}
        Affected Devices: {', '.join(incident.affected_devices)}
        
        {context}
        
        Patterns found:
        {chr(10).join(f'- {p}' for p in patterns)}
        
        Based on the incident details and similar historical incidents, what is the likely root cause?
        Provide a concise analysis.
        """
        
        try:
            # Call LLM (placeholder - integrate with Claude/GPT-4)
            response = self.llm_client.generate(prompt)
            return response
        except Exception as e:
            logger.error(f"LLM analysis failed: {e}")
            return self._heuristic_analysis(incident, patterns)
    
    def _calculate_confidence(self, similar: List[Incident], patterns: List[str]) -> float:
        """Calculate confidence score for root cause analysis"""
        confidence = 0.0
        
        # More similar incidents = higher confidence
        if len(similar) >= 3:
            confidence += 0.4
        elif len(similar) >= 1:
            confidence += 0.2
        
        # Resolved incidents = higher confidence
        resolved = [s for s in similar if s.resolution]
        if len(resolved) >= 2:
            confidence += 0.3
        elif len(resolved) >= 1:
            confidence += 0.15
        
        # Patterns found = higher confidence
        if len(patterns) >= 2:
            confidence += 0.3
        elif len(patterns) >= 1:
            confidence += 0.15
        
        return min(confidence, 1.0)


class RegressionTestGenerator:
    """
    Generates pyATS/NUTS regression tests from incidents
    Prevents similar issues from recurring
    """
    
    def generate_test(self, incident: Incident) -> str:
        """
        Generate regression test for incident
        
        Args:
            incident: Incident to create test for
            
        Returns:
            Test code as string
        """
        logger.info(f"Generating regression test for {incident.id}...")
        
        # Select test template based on category
        if incident.category == 'config_error':
            return self._generate_config_test(incident)
        elif incident.category == 'routing_issue':
            return self._generate_routing_test(incident)
        elif incident.category == 'deployment_failure':
            return self._generate_deployment_test(incident)
        else:
            return self._generate_generic_test(incident)
    
    def _generate_config_test(self, incident: Incident) -> str:
        """Generate test for configuration errors"""
        test = f'''"""
Regression test for incident {incident.id}
{incident.description}
Generated: {datetime.now().isoformat()}
"""
from pyats import aetest
from pyats.topology import loader

class Test{incident.id.replace("-", "_")}(aetest.Testcase):
    """Test that prevents recurrence of {incident.description}"""
    
    @aetest.setup
    def setup(self, testbed):
        self.devices = {{}}
        for device_name in {incident.affected_devices}:
            device = testbed.devices[device_name]
            device.connect()
            self.devices[device_name] = device
    
    @aetest.test
    def verify_configuration(self):
        """Verify configuration doesn't have the same error"""
        for device_name, device in self.devices.items():
            # Get running config
            output = device.execute('show running-config')
            
            # Check for known bad pattern
            # TODO: Customize based on root cause
            assert "invalid_config_pattern" not in output, \\
                f"Found invalid config on {{device_name}}"
    
    @aetest.cleanup
    def cleanup(self):
        for device in self.devices.values():
            device.disconnect()

if __name__ == '__main__':
    import sys
    from pyats.topology import loader
    
    testbed = loader.load('testbed.yaml')
    aetest.main(testbed=testbed)
'''
        return test
    
    def _generate_routing_test(self, incident: Incident) -> str:
        """Generate test for routing issues"""
        test = f'''"""
Regression test for routing incident {incident.id}
{incident.description}
"""
from pyats import aetest

class TestRouting{incident.id.replace("-", "_")}(aetest.Testcase):
    """Prevent routing loops and blackholes"""
    
    @aetest.test
    def verify_no_routing_loops(self, testbed):
        """Check for routing loops"""
        for device in {incident.affected_devices}:
            dev = testbed.devices[device]
            dev.connect()
            
            # Get routing table
            routes = dev.parse('show ip route')
            
            # Check for loops (routes pointing to themselves)
            for route in routes.get('vrf', {{}}).get('default', {{}}).get('routes', {{}}).values():
                next_hop = route.get('next_hop', {{}})
                assert next_hop != dev.name, f"Routing loop detected on {{device}}"
            
            dev.disconnect()
    
    @aetest.test  
    def verify_reachability(self, testbed):
        """Verify critical subnets are reachable"""
        # TODO: Add reachability checks for affected subnets
        pass

if __name__ == '__main__':
    aetest.main()
'''
        return test
    
    def _generate_deployment_test(self, incident: Incident) -> str:
        """Generate test for deployment failures"""
        return f'''"""
Regression test for deployment failure {incident.id}
Validates deployment process before executing
"""
import pytest

def test_pre_deployment_validation():
    """Pre-flight checks to prevent {incident.description}"""
    # TODO: Add validation checks
    assert True, "Validation passed"

def test_config_syntax():
    """Verify config syntax is valid"""
    # TODO: Parse and validate configs
    assert True, "Syntax valid"
'''
    
    def _generate_generic_test(self, incident: Incident) -> str:
        """Generic test template"""
        return f'''"""
Regression test for {incident.id}
{incident.description}
"""
import pytest

def test_{incident.id.replace("-", "_")}():
    """Prevent recurrence of {incident.description}"""
    # TODO: Implement test logic
    assert True, "Test passed"
'''


def capture_deployment_failure(
    description: str,
    network_model: Dict[str, Any],
    validation_errors: List[Dict[str, Any]],
    affected_devices: List[str]
) -> Incident:
    """
    Capture deployment failure as incident
    
    Args:
        description: What went wrong
        network_model: Network model that failed
        validation_errors: Validation errors encountered
        affected_devices: List of affected device names
        
    Returns:
        Created incident
    """
    incident_id = f"deploy-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
    
    incident = Incident(
        id=incident_id,
        timestamp=datetime.now().isoformat(),
        severity='high',
        category='deployment_failure',
        description=description,
        affected_devices=affected_devices,
        network_model=network_model,
        validation_errors=validation_errors
    )
    
    # Store in database
    db = IncidentDatabase()
    db.add_incident(incident)
    
    logger.warning(f"Captured deployment failure: {incident_id}")
    
    return incident


def learn_from_incident(incident: Incident) -> Dict[str, Any]:
    """
    Complete learning cycle: analyze, test, update
    
    Args:
        incident: Incident to learn from
        
    Returns:
        Learning outcomes
    """
    db = IncidentDatabase()
    analyzer = RootCauseAnalyzer(db)
    test_gen = RegressionTestGenerator()
    
    # 1. Root cause analysis
    analysis = analyzer.analyze(incident)
    logger.info(f"Root cause: {analysis['suggested_root_cause']}")
    
    # 2. Generate regression test
    test_code = test_gen.generate_test(incident)
    
    # Save test
    test_path = Path("tests") / "regression" / f"test_{incident.id}.py"
    test_path.parent.mkdir(parents=True, exist_ok=True)
    test_path.write_text(test_code)
    
    logger.info(f"Generated regression test: {test_path}")
    
    # 3. Update incident with learnings
    db.update_incident(incident.id, {
        'root_cause': analysis['suggested_root_cause'],
        'regression_test': str(test_path)
    })
    
    return {
        'incident_id': incident.id,
        'root_cause': analysis['suggested_root_cause'],
        'regression_test': str(test_path),
        'similar_incidents': analysis['similar_incidents'],
        'confidence': analysis['confidence']
    }