"""
Category cache system for consistent categorization across mailboxes.

Stores discovered categories and provides semantic matching to snap
new discoveries to existing categories for cross-mailbox consistency.
"""

import json
import logging
import numpy as np
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from datetime import datetime

logger = logging.getLogger(__name__)


class CategoryCache:
    """
    Manages cached categories for consistent email classification.

    Features:
    - Persistent storage of discovered categories
    - Semantic matching to snap new categories to cached ones
    - Usage tracking for category popularity
    - Support for mailbox-specific overrides
    """

    def __init__(self, cache_path: str = "src/models/category_cache.json", embedding_model=None, llm_provider=None):
        self.cache_path = Path(cache_path)
        self.cache: Dict[str, dict] = {}
        self.embedding_model = embedding_model
        self.llm_provider = llm_provider
        self.load()

    def load(self) -> None:
        """Load category cache from disk."""
        if self.cache_path.exists():
            try:
                with open(self.cache_path, 'r') as f:
                    data = json.load(f)
                    self.cache = data.get('categories', {})
                logger.info(f"Loaded {len(self.cache)} cached categories from {self.cache_path}")
            except Exception as e:
                logger.error(f"Failed to load category cache: {e}")
                self.cache = {}
        else:
            logger.info("No category cache found, starting fresh")
            self.cache = {}

    def save(self) -> None:
        """Save category cache to disk."""
        try:
            self.cache_path.parent.mkdir(parents=True, exist_ok=True)
            data = {
                'version': '1.0',
                'updated': datetime.now().isoformat(),
                'categories': self.cache
            }
            with open(self.cache_path, 'w') as f:
                json.dump(data, f, indent=2)
            logger.info(f"Saved {len(self.cache)} categories to cache")
        except Exception as e:
            logger.error(f"Failed to save category cache: {e}")

    def get_cached_categories(self) -> Dict[str, str]:
        """Get all cached categories as {name: description}."""
        return {name: info['description'] for name, info in self.cache.items()}

    def snap_to_cache(
        self,
        discovered: Dict[str, str],
        similarity_threshold: float = 0.7,
        allow_new: bool = True,
        max_new: int = 3,
        llm_review_threshold: float = 0.5
    ) -> Tuple[Dict[str, str], Dict[str, str], List[Tuple[str, str]]]:
        """
        Snap discovered categories to cached ones using multi-stage matching.

        Matching Strategy:
        1. Stage 1: Exact name match (score=1.0) → instant snap
        2. Stage 2: Embedding similarity ≥ threshold → confident snap
        3. Stage 3: Similarity < threshold but ≥ llm_review_threshold → LLM review
        4. Stage 4: Very low similarity → auto-accept as new (if allowed)
        5. Stage 5: Exceeded max_new → force snap to closest

        Args:
            discovered: Newly discovered categories {name: description}
            similarity_threshold: Confident match threshold (default: 0.7)
            allow_new: Whether to allow new categories
            max_new: Maximum new categories per run
            llm_review_threshold: Min score to trigger LLM review (default: 0.5)

        Returns:
            (snapped_categories, mapping, cache_worthy) where:
            - snapped_categories: Final category set (from cache + new)
            - mapping: {discovered_name: final_name} for all discovered
            - cache_worthy: List of (name, desc) for categories to add to persistent cache
        """
        if not self.cache:
            logger.info("Empty cache, using all discovered categories")
            # First run - all categories can be cache-worthy
            return discovered, {name: name for name in discovered}, list(discovered.items())

        snapped = {}
        mapping = {}
        new_categories = []
        ambiguous_cases = []  # For LLM review

        cached_cats = self.get_cached_categories()

        # Stage 1-4: Classify each discovered category
        for disc_name, disc_desc in discovered.items():
            best_match, best_score = self._find_best_match(
                disc_name, disc_desc, cached_cats
            )

            if best_score >= similarity_threshold:
                # Stage 2: Confident match → snap immediately
                mapping[disc_name] = best_match
                if best_match not in snapped:
                    snapped[best_match] = cached_cats[best_match]
                logger.debug(f"Confident snap '{disc_name}' → '{best_match}' (similarity: {best_score:.2f})")

            elif best_score >= llm_review_threshold:
                # Stage 3: Ambiguous → save for LLM review
                ambiguous_cases.append((disc_name, disc_desc, best_match, best_score))
                logger.info(f"Ambiguous: '{disc_name}' ~ '{best_match}' (similarity: {best_score:.2f}) → LLM review")

            else:
                # Stage 4: Very low similarity → likely new category
                if allow_new and len(new_categories) < max_new:
                    new_categories.append((disc_name, disc_desc))
                    mapping[disc_name] = disc_name
                    snapped[disc_name] = disc_desc
                    logger.info(f"New category: '{disc_name}' (low similarity: {best_score:.2f})")
                else:
                    # Stage 5: Can't add new → force snap
                    ambiguous_cases.append((disc_name, disc_desc, best_match, best_score))
                    logger.warning(f"Force review: '{disc_name}' (max_new exceeded, score: {best_score:.2f})")

        # LLM Review for ambiguous cases
        cache_worthy_additions = []  # Track categories to add to persistent cache

        if ambiguous_cases and self.llm_provider:
            logger.info(f"Requesting LLM review for {len(ambiguous_cases)} ambiguous cases...")
            llm_decisions = self._llm_review_ambiguous(ambiguous_cases, cached_cats, allow_new, len(new_categories), max_new)

            for disc_name, decision in llm_decisions.items():
                if decision['action'] == 'snap':
                    target = decision['target']
                    mapping[disc_name] = target
                    if target not in snapped:
                        snapped[target] = cached_cats[target]
                    reasoning = decision.get('reasoning', 'similar to existing')
                    logger.info(f"LLM snap: '{disc_name}' → '{target}' ({reasoning})")

                elif decision['action'] == 'new':
                    # Find original description
                    disc_desc = next(desc for name, desc, _, _ in ambiguous_cases if name == disc_name)
                    new_categories.append((disc_name, disc_desc))
                    mapping[disc_name] = disc_name
                    snapped[disc_name] = disc_desc

                    # Check if LLM recommends adding to cache
                    is_cache_worthy = decision.get('cache_worthy', False)
                    reasoning = decision.get('reasoning', 'new category')

                    if is_cache_worthy:
                        cache_worthy_additions.append((disc_name, disc_desc))
                        logger.info(f"LLM approved new + CACHE: '{disc_name}' ({reasoning})")
                    else:
                        logger.info(f"LLM approved new (run-only): '{disc_name}' ({reasoning})")

        elif ambiguous_cases:
            # No LLM available → use heuristic fallback
            logger.warning(f"No LLM for review, using heuristic for {len(ambiguous_cases)} ambiguous cases")
            for disc_name, disc_desc, best_match, best_score in ambiguous_cases:
                if best_score >= 0.6:  # Medium-high score → snap
                    mapping[disc_name] = best_match
                    if best_match not in snapped:
                        snapped[best_match] = cached_cats[best_match]
                elif allow_new and len(new_categories) < max_new:  # Low score + space → new
                    new_categories.append((disc_name, disc_desc))
                    mapping[disc_name] = disc_name
                    snapped[disc_name] = disc_desc
                else:  # Force snap
                    mapping[disc_name] = best_match
                    if best_match not in snapped:
                        snapped[best_match] = cached_cats[best_match]

        logger.info(f"Snapping result: {len(snapped)} final categories ({len(new_categories)} new, {len(cache_worthy_additions)} cache-worthy)")
        return snapped, mapping, cache_worthy_additions

    def _find_best_match(
        self,
        name: str,
        description: str,
        cached: Dict[str, str]
    ) -> Tuple[Optional[str], float]:
        """
        Find best matching cached category using embedding-based semantic similarity.

        Uses embeddings + cosine similarity for true semantic matching.
        Returns (best_category_name, similarity_score).
        """
        if not cached:
            return None, 0.0

        # Exact name match always wins
        name_lower = name.lower()
        for cached_name in cached.keys():
            if name_lower == cached_name.lower():
                return cached_name, 1.0

        # Use embeddings if available
        if self.embedding_model:
            try:
                # Combine name and description for richer semantic representation
                discovered_text = f"{name}: {description}"
                response = self.embedding_model.embeddings(
                    model='all-minilm:l6-v2',
                    prompt=discovered_text
                )
                discovered_emb = np.array(response['embedding'], dtype=np.float32)

                best_match = None
                best_score = 0.0

                for cached_name, cached_desc in cached.items():
                    cached_text = f"{cached_name}: {cached_desc}"
                    response = self.embedding_model.embeddings(
                        model='all-minilm:l6-v2',
                        prompt=cached_text
                    )
                    cached_emb = np.array(response['embedding'], dtype=np.float32)

                    # Cosine similarity
                    similarity = self._cosine_similarity(discovered_emb, cached_emb)

                    if similarity > best_score:
                        best_score = similarity
                        best_match = cached_name

                return best_match, best_score
            except Exception as e:
                logger.warning(f"Embedding-based matching failed: {e}, falling back to partial name match")
                # Fall through to partial matching below

        # Fallback to partial name matching if no embeddings
        best_match = None
        best_score = 0.0

        for cached_name in cached.keys():
            if name_lower in cached_name.lower() or cached_name.lower() in name_lower:
                score = 0.8
                if score > best_score:
                    best_score = score
                    best_match = cached_name

        return best_match if best_match else list(cached.keys())[0], best_score

    def _cosine_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
        """Calculate cosine similarity between two vectors."""
        dot_product = np.dot(vec1, vec2)
        norm_a = np.linalg.norm(vec1)
        norm_b = np.linalg.norm(vec2)

        if norm_a == 0 or norm_b == 0:
            return 0.0

        return float(dot_product / (norm_a * norm_b))

    def _llm_review_ambiguous(
        self,
        ambiguous_cases: List[Tuple[str, str, str, float]],
        cached_cats: Dict[str, str],
        allow_new: bool,
        current_new_count: int,
        max_new: int
    ) -> Dict[str, Dict]:
        """
        Use LLM to review ambiguous category matches.

        Args:
            ambiguous_cases: List of (disc_name, disc_desc, best_match, score)
            cached_cats: All cached categories
            allow_new: Whether new categories are allowed
            current_new_count: How many new categories already accepted
            max_new: Maximum new categories allowed

        Returns:
            Dict of {disc_name: {'action': 'snap'|'new', 'target': category_name}}
        """
        import json
        import re

        # Build ambiguous cases summary
        cases_summary = "\n".join([
            f"- \"{name}\": {desc}\n  Best cache match: \"{match}\" (similarity: {score:.2f})"
            for name, desc, match, score in ambiguous_cases
        ])

        # Build cached categories list
        cached_summary = "\n".join([
            f"- \"{name}\": {desc}"
            for name, desc in cached_cats.items()
        ])

        remaining_slots = max(0, max_new - current_new_count)

        prompt = f"""<no_think>You are reviewing ambiguous category matches for an email classification system.

CACHED CATEGORIES (established):
{cached_summary}

AMBIGUOUS CASES (need decision):
{cases_summary}

CONTEXT:
- Embedding similarity scores: 1.0=identical, 0.7+=very similar, 0.5-0.7=ambiguous, <0.5=different
- Allow new categories: {allow_new}
- Remaining slots for new categories: {remaining_slots}/{max_new}

TASK:
For each ambiguous case, decide:
1. "snap" - If semantically similar enough to cached category
2. "new" - If genuinely distinct and worth adding for THIS RUN

ADDITIONALLY:
For "new" categories, decide if they should be added to the CACHE for future mailboxes:
- "cache_worthy": true - High-quality, reusable category (e.g., "Customer Support", "Sales")
- "cache_worthy": false - Mailbox-specific, not broadly useful (e.g., "Project X Updates")

GUIDELINES:
- PREFER snapping to maintain consistency across mailboxes
- Only approve "new" if category serves a clearly distinct purpose
- Be VERY selective with cache_worthy - only approve universally useful categories
- Consider: Would this category be useful across different users' mailboxes?
- Mailbox-specific categories can be "new" without being cache_worthy

Return JSON:
{{
  "CategoryName": {{
    "action": "snap"|"new",
    "target": "CachedCategoryName",
    "cache_worthy": true|false,
    "reasoning": "brief explanation"
  }},
  ...
}}

For "snap": target = cached category to snap to, cache_worthy = false (not applicable)
For "new": target = same as CategoryName, cache_worthy = true/false based on reusability

JSON:
"""

        try:
            response = self.llm_provider.complete(prompt, temperature=0.1, max_tokens=2000)

            # Parse JSON
            cleaned = re.sub(r'<think>.*?</think>', '', response, flags=re.DOTALL)
            json_match = re.search(r'\{.*\}', cleaned, re.DOTALL)

            if json_match:
                decisions = json.loads(json_match.group())
                logger.info(f"LLM reviewed {len(decisions)} categories")
                return decisions
            else:
                logger.error("LLM review: no JSON found in response")
                return {}

        except Exception as e:
            logger.error(f"LLM review failed: {e}")
            return {}

    def update_cache(
        self,
        categories: Dict[str, str],
        usage_count: Optional[Dict[str, int]] = None
    ) -> None:
        """
        Update cache with new/refined categories.

        Args:
            categories: Categories to add/update {name: description}
            usage_count: Optional email counts per category
        """
        for name, desc in categories.items():
            if name in self.cache:
                # Update existing
                self.cache[name]['description'] = desc
                self.cache[name]['last_seen'] = datetime.now().isoformat()
                if usage_count and name in usage_count:
                    self.cache[name]['total_emails'] = self.cache[name].get('total_emails', 0) + usage_count[name]
            else:
                # Add new
                self.cache[name] = {
                    'description': desc,
                    'created': datetime.now().isoformat(),
                    'last_seen': datetime.now().isoformat(),
                    'total_emails': usage_count.get(name, 0) if usage_count else 0
                }

        self.save()
        logger.info(f"Updated cache with {len(categories)} categories")

    def get_stats(self) -> Dict:
        """Get cache statistics."""
        if not self.cache:
            return {'total_categories': 0}

        total_emails = sum(info.get('total_emails', 0) for info in self.cache.values())
        sorted_by_usage = sorted(
            self.cache.items(),
            key=lambda x: x[1].get('total_emails', 0),
            reverse=True
        )

        return {
            'total_categories': len(self.cache),
            'total_emails_classified': total_emails,
            'top_categories': [
                (name, info.get('total_emails', 0))
                for name, info in sorted_by_usage[:10]
            ]
        }
