"""Data quality validation for prompt optimization."""

from typing import List, Dict, Any, Optional, Set
import re
from collections import Counter
import statistics

from ..core.base import Dataset, Example


class DataQualityValidator:
    """Validate data quality for prompt optimization."""
    
    def __init__(self):
        self.validation_results = {}
    
    def validate_dataset(self, dataset: Dataset, 
                        check_pii: bool = True,
                        check_bias: bool = True,
                        check_format: bool = True) -> Dict[str, Any]:
        """Comprehensive dataset validation."""
        results = {
            "total_examples": len(dataset),
            "passed": True,
            "warnings": [],
            "errors": [],
            "statistics": self._calculate_statistics(dataset),
            "quality_scores": {}
        }
        
        # Basic validation
        if len(dataset) == 0:
            results["errors"].append("Dataset is empty")
            results["passed"] = False
            return results
        
        # Format validation
        if check_format:
            format_results = self._validate_format(dataset)
            results["quality_scores"]["format"] = format_results["score"]
            results["warnings"].extend(format_results["warnings"])
            results["errors"].extend(format_results["errors"])
        
        # PII detection
        if check_pii:
            pii_results = self._check_pii(dataset)
            results["quality_scores"]["privacy"] = pii_results["score"]
            results["warnings"].extend(pii_results["warnings"])
        
        # Bias detection
        if check_bias:
            bias_results = self._check_bias(dataset)
            results["quality_scores"]["bias"] = bias_results["score"]
            results["warnings"].extend(bias_results["warnings"])
        
        # Data quality checks
        quality_results = self._check_data_quality(dataset)
        results["quality_scores"]["quality"] = quality_results["score"]
        results["warnings"].extend(quality_results["warnings"])
        
        # Overall score
        if results["quality_scores"]:
            results["overall_score"] = sum(results["quality_scores"].values()) / len(results["quality_scores"])
        else:
            results["overall_score"] = 0.0
        
        # Determine if validation passed
        results["passed"] = len(results["errors"]) == 0 and results["overall_score"] >= 0.7
        
        return results
    
    def _calculate_statistics(self, dataset: Dataset) -> Dict[str, Any]:
        """Calculate dataset statistics."""
        input_lengths = [len(ex.input.split()) for ex in dataset.examples]
        output_lengths = [len(ex.output.split()) for ex in dataset.examples]
        
        return {
            "input_length": {
                "mean": statistics.mean(input_lengths) if input_lengths else 0,
                "median": statistics.median(input_lengths) if input_lengths else 0,
                "min": min(input_lengths) if input_lengths else 0,
                "max": max(input_lengths) if input_lengths else 0,
                "std": statistics.stdev(input_lengths) if len(input_lengths) > 1 else 0
            },
            "output_length": {
                "mean": statistics.mean(output_lengths) if output_lengths else 0,
                "median": statistics.median(output_lengths) if output_lengths else 0,
                "min": min(output_lengths) if output_lengths else 0,
                "max": max(output_lengths) if output_lengths else 0,
                "std": statistics.stdev(output_lengths) if len(output_lengths) > 1 else 0
            },
            "unique_inputs": len(set(ex.input for ex in dataset.examples)),
            "unique_outputs": len(set(ex.output for ex in dataset.examples)),
            "duplicate_rate": 1 - (len(set(ex.input for ex in dataset.examples)) / len(dataset))
        }
    
    def _validate_format(self, dataset: Dataset) -> Dict[str, Any]:
        """Validate data format consistency."""
        warnings = []
        errors = []
        score = 1.0
        
        # Check for empty inputs/outputs
        empty_inputs = sum(1 for ex in dataset.examples if not ex.input.strip())
        empty_outputs = sum(1 for ex in dataset.examples if not ex.output.strip())
        
        if empty_inputs > 0:
            errors.append(f"{empty_inputs} examples have empty inputs")
            score -= 0.3
        
        if empty_outputs > 0:
            errors.append(f"{empty_outputs} examples have empty outputs")
            score -= 0.3
        
        # Check for extremely short or long examples
        very_short = sum(1 for ex in dataset.examples 
                        if len(ex.input.split()) < 3 or len(ex.output.split()) < 3)
        very_long = sum(1 for ex in dataset.examples 
                       if len(ex.input.split()) > 500 or len(ex.output.split()) > 500)
        
        if very_short > len(dataset) * 0.1:
            warnings.append(f"{very_short} examples are very short (< 3 words)")
            score -= 0.1
        
        if very_long > 0:
            warnings.append(f"{very_long} examples are very long (> 500 words)")
            score -= 0.1
        
        return {
            "score": max(0, score),
            "warnings": warnings,
            "errors": errors
        }
    
    def _check_pii(self, dataset: Dataset) -> Dict[str, Any]:
        """Check for potential PII in dataset."""
        warnings = []
        score = 1.0
        
        pii_patterns = {
            "email": r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',
            "phone": r'\b\d{3}[-.\s]?\d{3}[-.\s]?\d{4}\b',
            "ssn": r'\b\d{3}-\d{2}-\d{4}\b',
            "credit_card": r'\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b',
            "ip_address": r'\b(?:\d{1,3}\.){3}\d{1,3}\b'
        }
        
        pii_found = {pii_type: 0 for pii_type in pii_patterns}
        
        for example in dataset.examples:
            text = example.input + " " + example.output
            for pii_type, pattern in pii_patterns.items():
                if re.search(pattern, text):
                    pii_found[pii_type] += 1
        
        for pii_type, count in pii_found.items():
            if count > 0:
                warnings.append(f"Found {count} potential {pii_type.replace('_', ' ')} instances")
                score -= 0.2
        
        return {
            "score": max(0, score),
            "warnings": warnings,
            "errors": []
        }
    
    def _check_bias(self, dataset: Dataset) -> Dict[str, Any]:
        """Check for potential bias in dataset."""
        warnings = []
        score = 1.0
        
        # Gender bias indicators
        gender_terms = {
            "male": ["he", "him", "his", "man", "men", "boy", "male"],
            "female": ["she", "her", "hers", "woman", "women", "girl", "female"]
        }
        
        gender_counts = {"male": 0, "female": 0}
        
        for example in dataset.examples:
            text = (example.input + " " + example.output).lower()
            for gender, terms in gender_terms.items():
                for term in terms:
                    if f" {term} " in f" {text} ":
                        gender_counts[gender] += 1
                        break
        
        total_gendered = sum(gender_counts.values())
        if total_gendered > 0:
            male_ratio = gender_counts["male"] / total_gendered
            female_ratio = gender_counts["female"] / total_gendered
            
            if abs(male_ratio - female_ratio) > 0.3:
                warnings.append(f"Gender imbalance detected: {male_ratio:.1%} male, {female_ratio:.1%} female")
                score -= 0.2
        
        # Check for other bias indicators
        bias_terms = [
            "ethnic", "racial", "religious", "political",
            "nationality", "disability", "age-related"
        ]
        
        # This is a simplified check - in production would use more sophisticated methods
        
        return {
            "score": max(0, score),
            "warnings": warnings,
            "errors": []
        }
    
    def _check_data_quality(self, dataset: Dataset) -> Dict[str, Any]:
        """Check overall data quality."""
        warnings = []
        score = 1.0
        
        # Check for high duplication
        stats = self._calculate_statistics(dataset)
        if stats["duplicate_rate"] > 0.1:
            warnings.append(f"High duplication rate: {stats['duplicate_rate']:.1%}")
            score -= 0.2
        
        # Check for good diversity
        vocab_diversity = self._calculate_vocabulary_diversity(dataset)
        if vocab_diversity < 0.3:
            warnings.append(f"Low vocabulary diversity: {vocab_diversity:.2f}")
            score -= 0.2
        
        # Check for balanced lengths
        input_cv = stats["input_length"]["std"] / stats["input_length"]["mean"] if stats["input_length"]["mean"] > 0 else 0
        output_cv = stats["output_length"]["std"] / stats["output_length"]["mean"] if stats["output_length"]["mean"] > 0 else 0
        
        if input_cv > 1.0 or output_cv > 1.0:
            warnings.append("High variance in example lengths")
            score -= 0.1
        
        return {
            "score": max(0, score),
            "warnings": warnings,
            "errors": []
        }
    
    def _calculate_vocabulary_diversity(self, dataset: Dataset) -> float:
        """Calculate vocabulary diversity score."""
        all_words = []
        for example in dataset.examples:
            words = example.input.lower().split() + example.output.lower().split()
            all_words.extend(words)
        
        if not all_words:
            return 0.0
        
        unique_words = len(set(all_words))
        total_words = len(all_words)
        
        # Type-token ratio
        return unique_words / total_words


class ComplianceValidator:
    """Validate data compliance with business requirements."""
    
    def __init__(self, compliance_rules: Optional[List[str]] = None):
        self.compliance_rules = compliance_rules or []
    
    def validate_for_compliance(self, dataset: Dataset, 
                              business_context: Dict[str, Any]) -> Dict[str, Any]:
        """Validate dataset against compliance requirements."""
        results = {
            "compliant": True,
            "violations": [],
            "warnings": []
        }
        
        # Check each compliance requirement
        for rule in self.compliance_rules:
            if rule == "GDPR":
                gdpr_results = self._check_gdpr_compliance(dataset)
                results["violations"].extend(gdpr_results["violations"])
                results["warnings"].extend(gdpr_results["warnings"])
            
            elif rule == "HIPAA":
                hipaa_results = self._check_hipaa_compliance(dataset)
                results["violations"].extend(hipaa_results["violations"])
                results["warnings"].extend(hipaa_results["warnings"])
            
            elif rule == "SOC2":
                soc2_results = self._check_soc2_compliance(dataset)
                results["violations"].extend(soc2_results["violations"])
                results["warnings"].extend(soc2_results["warnings"])
        
        results["compliant"] = len(results["violations"]) == 0
        return results
    
    def _check_gdpr_compliance(self, dataset: Dataset) -> Dict[str, Any]:
        """Check GDPR compliance."""
        violations = []
        warnings = []
        
        # Check for personal data
        personal_data_patterns = {
            "names": r'\b[A-Z][a-z]+ [A-Z][a-z]+\b',
            "addresses": r'\b\d+\s+[A-Za-z\s]+(?:Street|St|Avenue|Ave|Road|Rd|Boulevard|Blvd)\b',
            "ids": r'\b(?:ID|id):\s*\w+\b'
        }
        
        for pattern_name, pattern in personal_data_patterns.items():
            found = False
            for example in dataset.examples:
                if re.search(pattern, example.input + " " + example.output):
                    found = True
                    break
            
            if found:
                warnings.append(f"Potential {pattern_name} found - ensure proper consent")
        
        return {"violations": violations, "warnings": warnings}
    
    def _check_hipaa_compliance(self, dataset: Dataset) -> Dict[str, Any]:
        """Check HIPAA compliance."""
        violations = []
        warnings = []
        
        # Check for health information
        health_terms = ["diagnosis", "treatment", "medical", "patient", "prescription"]
        
        for example in dataset.examples:
            text = (example.input + " " + example.output).lower()
            if any(term in text for term in health_terms):
                warnings.append("Contains potential health information - ensure PHI is properly handled")
                break
        
        return {"violations": violations, "warnings": warnings}
    
    def _check_soc2_compliance(self, dataset: Dataset) -> Dict[str, Any]:
        """Check SOC2 compliance."""
        violations = []
        warnings = []
        
        # Check for security-related data
        security_terms = ["password", "secret", "key", "token", "credential"]
        
        for example in dataset.examples:
            text = (example.input + " " + example.output).lower()
            if any(term in text for term in security_terms):
                violations.append("Contains potential security credentials")
                break
        
        return {"violations": violations, "warnings": warnings}