"""
Migration Commands

Commands for migrating data from the existing web application database
to the new file-based directory structure.
"""

import asyncio
import json
import os
import sys
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Any
import click

# Import the base command class
from .base_command import BaseCommand

# Import database models from the existing backend
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'backend'))

class MigrateCommands(BaseCommand):
    """Commands for migrating from database to file structure"""
    
    async def migrate_from_database(self, database_url: Optional[str], batch_size: int, dry_run: bool):
        """Migrate data from existing database to file structure"""
        
        try:
            # Import database components  
            from database_models import CVE, SigmaRule, RuleTemplate, SessionLocal
            
            # Use existing database session factory
            if database_url:
                self.info(f"Using provided database URL")
                # Create new engine with provided URL
                from sqlalchemy import create_engine
                from sqlalchemy.orm import sessionmaker
                engine = create_engine(database_url)
                SessionFactory = sessionmaker(autocommit=False, autoflush=False, bind=engine)
                db = SessionFactory()
            else:
                # Use default session factory
                db = SessionLocal()
            
            # Get total counts
            cve_count = db.query(CVE).count()
            rule_count = db.query(SigmaRule).count()
            template_count = db.query(RuleTemplate).count()
            
            self.info(f"Found {cve_count} CVEs, {rule_count} SIGMA rules, {template_count} templates")
            
            if dry_run:
                self.warning("DRY RUN MODE - No files will be created")
                
                # Show what would be migrated
                sample_cves = db.query(CVE).limit(5).all()
                for cve in sample_cves:
                    cve_dir = self.get_cve_directory(cve.cve_id)
                    self.info(f"Would create: {cve_dir}")
                    
                    # Count rules for this CVE
                    rules = db.query(SigmaRule).filter(SigmaRule.cve_id == cve.cve_id).all()
                    self.info(f"  - Would migrate {len(rules)} SIGMA rules")
                
                return
            
            # Migrate CVEs and rules
            migrated_cves = 0
            migrated_rules = 0
            
            # Process CVEs in batches
            offset = 0
            while offset < cve_count:
                batch_cves = db.query(CVE).offset(offset).limit(batch_size).all()
                
                for cve in batch_cves:
                    try:
                        await self._migrate_cve(db, cve)
                        migrated_cves += 1
                        
                        # Migrate associated rules
                        rules = db.query(SigmaRule).filter(SigmaRule.cve_id == cve.cve_id).all()
                        for rule in rules:
                            if await self._migrate_sigma_rule(cve.cve_id, rule):
                                migrated_rules += 1
                        
                        if migrated_cves % 10 == 0:
                            self.info(f"Migrated {migrated_cves}/{cve_count} CVEs...")
                            
                    except Exception as e:
                        self.error(f"Error migrating {cve.cve_id}: {e}")
                
                offset += batch_size
            
            # Migrate templates to new location
            template_dir = self.config.base_dir / "backend" / "templates"
            template_dir.mkdir(exist_ok=True)
            
            templates = db.query(RuleTemplate).all()
            for template in templates:
                template_file = template_dir / f"{template.template_name.lower().replace(' ', '_')}.yaml"
                if not template_file.exists():
                    try:
                        with open(template_file, 'w') as f:
                            f.write(template.template_content)
                        self.info(f"Migrated template: {template.template_name}")
                    except Exception as e:
                        self.error(f"Error migrating template {template.template_name}: {e}")
            
            db.close()
            
            self.success(f"Migration completed!")
            self.success(f"Migrated {migrated_cves} CVEs and {migrated_rules} SIGMA rules")
            
        except ImportError as e:
            self.error(f"Could not import database models: {e}")
            self.error("Make sure you're running from the project root directory")
        except Exception as e:
            self.error(f"Migration failed: {e}")
            import traceback
            traceback.print_exc()
    
    async def _migrate_cve(self, db, cve) -> bool:
        """Migrate a single CVE to file structure"""
        try:
            # Create CVE metadata
            metadata = {
                "cve_info": {
                    "cve_id": cve.cve_id,
                    "description": cve.description,
                    "cvss_score": float(cve.cvss_score) if cve.cvss_score else None,
                    "severity": cve.severity,
                    "published_date": cve.published_date.isoformat() if cve.published_date else None,
                    "modified_date": cve.modified_date.isoformat() if cve.modified_date else None,
                    "affected_products": cve.affected_products or [],
                    "reference_urls": cve.reference_urls or []
                },
                "poc_data": {
                    "poc_count": getattr(cve, 'poc_count', 0),
                    "poc_data": getattr(cve, 'poc_data', {}),
                    "nomi_sec_data": getattr(cve, 'poc_data', {}).get('nomi_sec', []) if getattr(cve, 'poc_data', {}) else [],
                    "github_pocs": getattr(cve, 'poc_data', {}).get('github', []) if getattr(cve, 'poc_data', {}) else []
                },
                "processing": {
                    "data_source": getattr(cve, 'data_source', 'nvd_api'),
                    "bulk_processed": getattr(cve, 'bulk_processed', False),
                    "reference_sync_status": getattr(cve, 'reference_sync_status', 'pending')
                },
                "file_manifest": [],
                "rule_generation": {},
                "created_at": cve.created_at.isoformat() if cve.created_at else datetime.utcnow().isoformat(),
                "updated_at": datetime.utcnow().isoformat(),
                "migrated_at": datetime.utcnow().isoformat()
            }
            
            # Save PoC analysis if available
            if hasattr(cve, 'poc_data') and cve.poc_data:
                cve_dir = self.ensure_cve_directory(cve.cve_id)
                poc_analysis_file = cve_dir / "poc_analysis.json"
                
                with open(poc_analysis_file, 'w') as f:
                    json.dump(cve.poc_data, f, indent=2, default=str)
                
                metadata["file_manifest"].append("poc_analysis.json")
            
            # Save metadata
            return self.save_cve_metadata(cve.cve_id, metadata)
            
        except Exception as e:
            self.error(f"Error migrating CVE {cve.cve_id}: {e}")
            return False
    
    async def _migrate_sigma_rule(self, cve_id: str, rule) -> bool:
        """Migrate a single SIGMA rule to file structure"""
        try:
            # Determine rule filename based on generation method/source
            if hasattr(rule, 'poc_source') and rule.poc_source:
                if 'llm' in rule.poc_source.lower() or 'openai' in rule.poc_source.lower():
                    filename = "rule_llm_openai.sigma"
                elif 'anthropic' in rule.poc_source.lower():
                    filename = "rule_llm_anthropic.sigma" 
                elif 'hybrid' in rule.poc_source.lower():
                    filename = "rule_hybrid.sigma"
                else:
                    filename = "rule_template.sigma"
            else:
                # Default to template-based
                filename = "rule_template.sigma"
            
            # Check if we already have a rule with this name, if so append a suffix
            existing_rules = self.list_cve_rules(cve_id)
            if filename in existing_rules:
                base_name = filename.replace('.sigma', '')
                counter = 1
                while f"{base_name}_{counter}.sigma" in existing_rules:
                    counter += 1
                filename = f"{base_name}_{counter}.sigma"
            
            # Save the rule content
            if self.save_sigma_rule(cve_id, filename, rule.rule_content):
                # Update metadata with additional rule information
                metadata = self.load_cve_metadata(cve_id)
                if metadata:
                    rule_info = {
                        "rule_name": rule.rule_name,
                        "detection_type": getattr(rule, 'detection_type', ''),
                        "log_source": getattr(rule, 'log_source', ''),
                        "confidence_level": getattr(rule, 'confidence_level', ''),
                        "auto_generated": getattr(rule, 'auto_generated', True),
                        "exploit_based": getattr(rule, 'exploit_based', False),
                        "poc_source": getattr(rule, 'poc_source', 'template'),
                        "poc_quality_score": getattr(rule, 'poc_quality_score', 0),
                        "github_repos": getattr(rule, 'github_repos', []),
                        "created_at": rule.created_at.isoformat() if rule.created_at else None,
                        "migrated_at": datetime.utcnow().isoformat()
                    }
                    
                    method_key = filename.replace('rule_', '').replace('.sigma', '')
                    if 'rule_generation' not in metadata:
                        metadata['rule_generation'] = {}
                    
                    metadata['rule_generation'][method_key] = rule_info
                    self.save_cve_metadata(cve_id, metadata)
                
                return True
            
        except Exception as e:
            self.error(f"Error migrating rule for {cve_id}: {e}")
            return False
        
        return False
    
    async def validate_migration(self, year: Optional[int] = None):
        """Validate migrated data integrity"""
        self.info("Validating migrated data...")
        
        issues = []
        validated_cves = 0
        validated_rules = 0
        
        # Get CVEs to validate
        cves_to_check = self.get_all_cves(year)
        
        for cve_id in cves_to_check:
            try:
                # Check if metadata exists and is valid
                metadata = self.load_cve_metadata(cve_id)
                if not metadata:
                    issues.append(f"{cve_id}: Missing metadata.json")
                    continue
                
                # Validate required metadata fields
                required_fields = ['cve_info', 'poc_data', 'processing']
                for field in required_fields:
                    if field not in metadata:
                        issues.append(f"{cve_id}: Missing metadata field '{field}'")
                
                # Validate CVE info
                if 'cve_info' in metadata:
                    cve_info = metadata['cve_info']
                    if not cve_info.get('cve_id'):
                        issues.append(f"{cve_id}: Missing cve_id in metadata")
                    elif cve_info['cve_id'] != cve_id:
                        issues.append(f"{cve_id}: CVE ID mismatch in metadata")
                
                # Validate file manifest
                file_manifest = metadata.get('file_manifest', [])
                cve_dir = self.get_cve_directory(cve_id)
                
                for file_name in file_manifest:
                    file_path = cve_dir / file_name
                    if not file_path.exists():
                        issues.append(f"{cve_id}: Referenced file '{file_name}' does not exist")
                
                # Check for SIGMA rule files
                rule_files = self.list_cve_rules(cve_id)
                for rule_file in rule_files:
                    rule_content = self.load_sigma_rule(cve_id, rule_file)
                    if not rule_content:
                        issues.append(f"{cve_id}: Could not load rule file '{rule_file}'")
                    elif not rule_content.strip():
                        issues.append(f"{cve_id}: Empty rule file '{rule_file}'")
                    else:
                        # Basic YAML validation for SIGMA rules
                        if not rule_content.strip().startswith('title:'):
                            issues.append(f"{cve_id}: Rule '{rule_file}' doesn't appear to be valid SIGMA format")
                        validated_rules += 1
                
                validated_cves += 1
                
                if validated_cves % 100 == 0:
                    self.info(f"Validated {validated_cves} CVEs...")
                    
            except Exception as e:
                issues.append(f"{cve_id}: Validation error - {e}")
        
        # Print validation results
        self.info(f"\nValidation completed:")
        self.info(f"- Validated {validated_cves} CVEs")
        self.info(f"- Validated {validated_rules} SIGMA rules")
        
        if issues:
            self.warning(f"Found {len(issues)} validation issues:")
            for issue in issues[:20]:  # Show first 20 issues
                self.error(f"  {issue}")
            
            if len(issues) > 20:
                self.warning(f"  ... and {len(issues) - 20} more issues")
        else:
            self.success("No validation issues found!")
    
    async def cleanup_migration(self):
        """Clean up migration artifacts and temporary files"""
        self.info("Cleaning up migration artifacts...")
        
        # Remove empty directories
        for year_dir in self.config.cves_dir.iterdir():
            if year_dir.is_dir():
                for cve_dir in year_dir.iterdir():
                    if cve_dir.is_dir():
                        # Check if directory is empty
                        if not any(cve_dir.iterdir()):
                            cve_dir.rmdir()
                            self.info(f"Removed empty directory: {cve_dir}")
                
                # Check if year directory is now empty
                if not any(year_dir.iterdir()):
                    year_dir.rmdir()
                    self.info(f"Removed empty year directory: {year_dir}")
        
        self.success("Cleanup completed!")
    
    async def migration_stats(self):
        """Show migration statistics"""
        self.info("Migration Statistics:")
        
        years = self.get_years_with_data()
        total_cves = 0
        total_rules = 0
        
        stats_by_year = {}
        
        for year in years:
            cves = self.get_all_cves(year)
            year_cves = len(cves)
            year_rules = 0
            
            for cve_id in cves:
                rules = self.list_cve_rules(cve_id)
                year_rules += len(rules)
            
            stats_by_year[year] = {
                'cves': year_cves,
                'rules': year_rules
            }
            
            total_cves += year_cves
            total_rules += year_rules
        
        # Print statistics table
        headers = ["Year", "CVEs", "Rules", "Avg Rules/CVE"]
        rows = []
        
        for year in sorted(years):
            stats = stats_by_year[year]
            avg_rules = stats['rules'] / stats['cves'] if stats['cves'] > 0 else 0
            rows.append([
                str(year),
                str(stats['cves']),
                str(stats['rules']),
                f"{avg_rules:.1f}"
            ])
        
        # Add totals
        avg_total = total_rules / total_cves if total_cves > 0 else 0
        rows.append(["TOTAL", str(total_cves), str(total_rules), f"{avg_total:.1f}"])
        
        self.print_table(headers, rows, "Migration Statistics by Year")