"""RAG Knowledge Base service for storing and retrieving bug-fix examples."""

from __future__ import annotations

import ast
import difflib
import hashlib
import json
import re
import sqlite3
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import numpy as np

from ai_orchestrator.config.settings import Settings
from ai_orchestrator.schemas.rag_models import (
    BugFixExample,
    DatasetStatistics,
    RAGRetrievalResult,
    RetrievedExample,
    RetrievalQuery,
)
from ai_orchestrator.schemas.task_state import TaskState
from ai_orchestrator.services.embedding_service import EmbeddingService


@dataclass(slots=True, frozen=True)
class ParsedDiffHunk:
    line_number: int | None
    before_text: str
    after_text: str
    diff_text: str


@dataclass(slots=True, frozen=True)
class ParsedDiffFile:
    file_path: str
    before_text: str
    after_text: str
    diff_text: str
    line_number: int | None
    hunks: list[ParsedDiffHunk]


class RAGKnowledgeBase:
    """Knowledge base for bug-fix examples with hybrid retrieval and snippet indexing."""

    ERROR_CLASS_HINTS = {
        "attributeerror": {
            "display": "AttributeError",
            "category": "attribute-error",
            "explicit_patterns": [r"\battributeerror\b", r"has no attribute", r"\bnonetype\b"],
            "code_patterns": [],
            "required_terms": ["AttributeError", "has no attribute"],
            "preferred_datasets": ["AIIQ", "StackOverflow", "BugsInPy", "SWE-bench"],
        },
        "keyerror": {
            "display": "KeyError",
            "category": "key-error",
            "explicit_patterns": [r"\bkeyerror\b", r"missing key", r"dictionary key"],
            "code_patterns": [r"\[[\"'][A-Za-z_][^\"']*[\"']\]"],
            "required_terms": ["KeyError", "missing key"],
            "preferred_datasets": ["AIIQ", "StackOverflow", "BugsInPy", "SWE-bench"],
        },
        "unicodedecodeerror": {
            "display": "UnicodeDecodeError",
            "category": "encoding-issue",
            "explicit_patterns": [
                r"\bunicodedecodeerror\b",
                r"codec.+decode",
                r"wrong encoding",
                r"decode error",
            ],
            "code_patterns": [r"\.decode\(", r"\.encode\(", r"encoding\s*="],
            "required_terms": ["UnicodeDecodeError", "encoding", "decode"],
            "preferred_datasets": ["AIIQ", "StackOverflow", "BugsInPy", "SWE-bench"],
        },
    }
    CATEGORY_TO_ERROR_CLASS = {
        "attribute-error": "attributeerror",
        "key-error": "keyerror",
        "encoding-issue": "unicodedecodeerror",
    }
    QUERY_STOPWORDS = {
        "the",
        "and",
        "for",
        "with",
        "this",
        "that",
        "from",
        "into",
        "when",
        "where",
        "while",
        "return",
        "value",
        "values",
        "items",
        "python",
        "file",
        "code",
        "function",
        "please",
        "review",
        "found",
        "read",
        "your",
        "task",
        "goal",
    }
    TASK_CATEGORY_HINTS = {
        "caller_contract_key_access_without_guard": "key-error",
        "caller_contract_optional_attribute_access": "attribute-error",
        "caller_contract_text_read_without_encoding": "encoding-issue",
        "PotentialKeyError": "key-error",
        "PotentialAttributeError": "attribute-error",
        "PotentialUnicodeDecodeError": "encoding-issue",
    }

    def __init__(self, settings: Settings) -> None:
        self.settings = settings
        self.db_path = settings.data_dir / "rag_knowledge_base.db"
        self._ensure_db_schema()

    def _ensure_db_schema(self) -> None:
        self.db_path.parent.mkdir(parents=True, exist_ok=True)

        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS bug_fix_examples (
                    id TEXT PRIMARY KEY,
                    dataset TEXT NOT NULL,
                    repository TEXT NOT NULL,
                    example_kind TEXT DEFAULT 'full_example',
                    parent_example_id TEXT,
                    snippet_role TEXT,
                    snippet_index INTEGER,
                    commit_hash TEXT,
                    issue_url TEXT,
                    bug_description TEXT NOT NULL,
                    bug_summary TEXT NOT NULL,
                    bug_category TEXT NOT NULL,
                    file_path TEXT NOT NULL,
                    function_name TEXT,
                    line_number INTEGER,
                    buggy_code TEXT NOT NULL,
                    fixed_code TEXT NOT NULL,
                    diff TEXT NOT NULL,
                    language TEXT DEFAULT 'Python',
                    severity TEXT DEFAULT 'medium',
                    complexity TEXT DEFAULT 'simple',
                    root_cause TEXT,
                    symptoms TEXT,
                    related_patterns TEXT,
                    task_mode TEXT,
                    validation_failures TEXT,
                    source_issue_types TEXT,
                    retrieval_tags TEXT,
                    embedding BLOB,
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                    indexed_at TEXT
                )
                """
            )

            self._ensure_column(conn, "bug_fix_examples", "example_kind", "TEXT DEFAULT 'full_example'")
            self._ensure_column(conn, "bug_fix_examples", "parent_example_id", "TEXT")
            self._ensure_column(conn, "bug_fix_examples", "snippet_role", "TEXT")
            self._ensure_column(conn, "bug_fix_examples", "snippet_index", "INTEGER")
            self._ensure_column(conn, "bug_fix_examples", "task_mode", "TEXT")
            self._ensure_column(conn, "bug_fix_examples", "validation_failures", "TEXT")
            self._ensure_column(conn, "bug_fix_examples", "source_issue_types", "TEXT")
            self._ensure_column(conn, "bug_fix_examples", "retrieval_tags", "TEXT")

            conn.execute("CREATE INDEX IF NOT EXISTS idx_category ON bug_fix_examples(bug_category)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_dataset ON bug_fix_examples(dataset)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_severity ON bug_fix_examples(severity)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_repository ON bug_fix_examples(repository)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_example_kind ON bug_fix_examples(example_kind)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_parent_example ON bug_fix_examples(parent_example_id)")

            conn.execute(
                """
                CREATE VIRTUAL TABLE IF NOT EXISTS bug_fix_fts USING fts5(
                    id, bug_summary, bug_description, root_cause,
                    content=bug_fix_examples,
                    content_rowid=rowid
                )
                """
            )
            conn.execute(
                """
                CREATE TRIGGER IF NOT EXISTS bug_fix_ai AFTER INSERT ON bug_fix_examples BEGIN
                    INSERT INTO bug_fix_fts(rowid, id, bug_summary, bug_description, root_cause)
                    VALUES (new.rowid, new.id, new.bug_summary, new.bug_description, new.root_cause);
                END
                """
            )
            conn.execute(
                """
                CREATE TRIGGER IF NOT EXISTS bug_fix_ad AFTER DELETE ON bug_fix_examples BEGIN
                    DELETE FROM bug_fix_fts WHERE rowid = old.rowid;
                END
                """
            )
            conn.execute(
                """
                CREATE TRIGGER IF NOT EXISTS bug_fix_au AFTER UPDATE ON bug_fix_examples BEGIN
                    UPDATE bug_fix_fts
                    SET bug_summary = new.bug_summary,
                        bug_description = new.bug_description,
                        root_cause = new.root_cause
                    WHERE rowid = new.rowid;
                END
                """
            )
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS rag_retrieval_telemetry (
                    event_id TEXT PRIMARY KEY,
                    task_id TEXT,
                    stage TEXT NOT NULL,
                    file_path TEXT,
                    function_name TEXT,
                    query_digest TEXT NOT NULL,
                    query_text TEXT NOT NULL,
                    example_ids TEXT NOT NULL,
                    example_scores TEXT NOT NULL,
                    outcome TEXT,
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
                """
            )
            conn.execute(
                "CREATE INDEX IF NOT EXISTS idx_rag_retrieval_task ON rag_retrieval_telemetry(task_id)"
            )
            conn.commit()

    @staticmethod
    def _ensure_column(conn: sqlite3.Connection, table_name: str, column_name: str, column_spec: str) -> None:
        existing_columns = {row[1] for row in conn.execute(f"PRAGMA table_info({table_name})")}
        if column_name in existing_columns:
            return
        conn.execute(f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_spec}")

    def add_example(self, example: BugFixExample) -> None:
        self.add_examples_batch([example])

    def add_examples_batch(self, examples: list[BugFixExample]) -> None:
        if not examples:
            return
        with sqlite3.connect(self.db_path) as conn:
            conn.executemany(
                """
                INSERT OR REPLACE INTO bug_fix_examples
                (id, dataset, repository, example_kind, parent_example_id, snippet_role, snippet_index,
                 commit_hash, issue_url, bug_description, bug_summary, bug_category, file_path,
                 function_name, line_number, buggy_code, fixed_code, diff, language, severity,
                 complexity, root_cause, symptoms, related_patterns, task_mode, validation_failures,
                 source_issue_types, retrieval_tags, embedding, indexed_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                [self._prepare_example_record(example) for example in examples],
            )
            conn.commit()

    def ingest_example_family(self, example: BugFixExample, embedding_service: EmbeddingService) -> list[BugFixExample]:
        examples = self._family_examples(example, embedding_service)
        self.add_examples_batch(examples)
        return examples

    def ingest_task_result(
        self,
        *,
        task: TaskState,
        final_diff: str,
        embedding_service: EmbeddingService,
    ) -> list[BugFixExample]:
        parsed_files = self._parse_unified_diff(final_diff)
        if not parsed_files:
            return []

        preview_by_path = self._extract_preview_by_path(task)
        learned_examples: list[BugFixExample] = []
        for index, file_diff in enumerate(parsed_files, start=1):
            target_path = task.repo_path / file_diff.file_path
            try:
                fixed_text = target_path.read_text(encoding="utf-8")
            except (OSError, UnicodeDecodeError):
                fixed_text = file_diff.after_text
            preview_text = preview_by_path.get(file_diff.file_path, "")
            buggy_text = preview_text if preview_text and "...<truncated>" not in preview_text else file_diff.before_text
            if not buggy_text and not fixed_text:
                continue

            example = BugFixExample(
                id=self._stable_example_id("taskfix", task.task_id, file_diff.file_path, str(index)),
                dataset="AIIQ",
                repository=task.repo_path.name,
                example_kind="task_fix",
                bug_description=task.goal,
                bug_summary=(task.summary or f"Successful fix in {file_diff.file_path}")[:200],
                bug_category=self._infer_category_from_task(task, file_diff),
                file_path=file_diff.file_path,
                function_name=self._extract_primary_function_name(fixed_text or buggy_text),
                line_number=file_diff.line_number,
                buggy_code=buggy_text or file_diff.before_text,
                fixed_code=fixed_text or file_diff.after_text,
                diff=file_diff.diff_text,
                severity="medium" if task.attempt <= 1 else "high",
                complexity=self._assess_complexity(file_diff.diff_text),
                root_cause=self._task_root_cause(task),
                symptoms=self._task_symptoms(task),
                related_patterns=self._task_patterns(task),
                task_mode=task.task_mode,
                validation_failures=self._task_validation_failures(task),
                source_issue_types=self._task_patterns(task),
                retrieval_tags=self._task_retrieval_tags(task),
            )
            learned_examples.extend(self._family_examples(example, embedding_service))

        if learned_examples:
            self.add_examples_batch(learned_examples)
        return learned_examples

    def get_example(self, example_id: str) -> BugFixExample | None:
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            row = conn.execute("SELECT * FROM bug_fix_examples WHERE id = ?", (example_id,)).fetchone()
        if row is None:
            return None
        return self._row_to_example(row)

    def list_examples(
        self,
        *,
        example_ids: list[str] | None = None,
        dataset: str | None = None,
        repository: str | None = None,
        limit: int | None = None,
    ) -> list[BugFixExample]:
        where_clauses = ["1=1"]
        params: list[Any] = []
        if example_ids:
            placeholders = ",".join("?" for _ in example_ids)
            where_clauses.append(f"id IN ({placeholders})")
            params.extend(example_ids)
        if dataset:
            where_clauses.append("dataset = ?")
            params.append(dataset)
        if repository:
            where_clauses.append("repository = ?")
            params.append(repository)

        query = f"SELECT * FROM bug_fix_examples WHERE {' AND '.join(where_clauses)} ORDER BY id"
        if limit is not None:
            query += " LIMIT ?"
            params.append(limit)

        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            rows = conn.execute(query, params).fetchall()
        examples = [self._row_to_example(row) for row in rows]
        if not example_ids:
            return examples

        order = {example_id: index for index, example_id in enumerate(example_ids)}
        examples.sort(key=lambda example: order.get(example.id, len(order)))
        return examples

    def export_examples(
        self,
        *,
        example_ids: list[str] | None = None,
        dataset: str | None = None,
        repository: str | None = None,
        limit: int | None = None,
    ) -> list[BugFixExample]:
        return self.list_examples(
            example_ids=example_ids,
            dataset=dataset,
            repository=repository,
            limit=limit,
        )

    def import_examples(
        self,
        examples: list[BugFixExample],
        *,
        embedding_service: EmbeddingService | None = None,
    ) -> list[BugFixExample]:
        prepared: list[BugFixExample] = []
        for example in examples:
            item = example
            if item.embedding is None and embedding_service is not None:
                item = item.model_copy(update={"embedding": embedding_service.embed_code(self._example_embedding_text(item))})
            prepared.append(item)
        self.add_examples_batch(prepared)
        return prepared

    def import_examples_file(
        self,
        path: Path,
        *,
        embedding_service: EmbeddingService | None = None,
    ) -> list[BugFixExample]:
        files = self._example_files_from_path(path)
        imported: list[BugFixExample] = []
        for file_path in files:
            imported.extend(
                self.import_examples(
                    self._load_examples_from_file(file_path),
                    embedding_service=embedding_service,
                )
            )
        return imported

    def build_retrieval_query(
        self,
        *,
        query_code: str,
        bug_description: str | None = None,
        file_path: str | None = None,
        function_name: str | None = None,
        max_results: int = 5,
    ) -> RetrievalQuery:
        hints = self._infer_query_preferences(query_code=query_code, bug_description=bug_description)
        return RetrievalQuery(
            query_code=query_code,
            bug_description=bug_description,
            file_path=file_path,
            function_name=function_name,
            max_results=max_results,
            max_per_dataset=self.settings.rag_max_per_dataset,
            error_classes=hints["error_classes"],
            preferred_categories=hints["preferred_categories"],
            preferred_datasets=hints["preferred_datasets"],
            preferred_example_kinds=self._preferred_example_kinds(query_code, function_name),
            required_terms=hints["required_terms"],
        )

    def extract_query_units(self, query_code: str, *, file_path: str | None = None) -> list[dict[str, Any]]:
        units: list[dict[str, Any]] = [
            {
                "code": query_code,
                "file_path": file_path,
                "function_name": self._extract_primary_function_name(query_code),
                "line_number": None,
            }
        ]
        functions = self._parse_functions(query_code)
        for function_name, metadata in functions.items():
            units.append(
                {
                    "code": metadata["code"],
                    "file_path": file_path,
                    "function_name": function_name,
                    "line_number": metadata["line_number"],
                }
            )

        deduped: list[dict[str, Any]] = []
        seen: set[tuple[str, str | None]] = set()
        for unit in units:
            code = str(unit.get("code", "")).strip()
            key = (code, unit.get("function_name"))
            if not code or key in seen:
                continue
            seen.add(key)
            deduped.append(unit)
        return deduped[:4]

    def retrieve_similar(
        self,
        query: RetrievalQuery,
        query_embedding: list[float] | None = None,
    ) -> RAGRetrievalResult:
        query = self._merge_with_inferred_preferences(query)

        where_clauses: list[str] = ["1=1"]
        params: list[Any] = []

        if query.categories:
            placeholders = ",".join("?" for _ in query.categories)
            where_clauses.append(f"bug_category IN ({placeholders})")
            params.extend(query.categories)

        if query.severity:
            where_clauses.append("severity = ?")
            params.append(query.severity)

        if query.exclude_datasets:
            placeholders = ",".join("?" for _ in query.exclude_datasets)
            where_clauses.append(f"dataset NOT IN ({placeholders})")
            params.extend(query.exclude_datasets)

        where_clause = " AND ".join(where_clauses)
        candidate_limit = max(query.max_results * 12, 30)

        vector_examples = (
            self._vector_search(query_embedding, where_clause, params, candidate_limit)
            if query_embedding
            else []
        )
        keyword_examples = self._keyword_search(query, where_clause, params, candidate_limit)
        examples = self._rerank_candidates(query, query_embedding, vector_examples, keyword_examples)
        examples = self._diversify_ranked_examples(examples, query.max_results, query.max_per_dataset)

        all_patterns: set[str] = set()
        for item in examples:
            all_patterns.update(item.example.related_patterns)

        return RAGRetrievalResult(
            query=query.query_code[:200] if query.query_code else "",
            examples=examples[: query.max_results],
            total_found=len(examples),
            suggested_patterns=sorted(all_patterns),
            context_summary=self._build_context_summary(examples),
        )

    def _vector_search(
        self,
        query_embedding: list[float],
        where_clause: str,
        params: list[Any],
        limit: int,
    ) -> list[RetrievedExample]:
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            rows = conn.execute(
                f"SELECT * FROM bug_fix_examples WHERE {where_clause} AND embedding IS NOT NULL",
                params,
            ).fetchall()

        query_vec = np.array(query_embedding)
        examples: list[RetrievedExample] = []
        for row in rows:
            example = self._row_to_example(row)
            if not example.embedding:
                continue
            similarity = self._cosine_similarity(query_vec, np.array(example.embedding))
            examples.append(
                RetrievedExample(
                    example=example,
                    relevance_score=float(similarity),
                    match_reason="vector_similarity",
                )
            )
        examples.sort(key=lambda item: item.relevance_score, reverse=True)
        return examples[:limit]

    def _keyword_search(
        self,
        query: RetrievalQuery,
        where_clause: str,
        params: list[Any],
        limit: int,
    ) -> list[RetrievedExample]:
        search_terms = self._build_search_terms(query)
        if not search_terms:
            return []
        search_query = " OR ".join(self._quote_fts_term(term) for term in search_terms[:20])

        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            try:
                fts_rows = conn.execute(
                    f"""
                    SELECT b.* FROM bug_fix_examples b
                    INNER JOIN bug_fix_fts f ON b.id = f.id
                    WHERE bug_fix_fts MATCH ? AND {where_clause}
                    LIMIT ?
                    """,
                    [search_query, *params, limit],
                ).fetchall()
            except (sqlite3.OperationalError, sqlite3.DatabaseError):
                return self._fallback_keyword_search(query, where_clause, params, limit)

            direct_terms = search_terms[:6]
            direct_rows: list[sqlite3.Row] = []
            if direct_terms or query.file_path:
                placeholders = ",".join("?" for _ in direct_terms) or "''"
                direct_rows = conn.execute(
                    f"""
                    SELECT * FROM bug_fix_examples
                    WHERE {where_clause}
                      AND (
                        function_name IN ({placeholders})
                        OR file_path LIKE ?
                      )
                    LIMIT ?
                    """,
                    [*params, *direct_terms, f"%{query.file_path or ''}%", limit],
                ).fetchall()

        results_by_id: dict[str, RetrievedExample] = {}
        for row in [*fts_rows, *direct_rows]:
            example = self._row_to_example(row)
            score = self._keyword_relevance_score(query, example)
            if score <= 0:
                continue
            current = results_by_id.get(example.id)
            if current is None or score > current.relevance_score:
                results_by_id[example.id] = RetrievedExample(
                    example=example,
                    relevance_score=score,
                    match_reason="keyword_match",
                )
        examples = list(results_by_id.values())
        examples.sort(key=lambda item: item.relevance_score, reverse=True)
        return examples[:limit]

    def _fallback_keyword_search(
        self,
        query: RetrievalQuery,
        where_clause: str,
        params: list[Any],
        limit: int,
    ) -> list[RetrievedExample]:
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            rows = conn.execute(f"SELECT * FROM bug_fix_examples WHERE {where_clause}", params).fetchall()

        examples: list[RetrievedExample] = []
        for row in rows:
            example = self._row_to_example(row)
            score = self._keyword_relevance_score(query, example)
            if score <= 0:
                continue
            examples.append(
                RetrievedExample(
                    example=example,
                    relevance_score=score,
                    match_reason="keyword_fallback",
                )
            )
        examples.sort(key=lambda item: item.relevance_score, reverse=True)
        return examples[:limit]

    def _keyword_relevance_score(self, query: RetrievalQuery, example: BugFixExample) -> float:
        query_terms = {term.lower() for term in self._build_search_terms(query)}
        if not query_terms:
            return 0.0
        example_text = self._example_search_text(example)
        overlap = [term for term in query_terms if term in example_text]
        score = len(overlap) / len(query_terms)

        exact_hits = sum(1 for term in query.required_terms if term.lower() in example_text)
        if exact_hits:
            score += min(
                self.settings.rag_keyword_exact_hit_boost_cap,
                exact_hits * self.settings.rag_keyword_exact_hit_boost_per_term,
            )
        if query.error_classes and self._example_matches_error_classes(example, query.error_classes):
            score += self.settings.rag_keyword_error_class_boost
        if query.function_name and example.function_name == query.function_name:
            score += self.settings.rag_keyword_function_name_boost
        if example.example_kind in query.preferred_example_kinds:
            score += self._example_kind_adjustment(query, example)
        return min(score, 1.0)

    def _merge_with_inferred_preferences(self, query: RetrievalQuery) -> RetrievalQuery:
        hints = self._infer_query_preferences(query_code=query.query_code, bug_description=query.bug_description)
        return query.model_copy(
            update={
                "error_classes": self._unique_list([*query.error_classes, *hints["error_classes"]]),
                "preferred_categories": self._unique_list([*query.preferred_categories, *hints["preferred_categories"]]),
                "preferred_datasets": self._unique_list([*query.preferred_datasets, *hints["preferred_datasets"]]),
                "preferred_example_kinds": self._unique_list(
                    [*query.preferred_example_kinds, *self._preferred_example_kinds(query.query_code, query.function_name)]
                ),
                "exclude_datasets": self._unique_list(query.exclude_datasets),
                "required_terms": self._unique_list([*query.required_terms, *hints["required_terms"]]),
            }
        )

    def _infer_query_preferences(self, *, query_code: str, bug_description: str | None) -> dict[str, list[str]]:
        description = bug_description or ""
        combined = f"{description}\n{query_code}"
        preferred_categories: list[str] = []
        preferred_datasets: list[str] = []
        error_classes: list[str] = []
        required_terms: list[str] = []

        for normalized_error, metadata in self.ERROR_CLASS_HINTS.items():
            explicit_hit = any(re.search(pattern, combined, re.IGNORECASE) for pattern in metadata["explicit_patterns"])
            heuristic_hit = bool(query_code) and any(
                re.search(pattern, query_code, re.IGNORECASE) for pattern in metadata["code_patterns"]
            )
            if not explicit_hit and not heuristic_hit:
                continue
            preferred_categories.append(str(metadata["category"]))
            preferred_datasets.extend(str(dataset) for dataset in metadata["preferred_datasets"])
            if explicit_hit:
                error_classes.append(str(metadata["display"]))
                required_terms.extend(str(term) for term in metadata["required_terms"])

        return {
            "error_classes": self._unique_list(error_classes),
            "preferred_categories": self._unique_list(preferred_categories),
            "preferred_datasets": self._unique_list(preferred_datasets),
            "required_terms": self._unique_list(required_terms),
        }

    def _rerank_candidates(
        self,
        query: RetrievalQuery,
        query_embedding: list[float] | None,
        vector_examples: list[RetrievedExample],
        keyword_examples: list[RetrievedExample],
    ) -> list[RetrievedExample]:
        candidates: dict[str, BugFixExample] = {}
        vector_scores = {item.example.id: max(item.relevance_score, 0.0) for item in vector_examples}
        keyword_scores = {item.example.id: max(item.relevance_score, 0.0) for item in keyword_examples}

        for item in [*vector_examples, *keyword_examples]:
            candidates[item.example.id] = item.example

        vector_weight = self.settings.rag_vector_weight if query_embedding else 0.0
        keyword_weight = self.settings.rag_keyword_weight if query_embedding else self.settings.rag_keyword_weight_no_embedding
        telemetry_adjustments = self._telemetry_adjustments(set(candidates))
        ranked: list[RetrievedExample] = []

        for example_id, example in candidates.items():
            vector_score = vector_scores.get(example_id, 0.0)
            keyword_score = keyword_scores.get(example_id, 0.0)
            score = (vector_weight * vector_score) + (keyword_weight * keyword_score)
            reasons: list[str] = []

            if vector_score:
                reasons.append("vector")
            if keyword_score:
                reasons.append("keyword")

            category_adjustment = self._category_match_adjustment(query, example)
            if category_adjustment:
                score += category_adjustment
                reasons.append("category" if category_adjustment > 0 else "category_mismatch")

            error_class_adjustment = self._error_class_adjustment(query, example)
            if error_class_adjustment:
                score += error_class_adjustment
                reasons.append("error_class" if error_class_adjustment > 0 else "error_class_mismatch")

            term_boost = self._required_term_boost(query, example)
            if term_boost:
                score += term_boost
                reasons.append("required_term")

            dataset_boost = self._dataset_preference_boost(query, example)
            if dataset_boost:
                score += dataset_boost
                reasons.append("dataset")

            kind_boost = self._example_kind_adjustment(query, example)
            if kind_boost:
                score += kind_boost
                reasons.append("snippet_kind")

            if query.function_name and example.function_name and query.function_name == example.function_name:
                score += self.settings.rag_rerank_function_name_boost
                reasons.append("function")

            synthetic_penalty = self._synthetic_penalty(query, example)
            if synthetic_penalty:
                score -= synthetic_penalty
                reasons.append("synthetic_penalty")

            telemetry_adjustment = telemetry_adjustments.get(example_id, 0.0)
            if telemetry_adjustment:
                score += telemetry_adjustment
                reasons.append("telemetry_success" if telemetry_adjustment > 0 else "telemetry_failure")

            ranked.append(
                RetrievedExample(
                    example=example,
                    relevance_score=score,
                    match_reason="+".join(reasons) or "hybrid",
                )
            )

        ranked.sort(key=lambda item: item.relevance_score, reverse=True)
        return ranked

    def _diversify_ranked_examples(
        self,
        examples: list[RetrievedExample],
        max_results: int,
        max_per_dataset: int | None,
    ) -> list[RetrievedExample]:
        if max_per_dataset is None:
            return examples[:max_results]

        selected: list[RetrievedExample] = []
        selected_ids: set[str] = set()
        counts: dict[str, int] = {}

        for item in examples:
            dataset = item.example.dataset
            if counts.get(dataset, 0) >= max_per_dataset:
                continue
            selected.append(item)
            selected_ids.add(item.example.id)
            counts[dataset] = counts.get(dataset, 0) + 1
            if len(selected) >= max_results:
                return selected

        for item in examples:
            if item.example.id in selected_ids:
                continue
            selected.append(item)
            if len(selected) >= max_results:
                break
        return selected

    def _category_match_adjustment(self, query: RetrievalQuery, example: BugFixExample) -> float:
        if example.bug_category in query.preferred_categories:
            return self.settings.rag_category_match_boost
        if query.preferred_categories:
            return -self.settings.rag_category_mismatch_penalty
        if query.categories and example.bug_category in query.categories:
            return self.settings.rag_query_category_match_boost
        return 0.0

    def _error_class_adjustment(self, query: RetrievalQuery, example: BugFixExample) -> float:
        if not query.error_classes:
            return 0.0
        if self._example_matches_error_classes(example, query.error_classes):
            return self.settings.rag_error_class_match_boost
        return -self.settings.rag_error_class_mismatch_penalty

    def _required_term_boost(self, query: RetrievalQuery, example: BugFixExample) -> float:
        if not query.required_terms:
            return 0.0
        example_text = self._example_search_text(example)
        hits = sum(1 for term in query.required_terms if term.lower() in example_text)
        if not hits:
            return 0.0
        return min(self.settings.rag_required_term_boost_cap, hits * self.settings.rag_required_term_boost_per_term)

    def _dataset_preference_boost(self, query: RetrievalQuery, example: BugFixExample) -> float:
        if example.dataset not in query.preferred_datasets:
            return 0.0
        index = query.preferred_datasets.index(example.dataset)
        return max(
            self.settings.rag_dataset_preference_min_boost,
            self.settings.rag_dataset_preference_base_boost - (index * self.settings.rag_dataset_preference_decay),
        )

    def _example_kind_adjustment(self, query: RetrievalQuery, example: BugFixExample) -> float:
        if example.example_kind not in query.preferred_example_kinds:
            return 0.0
        index = query.preferred_example_kinds.index(example.example_kind)
        base = max(
            self.settings.rag_example_kind_min_boost,
            self.settings.rag_example_kind_base_boost - (index * self.settings.rag_example_kind_decay),
        )
        if query.function_name and example.example_kind == "function_snippet":
            base += self.settings.rag_example_kind_function_bonus
        if len(query.query_code.splitlines()) <= self.settings.rag_short_query_line_limit and example.example_kind in {"function_snippet", "diff_hunk"}:
            base += self.settings.rag_example_kind_short_query_bonus
        return base

    def _synthetic_penalty(self, query: RetrievalQuery, example: BugFixExample) -> float:
        if example.dataset != "synthetic":
            return 0.0
        if query.preferred_datasets or query.error_classes or query.preferred_categories:
            return self.settings.rag_synthetic_penalty_preferred
        return self.settings.rag_synthetic_penalty_default

    def _telemetry_adjustments(self, example_ids: set[str]) -> dict[str, float]:
        if not example_ids:
            return {}

        scores: dict[str, list[int]] = {}
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            rows = conn.execute(
                """
                SELECT example_ids, outcome
                FROM rag_retrieval_telemetry
                WHERE outcome IS NOT NULL
                ORDER BY created_at DESC
                LIMIT 1000
                """
            ).fetchall()

        for row in rows:
            outcome = str(row["outcome"] or "").strip().lower()
            if outcome not in {"success", "failure"}:
                continue
            for example_id in json.loads(row["example_ids"] or "[]"):
                if example_id not in example_ids:
                    continue
                bucket = scores.setdefault(example_id, [0, 0])
                if outcome == "success":
                    bucket[0] += 1
                else:
                    bucket[1] += 1

        adjustments: dict[str, float] = {}
        for example_id, (success_count, failure_count) in scores.items():
            total = success_count + failure_count
            if total <= 0:
                continue
            success_ratio = success_count / total
            failure_ratio = failure_count / total
            adjustment = (
                success_ratio * self.settings.rag_telemetry_success_boost
                - failure_ratio * self.settings.rag_telemetry_failure_penalty
            )
            if adjustment:
                adjustments[example_id] = adjustment
        return adjustments

    def _example_matches_error_classes(self, example: BugFixExample, error_classes: list[str]) -> bool:
        normalized_targets = {self._normalize_error_class(value) for value in error_classes}
        normalized_targets.discard("")
        if not normalized_targets:
            return False

        example_matches = set()
        category_match = self.CATEGORY_TO_ERROR_CLASS.get(example.bug_category)
        if category_match:
            example_matches.add(category_match)

        example_text = self._example_search_text(example)
        for normalized_error, metadata in self.ERROR_CLASS_HINTS.items():
            if any(re.search(pattern, example_text, re.IGNORECASE) for pattern in metadata["explicit_patterns"]):
                example_matches.add(normalized_error)

        return bool(normalized_targets & example_matches)

    def _build_search_terms(self, query: RetrievalQuery) -> list[str]:
        terms: list[str] = []
        terms.extend(query.required_terms)
        terms.extend(query.error_classes)
        terms.extend(query.preferred_categories)
        terms.extend(query.preferred_example_kinds)
        if query.function_name:
            terms.append(query.function_name)
        if query.file_path:
            terms.append(query.file_path.split("/")[-1])

        if query.bug_description:
            terms.extend(
                token
                for token in re.findall(r"\b[a-zA-Z_][a-zA-Z0-9_]{2,}\b", query.bug_description)
                if token.lower() not in self.QUERY_STOPWORDS
            )

        if query.query_code:
            code_identifiers = re.findall(r"\b[a-zA-Z_][a-zA-Z0-9_]{2,}\b", query.query_code)
            terms.extend(token for token in code_identifiers[:20] if token.lower() not in self.QUERY_STOPWORDS)

        return self._unique_list(terms)

    @staticmethod
    def _quote_fts_term(term: str) -> str:
        escaped = term.replace('"', '""').strip()
        if " " in escaped or "-" in escaped:
            return f'"{escaped}"'
        return escaped

    def _example_search_text(self, example: BugFixExample) -> str:
        parts = [
            example.dataset,
            example.repository,
            example.example_kind,
            example.snippet_role or "",
            example.file_path,
            example.bug_category,
            example.bug_summary,
            example.bug_description,
            example.root_cause or "",
            example.function_name or "",
            " ".join(example.symptoms),
            " ".join(example.related_patterns),
            example.task_mode or "",
            " ".join(example.validation_failures),
            " ".join(example.source_issue_types),
            " ".join(example.retrieval_tags),
            example.buggy_code[:400],
            example.fixed_code[:400],
        ]
        return " ".join(part.lower() for part in parts if part)

    def _normalize_error_class(self, value: str) -> str:
        normalized = re.sub(r"[^a-z0-9]+", "", value.lower())
        if normalized in self.ERROR_CLASS_HINTS:
            return normalized
        for key, metadata in self.ERROR_CLASS_HINTS.items():
            if normalized == re.sub(r"[^a-z0-9]+", "", str(metadata["display"]).lower()):
                return key
        return ""

    @staticmethod
    def _unique_list(values: list[str]) -> list[str]:
        seen: set[str] = set()
        result: list[str] = []
        for value in values:
            normalized = str(value).strip()
            if not normalized or normalized in seen:
                continue
            seen.add(normalized)
            result.append(normalized)
        return result

    @staticmethod
    def _cosine_similarity(vec1: np.ndarray, vec2: np.ndarray) -> float:
        norm1 = np.linalg.norm(vec1)
        norm2 = np.linalg.norm(vec2)
        if norm1 == 0 or norm2 == 0:
            return 0.0
        return float(np.dot(vec1, vec2) / (norm1 * norm2))

    def _row_to_example(self, row: sqlite3.Row) -> BugFixExample:
        embedding = json.loads(row["embedding"]) if row["embedding"] else None
        symptoms = json.loads(row["symptoms"]) if row["symptoms"] else []
        patterns = json.loads(row["related_patterns"]) if row["related_patterns"] else []
        validation_failures = json.loads(row["validation_failures"]) if row["validation_failures"] else []
        source_issue_types = json.loads(row["source_issue_types"]) if row["source_issue_types"] else []
        retrieval_tags = json.loads(row["retrieval_tags"]) if row["retrieval_tags"] else []
        return BugFixExample(
            id=row["id"],
            dataset=row["dataset"],
            repository=row["repository"],
            example_kind=row["example_kind"] or "full_example",
            parent_example_id=row["parent_example_id"],
            snippet_role=row["snippet_role"],
            snippet_index=row["snippet_index"],
            commit_hash=row["commit_hash"],
            issue_url=row["issue_url"],
            bug_description=row["bug_description"],
            bug_summary=row["bug_summary"],
            bug_category=row["bug_category"],
            file_path=row["file_path"],
            function_name=row["function_name"],
            line_number=row["line_number"],
            buggy_code=row["buggy_code"],
            fixed_code=row["fixed_code"],
            diff=row["diff"],
            language=row["language"],
            severity=row["severity"],
            complexity=row["complexity"],
            root_cause=row["root_cause"],
            symptoms=symptoms,
            related_patterns=patterns,
            task_mode=row["task_mode"],
            validation_failures=validation_failures,
            source_issue_types=source_issue_types,
            retrieval_tags=retrieval_tags,
            embedding=embedding,
        )

    def _build_context_summary(self, examples: list[RetrievedExample]) -> str:
        if not examples:
            return "No similar bug-fix examples found."
        categories: dict[str, list[str]] = {}
        for item in examples:
            categories.setdefault(item.example.bug_category, []).append(item.example.repository)
        parts = []
        for category, repos in categories.items():
            unique_repos = sorted(set(repos))
            parts.append(f"{len(unique_repos)} {category} issue(s) from {', '.join(unique_repos[:3])}")
        return f"Found {len(examples)} similar examples: " + ", ".join(parts)

    def get_statistics(self) -> DatasetStatistics:
        with sqlite3.connect(self.db_path) as conn:
            total = conn.execute("SELECT COUNT(*) FROM bug_fix_examples").fetchone()[0]
            by_dataset = {
                row[0]: row[1]
                for row in conn.execute("SELECT dataset, COUNT(*) FROM bug_fix_examples GROUP BY dataset")
            }
            by_category = {
                row[0]: row[1]
                for row in conn.execute("SELECT bug_category, COUNT(*) FROM bug_fix_examples GROUP BY bug_category")
            }
            by_severity = {
                row[0]: row[1]
                for row in conn.execute("SELECT severity, COUNT(*) FROM bug_fix_examples GROUP BY severity")
            }
            last_updated = conn.execute("SELECT MAX(indexed_at) FROM bug_fix_examples").fetchone()[0]
        return DatasetStatistics(
            total_examples=total,
            by_dataset=by_dataset,
            by_category=by_category,
            by_severity=by_severity,
            last_updated=last_updated or datetime.now(UTC).isoformat(),
        )

    def record_retrieval_event(
        self,
        *,
        task_id: str | None,
        stage: str,
        query: RetrievalQuery,
        examples: list[RetrievedExample],
        file_path: str | None = None,
        function_name: str | None = None,
    ) -> str:
        event_id = self._stable_example_id(
            "rag-retrieval",
            task_id or "anonymous",
            stage,
            file_path or "",
            function_name or "",
            query.query_code[:80],
            datetime.now(UTC).isoformat(),
        )
        query_text = "\n".join(
            part
            for part in (
                query.bug_description or "",
                query.file_path or "",
                query.function_name or "",
                query.query_code[:1200],
            )
            if part
        )
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                """
                INSERT OR REPLACE INTO rag_retrieval_telemetry (
                    event_id, task_id, stage, file_path, function_name, query_digest,
                    query_text, example_ids, example_scores, outcome, created_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    event_id,
                    task_id,
                    stage,
                    file_path,
                    function_name,
                    hashlib.sha256(query_text.encode("utf-8")).hexdigest()[:16],
                    query_text,
                    json.dumps([item.example.id for item in examples]),
                    json.dumps({item.example.id: item.relevance_score for item in examples}),
                    None,
                    datetime.now(UTC).isoformat(),
                ),
            )
            conn.commit()
        return event_id

    def mark_task_retrieval_outcome(self, task_id: str, *, success: bool) -> None:
        outcome = "success" if success else "failure"
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                "UPDATE rag_retrieval_telemetry SET outcome = ? WHERE task_id = ?",
                (outcome, task_id),
            )
            conn.commit()

    def search_by_category(self, category: str, limit: int = 10) -> list[BugFixExample]:
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            rows = conn.execute(
                "SELECT * FROM bug_fix_examples WHERE bug_category = ? LIMIT ?",
                (category, limit),
            ).fetchall()
        return [self._row_to_example(row) for row in rows]

    def search_by_repository(self, repository: str, limit: int = 10) -> list[BugFixExample]:
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            rows = conn.execute(
                "SELECT * FROM bug_fix_examples WHERE repository = ? LIMIT ?",
                (repository, limit),
            ).fetchall()
        return [self._row_to_example(row) for row in rows]

    def clear_dataset(self, dataset: str) -> int:
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute("DELETE FROM bug_fix_examples WHERE dataset = ?", (dataset,))
            conn.commit()
        return cursor.rowcount

    def delete_example(self, example_id: str) -> bool:
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute("DELETE FROM bug_fix_examples WHERE id = ?", (example_id,))
            conn.commit()
        return cursor.rowcount > 0

    @staticmethod
    def _example_files_from_path(path: Path) -> list[Path]:
        if path.is_file():
            return [path]
        if not path.exists():
            raise FileNotFoundError(f"Seed examples path not found: {path}")
        return sorted(
            file_path
            for file_path in path.rglob("*")
            if file_path.is_file() and file_path.suffix.lower() in {".json", ".jsonl"}
        )

    def _load_examples_from_file(self, path: Path) -> list[BugFixExample]:
        if path.suffix.lower() == ".jsonl":
            examples: list[BugFixExample] = []
            for line in path.read_text(encoding="utf-8").splitlines():
                line = line.strip()
                if not line:
                    continue
                examples.append(BugFixExample.model_validate(json.loads(line)))
            return examples

        payload = json.loads(path.read_text(encoding="utf-8"))
        if isinstance(payload, dict):
            payload = payload.get("examples", [])
        if not isinstance(payload, list):
            raise ValueError(f"Unsupported seed example payload in {path}")
        return [BugFixExample.model_validate(item) for item in payload if isinstance(item, dict)]

    def _prepare_example_record(self, example: BugFixExample) -> tuple[Any, ...]:
        embedding_blob = json.dumps(example.embedding).encode() if example.embedding else None
        symptoms_json = json.dumps(example.symptoms) if example.symptoms else None
        patterns_json = json.dumps(example.related_patterns) if example.related_patterns else None
        validation_failures_json = json.dumps(example.validation_failures) if example.validation_failures else None
        source_issue_types_json = json.dumps(example.source_issue_types) if example.source_issue_types else None
        retrieval_tags_json = json.dumps(example.retrieval_tags) if example.retrieval_tags else None
        return (
            example.id,
            example.dataset,
            example.repository,
            example.example_kind,
            example.parent_example_id,
            example.snippet_role,
            example.snippet_index,
            example.commit_hash,
            example.issue_url,
            example.bug_description,
            example.bug_summary,
            example.bug_category,
            example.file_path,
            example.function_name,
            example.line_number,
            example.buggy_code,
            example.fixed_code,
            example.diff,
            example.language,
            example.severity,
            example.complexity,
            example.root_cause,
            symptoms_json,
            patterns_json,
            example.task_mode,
            validation_failures_json,
            source_issue_types_json,
            retrieval_tags_json,
            embedding_blob,
            datetime.now(UTC).isoformat(),
        )

    def _family_examples(self, example: BugFixExample, embedding_service: EmbeddingService) -> list[BugFixExample]:
        family = [example.model_copy(deep=True)]
        family.extend(self._build_function_snippet_examples(example))
        family.extend(self._build_diff_hunk_examples(example))

        prepared: list[BugFixExample] = []
        seen_ids: set[str] = set()
        for item in family:
            if item.id in seen_ids:
                continue
            seen_ids.add(item.id)
            if item.embedding is None:
                item = item.model_copy(
                    update={"embedding": embedding_service.embed_code(self._example_embedding_text(item))}
                )
            prepared.append(item)
        return prepared

    def _build_function_snippet_examples(self, example: BugFixExample) -> list[BugFixExample]:
        before_functions = self._parse_functions(example.buggy_code)
        after_functions = self._parse_functions(example.fixed_code)
        if not before_functions and not after_functions:
            return []

        snippets: list[BugFixExample] = []
        for index, function_name in enumerate(sorted(set(before_functions) | set(after_functions)), start=1):
            before_metadata = before_functions.get(function_name, {})
            after_metadata = after_functions.get(function_name, {})
            before_code = str(before_metadata.get("code", "")).strip()
            after_code = str(after_metadata.get("code", "")).strip()
            if not before_code and not after_code:
                continue
            if before_code == after_code:
                continue
            diff_text = self._generate_diff(before_code, after_code)
            snippets.append(
                BugFixExample(
                    id=self._stable_example_id(example.id, "function", function_name),
                    dataset=example.dataset,
                    repository=example.repository,
                    example_kind="function_snippet",
                    parent_example_id=example.id,
                    snippet_role="function",
                    snippet_index=index,
                    commit_hash=example.commit_hash,
                    issue_url=example.issue_url,
                    bug_description=example.bug_description,
                    bug_summary=f"{example.bug_summary} [{function_name}]",
                    bug_category=example.bug_category,
                    file_path=example.file_path,
                    function_name=function_name,
                    line_number=after_metadata.get("line_number") or before_metadata.get("line_number") or example.line_number,
                    buggy_code=before_code or example.buggy_code,
                    fixed_code=after_code or example.fixed_code,
                    diff=diff_text,
                    language=example.language,
                    severity=example.severity,
                    complexity=self._assess_complexity(diff_text),
                    root_cause=example.root_cause,
                    symptoms=list(example.symptoms),
                    related_patterns=list(example.related_patterns),
                    task_mode=example.task_mode,
                    validation_failures=list(example.validation_failures),
                    source_issue_types=list(example.source_issue_types),
                    retrieval_tags=list(example.retrieval_tags),
                )
            )
        return snippets

    def _build_diff_hunk_examples(self, example: BugFixExample) -> list[BugFixExample]:
        parsed_files = self._parse_unified_diff(example.diff)
        snippets: list[BugFixExample] = []
        if not parsed_files:
            return snippets

        index = 0
        for file_diff in parsed_files:
            for hunk in file_diff.hunks:
                if hunk.before_text.strip() == hunk.after_text.strip():
                    continue
                index += 1
                snippets.append(
                    BugFixExample(
                        id=self._stable_example_id(example.id, "hunk", str(index)),
                        dataset=example.dataset,
                        repository=example.repository,
                        example_kind="diff_hunk",
                        parent_example_id=example.id,
                        snippet_role="diff_hunk",
                        snippet_index=index,
                        commit_hash=example.commit_hash,
                        issue_url=example.issue_url,
                        bug_description=example.bug_description,
                        bug_summary=f"{example.bug_summary} [hunk {index}]",
                        bug_category=example.bug_category,
                        file_path=file_diff.file_path or example.file_path,
                        function_name=example.function_name,
                        line_number=hunk.line_number or example.line_number,
                        buggy_code=hunk.before_text or example.buggy_code,
                        fixed_code=hunk.after_text or example.fixed_code,
                        diff=hunk.diff_text,
                        language=example.language,
                        severity=example.severity,
                        complexity=self._assess_complexity(hunk.diff_text),
                        root_cause=example.root_cause,
                        symptoms=list(example.symptoms),
                        related_patterns=list(example.related_patterns),
                        task_mode=example.task_mode,
                        validation_failures=list(example.validation_failures),
                        source_issue_types=list(example.source_issue_types),
                        retrieval_tags=list(example.retrieval_tags),
                    )
                )
        return snippets

    def _example_embedding_text(self, example: BugFixExample) -> str:
        parts = [
            example.bug_summary,
            example.bug_category,
            example.example_kind,
            example.file_path,
            example.function_name or "",
            example.bug_description,
            example.root_cause or "",
            " ".join(example.symptoms),
            " ".join(example.related_patterns),
            example.task_mode or "",
            " ".join(example.validation_failures),
            " ".join(example.source_issue_types),
            " ".join(example.retrieval_tags),
            example.buggy_code[:800],
            example.fixed_code[:800],
            example.diff[:800],
        ]
        return "\n".join(part for part in parts if part)

    def _parse_functions(self, code: str) -> dict[str, dict[str, Any]]:
        try:
            tree = ast.parse(code)
        except SyntaxError:
            return {}

        lines = code.splitlines()
        functions: dict[str, dict[str, Any]] = {}
        for node in ast.walk(tree):
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            start_line = max(getattr(node, "lineno", 1), 1)
            end_line = max(getattr(node, "end_lineno", start_line), start_line)
            snippet = "\n".join(lines[start_line - 1 : end_line]).strip()
            if not snippet:
                continue
            functions[node.name] = {"code": snippet, "line_number": start_line}
        return functions

    @staticmethod
    def _generate_diff(before: str, after: str) -> str:
        return "".join(
            difflib.unified_diff(
                before.splitlines(keepends=True),
                after.splitlines(keepends=True),
                fromfile="before",
                tofile="after",
                lineterm="",
            )
        )

    def _parse_unified_diff(self, diff_text: str) -> list[ParsedDiffFile]:
        if not diff_text.strip():
            return []

        lines = diff_text.splitlines()
        files: list[ParsedDiffFile] = []
        current_file_path = ""
        current_before: list[str] = []
        current_after: list[str] = []
        current_diff_lines: list[str] = []
        hunks: list[ParsedDiffHunk] = []
        hunk_before: list[str] = []
        hunk_after: list[str] = []
        hunk_lines: list[str] = []
        hunk_line_number: int | None = None

        def flush_hunk() -> None:
            nonlocal hunk_before, hunk_after, hunk_lines, hunk_line_number
            if not hunk_lines:
                return
            hunks.append(
                ParsedDiffHunk(
                    line_number=hunk_line_number,
                    before_text="\n".join(hunk_before).strip(),
                    after_text="\n".join(hunk_after).strip(),
                    diff_text="\n".join(hunk_lines).strip(),
                )
            )
            hunk_before = []
            hunk_after = []
            hunk_lines = []
            hunk_line_number = None

        def flush_file() -> None:
            nonlocal current_file_path, current_before, current_after, current_diff_lines, hunks
            flush_hunk()
            if not current_file_path and not current_diff_lines:
                return
            files.append(
                ParsedDiffFile(
                    file_path=current_file_path,
                    before_text="\n".join(current_before).strip(),
                    after_text="\n".join(current_after).strip(),
                    diff_text="\n".join(current_diff_lines).strip(),
                    line_number=hunks[0].line_number if hunks else None,
                    hunks=list(hunks),
                )
            )
            current_file_path = ""
            current_before = []
            current_after = []
            current_diff_lines = []
            hunks = []

        for line in lines:
            if line.startswith("diff --git "):
                flush_file()
                continue
            if line.startswith("--- "):
                current_diff_lines.append(line)
                continue
            if line.startswith("+++ "):
                current_diff_lines.append(line)
                candidate = line[4:].strip()
                if candidate.startswith("b/"):
                    candidate = candidate[2:]
                current_file_path = candidate
                continue
            if line.startswith("@@ "):
                flush_hunk()
                current_diff_lines.append(line)
                hunk_lines.append(line)
                match = re.search(r"\+(\d+)", line)
                hunk_line_number = int(match.group(1)) if match else None
                continue

            current_diff_lines.append(line)
            if current_file_path and not hunk_lines:
                continue
            hunk_lines.append(line)
            if line.startswith("-") and not line.startswith("---"):
                current_before.append(line[1:])
                hunk_before.append(line[1:])
            elif line.startswith("+") and not line.startswith("+++"):
                current_after.append(line[1:])
                hunk_after.append(line[1:])
            elif line.startswith(" "):
                current_before.append(line[1:])
                current_after.append(line[1:])
                hunk_before.append(line[1:])
                hunk_after.append(line[1:])

        flush_file()
        return files

    @staticmethod
    def _stable_example_id(*parts: str) -> str:
        joined = "::".join(part.strip() for part in parts if part.strip())
        digest = hashlib.sha256(joined.encode("utf-8")).hexdigest()[:16]
        label = "-".join(part.strip().replace("/", "-").replace("\\", "-") for part in parts[:2] if part.strip())
        label = re.sub(r"[^a-zA-Z0-9._-]+", "-", label).strip("-")
        return f"{label}-{digest}" if label else digest

    def _infer_category_from_task(self, task: TaskState, file_diff: ParsedDiffFile) -> str:
        patterns = self._task_patterns(task)
        for pattern in patterns:
            category = self.TASK_CATEGORY_HINTS.get(pattern)
            if category:
                return category

        hints = self._infer_query_preferences(
            query_code=file_diff.diff_text,
            bug_description=f"{task.goal}\n{task.summary or ''}",
        )
        if hints["preferred_categories"]:
            return hints["preferred_categories"][0]

        diff_lower = file_diff.diff_text.lower()
        if ".get(" in diff_lower and "[" in diff_lower:
            return "key-error"
        if "encoding" in diff_lower or "read_text" in diff_lower or "open(" in diff_lower:
            return "encoding-issue"
        return "logic-error"

    @staticmethod
    def _task_root_cause(task: TaskState) -> str | None:
        issue_titles: list[str] = []
        if task.focus_issue_inventory is not None:
            issue_titles.extend(issue.title for issue in task.focus_issue_inventory.issues[:4])
        issue_titles.extend(task.reflection_issues[:2])
        issue_titles = [item.strip() for item in issue_titles if item.strip()]
        if not issue_titles:
            return None
        return " | ".join(issue_titles[:4])

    @staticmethod
    def _task_symptoms(task: TaskState) -> list[str]:
        symptoms: list[str] = []
        for summary in task.validation_history:
            symptoms.extend(summary.failed_checks)
        return sorted({item for item in symptoms if item})

    @staticmethod
    def _task_patterns(task: TaskState) -> list[str]:
        patterns: list[str] = []
        if task.focus_issue_inventory is not None:
            patterns.extend(issue.issue_type for issue in task.focus_issue_inventory.issues)
        patterns.extend(task.reflection_issues)
        return sorted({item for item in patterns if item})

    @staticmethod
    def _task_validation_failures(task: TaskState) -> list[str]:
        failures: list[str] = []
        for summary in task.validation_history:
            failures.extend(summary.failed_checks)
        return sorted({item for item in failures if item})

    @staticmethod
    def _task_retrieval_tags(task: TaskState) -> list[str]:
        tags = {task.task_mode}
        tags.update(task.reflection_issues)
        if task.focus_issue_inventory is not None:
            tags.update(issue.issue_type for issue in task.focus_issue_inventory.issues)
            tags.update(check for check in task.focus_issue_inventory.acceptance_checks if check)
        for summary in task.validation_history:
            tags.update(summary.failed_checks)
        return sorted(tag for tag in tags if tag)

    @staticmethod
    def _extract_preview_by_path(task: TaskState) -> dict[str, str]:
        previews: dict[str, str] = {}
        for preview in task.repo_context.get("file_previews", []):
            if not isinstance(preview, dict):
                continue
            path = str(preview.get("path", "")).strip()
            content = str(preview.get("content", ""))
            if path:
                previews[path] = content
        return previews

    @staticmethod
    def _extract_primary_function_name(code: str) -> str | None:
        try:
            tree = ast.parse(code)
        except SyntaxError:
            return None
        for node in ast.walk(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                return node.name
        return None

    @staticmethod
    def _assess_complexity(diff_text: str) -> str:
        added = len([line for line in diff_text.splitlines() if line.startswith("+") and not line.startswith("+++")])
        removed = len([line for line in diff_text.splitlines() if line.startswith("-") and not line.startswith("---")])
        total = added + removed
        if total <= 3:
            return "simple"
        if total <= 10:
            return "moderate"
        return "complex"

    def _preferred_example_kinds(self, query_code: str, function_name: str | None) -> list[str]:
        if function_name:
            return ["function_snippet", "diff_hunk", "full_example", "task_fix"]
        if len(query_code.splitlines()) <= self.settings.rag_short_query_line_limit:
            return ["function_snippet", "diff_hunk", "full_example", "task_fix"]
        return ["diff_hunk", "function_snippet", "full_example", "task_fix"]
