#!/usr/bin/env python3
"""
Smart Predictor V5 - Dual-mode prediction.

V4 insight: 18.4% zero predictions = too conservative
Many prompts don't have clear signals (no file mentions)

V5 Strategy:
- Confident mode: When we see file mentions, keywords, or strong sequences
- Fallback mode: Use recency + project context
- Never return empty - always predict something

Target: 35%+ F1
"""

import json
import pickle
import re
import sys
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

PROJECTS_DIR = Path.home() / ".claude" / "projects"
MODEL_CACHE_FILE = Path.home() / ".claude" / "telemetry" / "predictor_model.pkl"


@dataclass
class PredictorModelV5:
    co_occurrence: dict[str, Counter] = field(default_factory=lambda: defaultdict(Counter))
    project_files: dict[str, Counter] = field(default_factory=lambda: defaultdict(Counter))
    name_to_paths: dict[str, set[str]] = field(default_factory=lambda: defaultdict(set))
    sequences: dict[tuple[str, ...], Counter] = field(default_factory=lambda: defaultdict(Counter))
    strong_keywords: dict[str, str] = field(default_factory=dict)
    file_popularity: Counter = field(default_factory=Counter)
    # NEW: Project popularity (which files are most used in each project)
    project_popularity: dict[str, list[str]] = field(default_factory=dict)


def normalize_path(path: str) -> str:
    """Normalize path to a consistent format, stripping user home prefix."""
    path = path.replace("\\", "/")
    # Get home directory dynamically
    home = str(Path.home()).replace("\\", "/")
    # Common path prefixes to strip (covers Windows/Unix variations)
    prefixes = [
        home + "/",
        "/c" + home[2:] + "/",  # Git Bash style: /c/Users/...
    ]
    for prefix in prefixes:
        if path.lower().startswith(prefix.lower()):
            path = path[len(prefix):]
            break
    return path


def get_project_prefix(path: str) -> str:
    parts = path.split("/")
    return parts[0].lower() if parts else ""


def extract_file_mentions(prompt: str) -> set[str]:
    mentions = set()
    for match in re.findall(r'[A-Za-z]:[\\\/][^\s<>"|*?\n]+', prompt):
        mentions.add(normalize_path(match))
    for match in re.findall(r'\b([\w.-]+\.(?:py|js|ts|tsx|jsx|go|rs|java|md|json|html|css|yaml|yml|toml|dart|c|cpp|h))\b', prompt, re.IGNORECASE):
        mentions.add(match.lower())
    return mentions


def train_model_v5(max_sessions: int = 200) -> PredictorModelV5:
    print("Training smart predictor V5 model...")

    model = PredictorModelV5()
    keyword_file_counts: dict[str, Counter] = defaultdict(Counter)
    sessions_processed = 0
    total_turns = 0

    for project_dir in PROJECTS_DIR.iterdir():
        if not project_dir.is_dir():
            continue

        for session_file in project_dir.glob("*.jsonl"):
            if sessions_processed >= max_sessions:
                break
            if session_file.stat().st_size < 1000:
                continue

            try:
                with open(session_file, encoding='utf-8', errors='replace') as f:
                    lines = f.readlines()
            except Exception:
                continue

            session_files = []
            current_prompt = None
            current_files = set()

            for line in lines:
                try:
                    entry = json.loads(line.strip())
                except (json.JSONDecodeError, ValueError):
                    continue

                if entry.get("type") == "user":
                    if current_prompt and current_files:
                        words = set(re.findall(r'[a-zA-Z][a-zA-Z0-9_-]{2,}', current_prompt.lower()))
                        for word in words:
                            for f in current_files:
                                keyword_file_counts[word][f] += 1

                    message = entry.get("message", {})
                    content = message.get("content", "")
                    if isinstance(content, str):
                        current_prompt = content
                    elif isinstance(content, list):
                        current_prompt = " ".join(
                            c.get("text", "") for c in content if c.get("type") == "text"
                        )
                    current_files = set()

                elif entry.get("type") == "assistant":
                    content = entry.get("message", {}).get("content", [])
                    if isinstance(content, list):
                        turn_files = set()
                        for block in content:
                            if block.get("type") == "tool_use":
                                name = block.get("name", "")
                                inp = block.get("input", {})
                                if name in ("Read", "Edit", "Write"):
                                    fp = inp.get("file_path", "")
                                    if fp:
                                        norm = normalize_path(fp)
                                        turn_files.add(norm)
                                        current_files.add(norm)

                                        stem = Path(norm).stem.lower()
                                        model.name_to_paths[stem].add(norm)
                                        fname = Path(norm).name.lower()
                                        model.name_to_paths[fname].add(norm)

                                        prefix = get_project_prefix(norm)
                                        if prefix:
                                            model.project_files[prefix][norm] += 1

                                        model.file_popularity[norm] += 1

                        if turn_files:
                            total_turns += 1
                            for f1 in turn_files:
                                for f2 in turn_files:
                                    if f1 != f2:
                                        model.co_occurrence[f1][f2] += 1
                            session_files.extend(list(turn_files))

            for i in range(1, len(session_files)):
                ctx1 = (session_files[i-1],)
                model.sequences[ctx1][session_files[i]] += 1
                if i > 1:
                    ctx2 = (session_files[i-2], session_files[i-1])
                    model.sequences[ctx2][session_files[i]] += 1

            sessions_processed += 1

        if sessions_processed >= max_sessions:
            break

    # Strong keywords (70%+)
    print("  Building strong keyword associations...")
    for kw, file_counts in keyword_file_counts.items():
        if len(kw) < 4:
            continue
        total = sum(file_counts.values())
        if total < 3:
            continue
        top_file, top_count = file_counts.most_common(1)[0]
        if top_count / total >= 0.7:
            model.strong_keywords[kw] = top_file

    # Project popularity rankings
    print("  Building project popularity rankings...")
    for prefix, files in model.project_files.items():
        model.project_popularity[prefix] = [f for f, _ in files.most_common(20)]

    print(f"  Analyzed {sessions_processed} sessions, {total_turns} turns")
    print(f"  Found {len(model.strong_keywords)} strong keywords")
    print(f"  Indexed {len(model.project_popularity)} projects")

    return model


def predict_files_v5(
    prompt: str,
    model: PredictorModelV5,
    recent_files: list[str] = None,
) -> list[tuple[str, float]]:
    """
    Dual-mode prediction:
    - Confident mode: Use signals
    - Fallback mode: Use recency + project popularity
    """
    recent_files = recent_files or []
    scores: dict[str, float] = defaultdict(float)
    confident = False  # Track if we have high-confidence signals

    # Get project context from recent files
    recent_project = None
    if recent_files:
        recent_project = get_project_prefix(recent_files[-1])

    # === CONFIDENT MODE SIGNALS ===

    # Signal 1: File mentions (strongest)
    mentions = extract_file_mentions(prompt)
    for mention in mentions:
        confident = True  # We have explicit file references
        if "/" in mention:
            scores[mention] += 25.0

        stem = Path(mention).stem.lower()
        if stem in model.name_to_paths:
            for path in model.name_to_paths[stem]:
                scores[path] += 15.0

        name = Path(mention).name.lower()
        if name in model.name_to_paths:
            for path in model.name_to_paths[name]:
                scores[path] += 18.0

    # Signal 2: Strong keywords
    words = set(re.findall(r'[a-zA-Z][a-zA-Z0-9_-]{2,}', prompt.lower()))
    for word in words:
        if word in model.strong_keywords:
            confident = True
            scores[model.strong_keywords[word]] += 10.0

    # Signal 3: Sequence prediction
    if recent_files:
        ctx1 = (recent_files[-1],)
        if ctx1 in model.sequences:
            top_seqs = model.sequences[ctx1].most_common(5)
            if top_seqs and top_seqs[0][1] >= 3:  # Strong sequence
                confident = True
            for fpath, count in top_seqs:
                if count >= 2:
                    scores[fpath] += min(8.0, count * 0.6)

        if len(recent_files) >= 2:
            ctx2 = (recent_files[-2], recent_files[-1])
            if ctx2 in model.sequences:
                top_seqs = model.sequences[ctx2].most_common(3)
                if top_seqs and top_seqs[0][1] >= 2:
                    confident = True
                for fpath, count in top_seqs:
                    scores[fpath] += min(10.0, count * 1.0)

    # === FALLBACK MODE (always apply, but weighted less) ===

    # Recency boost (strong - this is the most reliable fallback)
    for i, recent in enumerate(reversed(recent_files[-5:])):
        boost = 6.0 / (i + 1)  # 6, 3, 2, 1.5, 1.2
        scores[recent] += boost

    # Co-occurrence with very recent files
    if recent_files:
        for recent in recent_files[-3:]:
            if recent in model.co_occurrence:
                for related, count in model.co_occurrence[recent].most_common(5):
                    if count >= 2:
                        scores[related] += min(4.0, count * 0.4)

    # Project popularity fallback (when we have no signals)
    if not confident and recent_project and recent_project in model.project_popularity:
        for i, popular_file in enumerate(model.project_popularity[recent_project][:5]):
            scores[popular_file] += 3.0 / (i + 1)  # 3, 1.5, 1, 0.75, 0.6

    # Project isolation
    if recent_project:
        for f in list(scores.keys()):
            if get_project_prefix(f) != recent_project:
                scores[f] *= 0.15  # Heavy penalty

    # Sort and select
    sorted_scores = sorted(scores.items(), key=lambda x: -x[1])

    if not sorted_scores:
        return []

    # Dynamic selection based on confidence
    if confident:
        # Be more selective when confident
        top_score = sorted_scores[0][1]
        results = [(f, s) for f, s in sorted_scores if s >= top_score * 0.4][:3]
    else:
        # Fallback: return top 2 based on recency/popularity
        results = sorted_scores[:2]

    return results


def benchmark_predictor_v5(model: PredictorModelV5, max_turns: int = 1500):
    print("\n" + "=" * 60)
    print("SMART PREDICTOR V5 BENCHMARK")
    print("=" * 60)

    test_data = []

    for project_dir in PROJECTS_DIR.iterdir():
        if not project_dir.is_dir():
            continue

        for session_file in project_dir.glob("*.jsonl"):
            if len(test_data) >= max_turns:
                break

            try:
                with open(session_file, encoding='utf-8', errors='replace') as f:
                    lines = f.readlines()
            except Exception:
                continue

            current_prompt = None
            recent_files = []
            current_files = set()

            for line in lines:
                try:
                    entry = json.loads(line.strip())
                except (json.JSONDecodeError, ValueError):
                    continue

                if entry.get("type") == "user":
                    if current_prompt and current_files:
                        test_data.append((current_prompt, list(recent_files), current_files))
                        recent_files.extend(list(current_files))
                        recent_files = recent_files[-20:]

                    message = entry.get("message", {})
                    content = message.get("content", "")
                    if isinstance(content, str):
                        current_prompt = content
                    elif isinstance(content, list):
                        current_prompt = " ".join(
                            c.get("text", "") for c in content if c.get("type") == "text"
                        )
                    current_files = set()

                elif entry.get("type") == "assistant":
                    content = entry.get("message", {}).get("content", [])
                    if isinstance(content, list):
                        for block in content:
                            if block.get("type") == "tool_use":
                                name = block.get("name", "")
                                inp = block.get("input", {})
                                if name in ("Read", "Edit", "Write"):
                                    fp = inp.get("file_path", "")
                                    if fp:
                                        current_files.add(normalize_path(fp))

    print(f"\nLoaded {len(test_data)} test turns")

    precisions = []
    recalls = []
    pred_counts = []
    actual_counts = []
    hits_at_1 = 0
    confident_count = 0
    fallback_count = 0

    for i, (prompt, recent, actual) in enumerate(test_data):
        if not actual or prompt.startswith("<"):
            continue

        predictions = predict_files_v5(prompt, model, recent)
        predicted = [p for p, _ in predictions]
        predicted_set = set(predicted)

        # Track if this was confident or fallback
        mentions = extract_file_mentions(prompt)
        words = set(re.findall(r'[a-zA-Z][a-zA-Z0-9_-]{2,}', prompt.lower()))
        has_strong_kw = any(w in model.strong_keywords for w in words)
        if mentions or has_strong_kw:
            confident_count += 1
        else:
            fallback_count += 1

        pred_counts.append(len(predicted))
        actual_counts.append(len(actual))

        if not predicted:
            precisions.append(0)
            recalls.append(0)
            continue

        intersection = predicted_set & actual
        precision = len(intersection) / len(predicted_set)
        recall = len(intersection) / len(actual)

        precisions.append(precision)
        recalls.append(recall)

        if predicted[0] in actual:
            hits_at_1 += 1

        if (i + 1) % 200 == 0:
            avg_p = sum(precisions[-200:]) / max(1, len(precisions[-200:]))
            avg_r = sum(recalls[-200:]) / max(1, len(recalls[-200:]))
            print(f"  {i+1}/{len(test_data)} - P: {avg_p:.1%}, R: {avg_r:.1%}")

    n = len(precisions)
    avg_precision = sum(precisions) / n if precisions else 0
    avg_recall = sum(recalls) / n if recalls else 0
    f1 = 2 * avg_precision * avg_recall / (avg_precision + avg_recall) if (avg_precision + avg_recall) > 0 else 0

    print("\n" + "=" * 60)
    print("SMART PREDICTOR V5 RESULTS")
    print("=" * 60)
    print("\n[ACCURACY]")
    print(f"   Precision: {avg_precision:.1%}")
    print(f"   Recall:    {avg_recall:.1%}")
    print(f"   F1 Score:  {f1:.1%}")

    print("\n[MODE ANALYSIS]")
    print(f"   Confident predictions: {confident_count} ({confident_count/n:.1%})")
    print(f"   Fallback predictions:  {fallback_count} ({fallback_count/n:.1%})")
    print(f"   Avg predicted: {sum(pred_counts)/len(pred_counts):.2f}")
    print(f"   Avg actual:    {sum(actual_counts)/len(actual_counts):.2f}")
    print(f"   Hit@1: {hits_at_1/n:.1%}")

    print("\n[PROGRESSION]")
    print("   Keyword:  7.4%")
    print("   V1:      17.5%")
    print("   V2:      16.9%")
    print("   V3:      26.1%")
    print("   V4:      27.3%")
    print(f"   V5:      {f1:.1%}")

    if f1 >= 0.4:
        grade = "A - Excellent"
    elif f1 >= 0.35:
        grade = "B+ - Very Good"
    elif f1 >= 0.3:
        grade = "B - Good"
    elif f1 >= 0.25:
        grade = "B- - Solid"
    else:
        grade = "C - Acceptable"
    print(f"   Grade: {grade}")

    # Calculate improvement over baseline
    baseline = 0.074
    improvement = (f1 / baseline - 1) * 100
    print(f"\n   Total improvement: +{improvement:.0f}% over keyword baseline")

    print("=" * 60)
    return f1


def save_model(model: PredictorModelV5, path: Path = MODEL_CACHE_FILE):
    """Save trained model to disk."""
    path.parent.mkdir(parents=True, exist_ok=True)
    try:
        with open(path, 'wb') as f:
            pickle.dump(model, f)
    except Exception as e:
        print(f"Failed to save predictor model: {e}", file=sys.stderr)


def load_model(path: Path = MODEL_CACHE_FILE) -> PredictorModelV5 | None:
    """Load trained model from disk."""
    if not path.exists():
        return None
    try:
        with open(path, 'rb') as f:
            return pickle.load(f)
    except Exception as e:
        print(f"Failed to load predictor model: {e}", file=sys.stderr)
        return None


class FilePredictor:
    """
    Production wrapper for the predictor model.

    Lazily loads the trained model on first use and caches it in memory.
    """
    def __init__(self):
        self._model: PredictorModelV5 | None = None

    def _ensure_model(self):
        """Lazy load model on first use."""
        if self._model is not None:
            return

        # Try to load cached model
        self._model = load_model()

        # If no model exists, train one and cache it
        if self._model is None:
            print("[predictor] No cached model found, training new model...", file=sys.stderr)
            self._model = train_model_v5(max_sessions=200)
            save_model(self._model)
            print("[predictor] Model trained and cached", file=sys.stderr)

    def predict(self, recent_files: list[str], top_k: int = 3) -> list[tuple[str, float]]:
        """
        Predict which files are likely to be needed next.

        Args:
            recent_files: List of recently accessed file paths
            top_k: Number of predictions to return

        Returns:
            List of (file_path, probability) tuples
        """
        self._ensure_model()
        if self._model is None:
            return []

        # Use empty prompt since we're doing sequence-based prediction
        predictions = predict_files_v5("", self._model, recent_files)
        return predictions[:top_k]

    def retrain(self, max_sessions: int = 200):
        """Retrain the model from scratch and save it."""
        print("[predictor] Retraining model...", file=sys.stderr)
        self._model = train_model_v5(max_sessions=max_sessions)
        save_model(self._model)
        print("[predictor] Model retrained and cached", file=sys.stderr)


def main():
    model = train_model_v5(max_sessions=200)
    save_model(model)
    print("\nModel saved to:", MODEL_CACHE_FILE)
    benchmark_predictor_v5(model, max_turns=1500)


if __name__ == "__main__":
    main()
