"""
Metrics module for cheat sheet quality evaluation.

Includes text quality metrics, LaTeX validation, and coverage metrics.
"""

import re
import subprocess
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Set

from ..infra.logging import get_logger
from ..infra.tracing import trace_function

logger = get_logger(__name__)


@dataclass
class CheatSheetMetrics:
    """Comprehensive cheat sheet quality metrics."""
    
    # Content metrics
    concept_count: int = 0
    formula_count: int = 0
    algorithm_count: int = 0
    model_count: int = 0
    
    # Coverage metrics
    topic_coverage: float = 0.0
    source_coverage: float = 0.0
    
    # Quality metrics
    latex_valid: bool = False
    latex_errors: List[str] = field(default_factory=list)
    compiles_successfully: bool = False
    
    # Text quality
    avg_concept_length: float = 0.0
    avg_formula_complexity: float = 0.0
    
    # Diversity
    unique_topics: int = 0
    vocabulary_richness: float = 0.0
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "content": {
                "concept_count": self.concept_count,
                "formula_count": self.formula_count,
                "algorithm_count": self.algorithm_count,
                "model_count": self.model_count,
            },
            "coverage": {
                "topic_coverage": self.topic_coverage,
                "source_coverage": self.source_coverage,
            },
            "quality": {
                "latex_valid": self.latex_valid,
                "compiles_successfully": self.compiles_successfully,
                "latex_errors": self.latex_errors[:5],  # Limit
            },
            "text_quality": {
                "avg_concept_length": self.avg_concept_length,
                "avg_formula_complexity": self.avg_formula_complexity,
            },
            "diversity": {
                "unique_topics": self.unique_topics,
                "vocabulary_richness": self.vocabulary_richness,
            },
        }


def calculate_bleu(reference: str, hypothesis: str, n_gram: int = 4) -> float:
    """
    Calculate BLEU score between reference and hypothesis.
    
    Args:
        reference: Reference text
        hypothesis: Generated text
        n_gram: Maximum n-gram to consider
        
    Returns:
        BLEU score between 0 and 1
    """
    try:
        from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
        
        reference_tokens = reference.lower().split()
        hypothesis_tokens = hypothesis.lower().split()
        
        # Use smoothing to handle zero counts
        smoother = SmoothingFunction().method1
        
        score = sentence_bleu(
            [reference_tokens],
            hypothesis_tokens,
            weights=[1/n_gram] * n_gram,
            smoothing_function=smoother
        )
        
        return score
        
    except ImportError:
        # Fallback to simple n-gram overlap
        return _simple_ngram_overlap(reference, hypothesis, n_gram)


def _simple_ngram_overlap(text1: str, text2: str, n: int = 4) -> float:
    """Simple n-gram overlap calculation."""
    def get_ngrams(text: str, n: int) -> Set[tuple]:
        tokens = text.lower().split()
        return set(tuple(tokens[i:i+n]) for i in range(len(tokens)-n+1))
    
    scores = []
    for k in range(1, n+1):
        ngrams1 = get_ngrams(text1, k)
        ngrams2 = get_ngrams(text2, k)
        
        if not ngrams2:
            continue
        
        overlap = len(ngrams1 & ngrams2)
        precision = overlap / len(ngrams2) if ngrams2 else 0.0
        scores.append(precision)
    
    if not scores:
        return 0.0
    
    # Geometric mean
    import math
    return math.exp(sum(math.log(max(s, 1e-10)) for s in scores) / len(scores))


def calculate_rouge(reference: str, hypothesis: str) -> Dict[str, float]:
    """
    Calculate ROUGE scores.
    
    Args:
        reference: Reference text
        hypothesis: Generated text
        
    Returns:
        Dictionary with rouge-1, rouge-2, rouge-l scores
    """
    try:
        from rouge_score import rouge_scorer
        
        scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2', 'rougeL'])
        scores = scorer.score(reference, hypothesis)
        
        return {
            "rouge-1": scores['rouge1'].fmeasure,
            "rouge-2": scores['rouge2'].fmeasure,
            "rouge-l": scores['rougeL'].fmeasure,
        }
        
    except ImportError:
        # Fallback
        return _simple_rouge(reference, hypothesis)


def _simple_rouge(reference: str, hypothesis: str) -> Dict[str, float]:
    """Simple ROUGE calculation."""
    ref_words = set(reference.lower().split())
    hyp_words = set(hypothesis.lower().split())
    
    if not hyp_words or not ref_words:
        return {"rouge-1": 0.0, "rouge-2": 0.0, "rouge-l": 0.0}
    
    # ROUGE-1 (unigram)
    overlap = len(ref_words & hyp_words)
    precision = overlap / len(hyp_words)
    recall = overlap / len(ref_words)
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
    
    return {
        "rouge-1": f1,
        "rouge-2": f1 * 0.8,  # Approximation
        "rouge-l": f1 * 0.9,  # Approximation
    }


@trace_function(name="evaluate_latex_quality")
def evaluate_latex_quality(latex_source: str) -> Dict[str, Any]:
    """
    Evaluate LaTeX source quality.
    
    Checks for:
    - Syntax validity
    - Common errors
    - Structure completeness
    - Compilation (if compiler available)
    
    Args:
        latex_source: LaTeX source string
        
    Returns:
        Quality evaluation results
    """
    results = {
        "valid": True,
        "errors": [],
        "warnings": [],
        "structure": {},
        "compiles": False,
    }
    
    # Check basic structure
    results["structure"] = _check_latex_structure(latex_source)
    
    # Check for common errors
    errors = _check_latex_errors(latex_source)
    results["errors"] = errors
    results["valid"] = len(errors) == 0
    
    # Check for warnings
    results["warnings"] = _check_latex_warnings(latex_source)
    
    # Try to compile
    results["compiles"] = _try_compile_latex(latex_source)
    
    return results


def _check_latex_structure(latex: str) -> Dict[str, bool]:
    """Check LaTeX document structure."""
    structure = {
        "has_documentclass": r"\documentclass" in latex,
        "has_begin_document": r"\begin{document}" in latex,
        "has_end_document": r"\end{document}" in latex,
        "balanced_braces": _check_balanced_braces(latex),
        "balanced_environments": _check_balanced_environments(latex),
    }
    return structure


def _check_balanced_braces(latex: str) -> bool:
    """Check if braces are balanced."""
    # Remove escaped braces
    cleaned = re.sub(r'\\[{}]', '', latex)
    
    count = 0
    for char in cleaned:
        if char == '{':
            count += 1
        elif char == '}':
            count -= 1
            if count < 0:
                return False
    
    return count == 0


def _check_balanced_environments(latex: str) -> bool:
    """Check if LaTeX environments are balanced."""
    # Find all \begin{...} and \end{...}
    begins = re.findall(r'\\begin\{(\w+)\}', latex)
    ends = re.findall(r'\\end\{(\w+)\}', latex)
    
    # Check counts match
    from collections import Counter
    begin_counts = Counter(begins)
    end_counts = Counter(ends)
    
    return begin_counts == end_counts


def _check_latex_errors(latex: str) -> List[str]:
    """Check for common LaTeX errors."""
    errors = []
    
    # Missing document class
    if r"\documentclass" not in latex:
        errors.append("Missing \\documentclass")
    
    # Missing document environment
    if r"\begin{document}" not in latex:
        errors.append("Missing \\begin{document}")
    if r"\end{document}" not in latex:
        errors.append("Missing \\end{document}")
    
    # Unbalanced braces
    if not _check_balanced_braces(latex):
        errors.append("Unbalanced braces")
    
    # Unbalanced environments
    if not _check_balanced_environments(latex):
        errors.append("Unbalanced environments")
    
    # Undefined control sequences (common ones)
    # Note: This is a simplified check
    
    return errors


def _check_latex_warnings(latex: str) -> List[str]:
    """Check for potential LaTeX issues."""
    warnings = []
    
    # Check for potentially unescaped special characters
    if '&' in latex and r'\&' not in latex and 'tabular' not in latex:
        warnings.append("Possible unescaped '&' character")
    
    if '%' in latex and r'\%' not in latex:
        # Check if it's a comment
        lines = latex.split('\n')
        for line in lines:
            if '%' in line:
                idx = line.index('%')
                if idx > 0 and line[idx-1] != '\\':
                    # It's likely a comment, which is fine
                    pass
    
    # Check for very long lines
    lines = latex.split('\n')
    long_lines = sum(1 for line in lines if len(line) > 200)
    if long_lines > 5:
        warnings.append(f"{long_lines} lines exceed 200 characters")
    
    return warnings


def _try_compile_latex(latex: str) -> bool:
    """Try to compile LaTeX to check validity."""
    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            tex_path = Path(tmpdir) / "test.tex"
            tex_path.write_text(latex, encoding='utf-8')
            
            # Try pdflatex with nonstop mode
            result = subprocess.run(
                ["pdflatex", "-interaction=nonstopmode", "-halt-on-error", "test.tex"],
                cwd=tmpdir,
                capture_output=True,
                timeout=60
            )
            
            # Check if PDF was generated
            pdf_path = Path(tmpdir) / "test.pdf"
            return pdf_path.exists()
            
    except (subprocess.TimeoutExpired, FileNotFoundError, Exception) as e:
        logger.debug(f"LaTeX compilation check failed: {e}")
        return False


@trace_function(name="evaluate_coverage")
def evaluate_coverage(
    cheat_sheet_content: Dict[str, Any],
    source_chunks: List[Dict[str, Any]],
    expected_topics: Optional[List[str]] = None
) -> Dict[str, float]:
    """
    Evaluate how well the cheat sheet covers the source material.
    
    Args:
        cheat_sheet_content: Generated cheat sheet content
        source_chunks: Original source chunks
        expected_topics: Expected topics to cover (optional)
        
    Returns:
        Coverage metrics
    """
    # Extract text from cheat sheet
    cs_text = _extract_text_from_content(cheat_sheet_content)
    cs_words = set(cs_text.lower().split())
    
    # Extract text from source
    source_text = " ".join(chunk.get("content", "") for chunk in source_chunks)
    source_words = set(source_text.lower().split())
    
    # Remove common words
    stopwords = {
        'the', 'a', 'an', 'is', 'are', 'was', 'were', 'be', 'been',
        'being', 'have', 'has', 'had', 'do', 'does', 'did', 'will',
        'would', 'could', 'should', 'may', 'might', 'must', 'shall',
        'can', 'need', 'dare', 'ought', 'used', 'to', 'of', 'in',
        'for', 'on', 'with', 'at', 'by', 'from', 'as', 'into', 'through',
        'during', 'before', 'after', 'above', 'below', 'between', 'under',
        'and', 'but', 'or', 'nor', 'so', 'yet', 'both', 'either', 'neither',
        'not', 'only', 'own', 'same', 'than', 'too', 'very', 'just',
        'also', 'now', 'here', 'there', 'where', 'when', 'why', 'how',
        'all', 'each', 'every', 'both', 'few', 'more', 'most', 'other',
        'some', 'such', 'no', 'any', 'this', 'that', 'these', 'those',
    }
    
    cs_words -= stopwords
    source_words -= stopwords
    
    # Calculate coverage
    if not source_words:
        return {"source_coverage": 0.0, "topic_coverage": 0.0}
    
    overlap = len(cs_words & source_words)
    source_coverage = overlap / len(source_words)
    
    # Topic coverage
    topic_coverage = 0.0
    if expected_topics:
        covered = sum(1 for topic in expected_topics 
                      if topic.lower() in cs_text.lower())
        topic_coverage = covered / len(expected_topics)
    
    return {
        "source_coverage": source_coverage,
        "topic_coverage": topic_coverage,
        "vocabulary_overlap": overlap,
        "unique_cs_terms": len(cs_words - source_words),
    }


def _extract_text_from_content(content: Dict[str, Any]) -> str:
    """Extract plain text from cheat sheet content."""
    texts = []
    
    # Title and overview
    if "title" in content:
        texts.append(content["title"])
    if "overview" in content:
        texts.append(content["overview"])
    
    # Concepts
    for concept in content.get("concepts", []):
        if isinstance(concept, dict):
            texts.append(concept.get("name", ""))
            texts.append(concept.get("definition", ""))
            texts.append(concept.get("intuition", ""))
        else:
            texts.append(str(concept))
    
    # Formulas
    for formula in content.get("formulas", []):
        if isinstance(formula, dict):
            texts.append(formula.get("name", ""))
            texts.append(formula.get("context", ""))
        else:
            texts.append(str(formula))
    
    # Algorithms
    for algo in content.get("algorithms", []):
        if isinstance(algo, dict):
            texts.append(algo.get("name", ""))
            texts.append(algo.get("purpose", ""))
            texts.append(algo.get("key_insight", ""))
    
    # Models
    for model in content.get("models", []):
        if isinstance(model, dict):
            texts.append(model.get("name", ""))
            texts.append(model.get("objective", ""))
    
    return " ".join(filter(None, texts))


def evaluate_cheatsheet(
    content: Dict[str, Any],
    latex_source: str,
    source_chunks: Optional[List[Dict[str, Any]]] = None
) -> CheatSheetMetrics:
    """
    Run complete cheat sheet evaluation.
    
    Args:
        content: Cheat sheet content dictionary
        latex_source: Generated LaTeX source
        source_chunks: Original source chunks (optional)
        
    Returns:
        CheatSheetMetrics
    """
    metrics = CheatSheetMetrics()
    
    # Content counts
    metrics.concept_count = len(content.get("concepts", []))
    metrics.formula_count = len(content.get("formulas", []))
    metrics.algorithm_count = len(content.get("algorithms", []))
    metrics.model_count = len(content.get("models", []))
    
    # LaTeX quality
    latex_eval = evaluate_latex_quality(latex_source)
    metrics.latex_valid = latex_eval["valid"]
    metrics.latex_errors = latex_eval["errors"]
    metrics.compiles_successfully = latex_eval["compiles"]
    
    # Text quality
    concepts = content.get("concepts", [])
    if concepts:
        lengths = []
        for c in concepts:
            if isinstance(c, dict):
                defn = c.get("definition", "")
                lengths.append(len(defn.split()) if defn else 0)
        metrics.avg_concept_length = sum(lengths) / len(lengths) if lengths else 0.0
    
    # Formula complexity (rough measure)
    formulas = content.get("formulas", [])
    if formulas:
        complexities = []
        for f in formulas:
            if isinstance(f, dict):
                latex = f.get("latex", "")
                # Count special characters as complexity indicator
                complexity = len(re.findall(r'[\\{}^_]', latex))
                complexities.append(complexity)
        metrics.avg_formula_complexity = sum(complexities) / len(complexities) if complexities else 0.0
    
    # Diversity - unique topics
    topics = set()
    for item in concepts:
        if isinstance(item, dict) and "topic" in item:
            topics.add(item["topic"])
    for item in formulas:
        if isinstance(item, dict) and "topic" in item:
            topics.add(item["topic"])
    metrics.unique_topics = len(topics)
    
    # Vocabulary richness
    text = _extract_text_from_content(content)
    words = text.lower().split()
    if words:
        unique_words = len(set(words))
        metrics.vocabulary_richness = unique_words / len(words)
    
    # Coverage (if source provided)
    if source_chunks:
        coverage = evaluate_coverage(content, source_chunks)
        metrics.source_coverage = coverage["source_coverage"]
        metrics.topic_coverage = coverage["topic_coverage"]
    
    return metrics


def compare_cheatsheets(
    content1: Dict[str, Any],
    content2: Dict[str, Any]
) -> Dict[str, Any]:
    """
    Compare two cheat sheets.
    
    Args:
        content1: First cheat sheet content
        content2: Second cheat sheet content
        
    Returns:
        Comparison results
    """
    text1 = _extract_text_from_content(content1)
    text2 = _extract_text_from_content(content2)
    
    # Calculate similarity
    words1 = set(text1.lower().split())
    words2 = set(text2.lower().split())
    
    intersection = len(words1 & words2)
    union = len(words1 | words2)
    jaccard = intersection / union if union > 0 else 0.0
    
    # Count differences
    return {
        "jaccard_similarity": jaccard,
        "concepts_diff": len(content1.get("concepts", [])) - len(content2.get("concepts", [])),
        "formulas_diff": len(content1.get("formulas", [])) - len(content2.get("formulas", [])),
        "algorithms_diff": len(content1.get("algorithms", [])) - len(content2.get("algorithms", [])),
        "unique_to_first": len(words1 - words2),
        "unique_to_second": len(words2 - words1),
    }
