coderag/ui/handlers.py

"""UI event handlers for Gradio interface."""

import json
from datetime import datetime
from pathlib import Path
from typing import Iterator, Optional

import torch

from coderag.config import get_settings
from coderag.generation.generator import ResponseGenerator
from coderag.indexing.embeddings import EmbeddingGenerator
from coderag.indexing.vectorstore import VectorStore
from coderag.ingestion.chunker import CodeChunker
from coderag.ingestion.filter import FileFilter
from coderag.ingestion.loader import RepositoryLoader
from coderag.ingestion.validator import GitHubURLValidator, ValidationError
from coderag.logging import get_logger
from coderag.models.chunk import Chunk
from coderag.models.document import Document
from coderag.models.query import Query
from coderag.models.repository import Repository, RepositoryStatus

logger = get_logger(__name__)


class UIHandlers:
    """Handlers for Gradio UI events."""

    def __init__(self) -> None:
        self.settings = get_settings()
        self.validator = GitHubURLValidator()
        self.loader = RepositoryLoader()
        self.filter = FileFilter()
        self.chunker = CodeChunker()
        self.embedder = EmbeddingGenerator()
        self.vectorstore = VectorStore()
        self.generator: Optional[ResponseGenerator] = None

        # Repository metadata storage
        self.repos_file = self.settings.data_dir / "repositories.json"
        self.repositories: dict[str, Repository] = self._load_repositories()

    def _load_repositories(self) -> dict[str, Repository]:
        if self.repos_file.exists():
            try:
                data = json.loads(self.repos_file.read_text())
                return {r["id"]: Repository.from_dict(r) for r in data}
            except Exception as e:
                logger.error("Failed to load repositories", error=str(e))
        return {}

    def _save_repositories(self) -> None:
        self.repos_file.parent.mkdir(parents=True, exist_ok=True)
        data = [r.to_dict() for r in self.repositories.values()]
        self.repos_file.write_text(json.dumps(data, indent=2))

    # =========================================================================
    # Streaming Methods (Nivel 1)
    # =========================================================================

    def _document_generator(
        self,
        files: list[Path],
        repo_path: Path,
        repo_id: str,
    ) -> Iterator[Document]:
        """Generate documents one by one without accumulating in memory."""
        for file_path in files:
            try:
                yield Document.from_file(file_path, repo_path, repo_id)
            except Exception as e:
                logger.warning("Failed to load file", path=str(file_path), error=str(e))

    def _process_batch(self, chunks: list[Chunk]) -> int:
        """Process a batch: embed + store + release memory."""
        if not chunks:
            return 0

        embedded = self.embedder.embed_chunks(chunks, show_progress=False)
        self.vectorstore.add_chunks(embedded)

        # Release memory
        del embedded
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        return len(chunks)

    def _stream_index_repository(
        self,
        documents: Iterator[Document],
        repo_id: str,
        batch_size: int = 100,
        progress_callback: Optional[callable] = None,
    ) -> int:
        """Index using streaming with batches."""
        total_chunks = 0
        batch: list[Chunk] = []
        doc_count = 0

        for doc in documents:
            doc_count += 1
            for chunk in self.chunker.chunk_document(doc):
                chunk.repo_id = repo_id
                batch.append(chunk)

                if len(batch) >= batch_size:
                    total_chunks += self._process_batch(batch)
                    logger.info("Batch processed", total_so_far=total_chunks, docs_processed=doc_count)
                    if progress_callback:
                        progress_callback(total_chunks, doc_count)
                    batch = []

        # Process final batch
        if batch:
            total_chunks += self._process_batch(batch)
            if progress_callback:
                progress_callback(total_chunks, doc_count)

        return total_chunks

    # =========================================================================
    # Validation Methods (Nivel 2)
    # =========================================================================

    def _estimate_repo_size(self, files: list[Path]) -> dict:
        """Estimate repository size before indexing."""
        total_size_kb = 0
        estimated_chunks = 0
        chunk_size = self.settings.ingestion.chunk_size

        for file_path in files:
            try:
                size_kb = file_path.stat().st_size / 1024
                total_size_kb += size_kb
                # Rough estimate: 1 chunk per chunk_size characters
                estimated_chunks += max(1, int(size_kb * 1024 / chunk_size))
            except Exception:
                continue

        return {
            "file_count": len(files),
            "total_size_kb": total_size_kb,
            "estimated_chunks": estimated_chunks,
            "exceeds_file_limit": len(files) > self.settings.ingestion.max_files_per_repo,
            "exceeds_chunk_limit": estimated_chunks > self.settings.ingestion.max_total_chunks,
            "warn_large_repo": len(files) > self.settings.ingestion.warn_files_threshold,
        }

    def _validate_repo_size(self, estimate: dict) -> tuple[bool, str]:
        """Validate if the repository can be indexed."""
        if estimate["exceeds_file_limit"]:
            return False, f"Repository exceeds file limit ({estimate['file_count']} > {self.settings.ingestion.max_files_per_repo})"
        if estimate["exceeds_chunk_limit"]:
            return False, f"Repository exceeds chunk limit (~{estimate['estimated_chunks']} > {self.settings.ingestion.max_total_chunks})"

        warning = ""
        if estimate["warn_large_repo"]:
            warning = f"Large repository ({estimate['file_count']} files, ~{estimate['estimated_chunks']} chunks). Processing may take several minutes."

        return True, warning

    # =========================================================================
    # Incremental Indexing Methods (Nivel 3)
    # =========================================================================

    def _get_current_commit(self, repo_path: Path) -> str:
        """Get the SHA of the current commit."""
        from git import Repo
        git_repo = Repo(repo_path)
        return git_repo.head.commit.hexsha

    def _get_changed_files(
        self,
        repo_path: Path,
        last_commit: str,
        current_commit: str,
    ) -> tuple[set[str], set[str], set[str]]:
        """Get files that were added, modified, or deleted between commits."""
        from git import Repo
        git_repo = Repo(repo_path)

        diff = git_repo.commit(last_commit).diff(current_commit)

        added: set[str] = set()
        modified: set[str] = set()
        deleted: set[str] = set()

        for d in diff:
            if d.new_file:
                added.add(d.b_path)
            elif d.deleted_file:
                deleted.add(d.a_path)
            elif d.renamed:
                deleted.add(d.a_path)
                added.add(d.b_path)
            else:
                modified.add(d.b_path or d.a_path)

        return added, modified, deleted

    def index_repository_incremental(self, repo_id: str) -> str:
        """Update only modified files since last indexing (incremental update)."""
        # Find repository by full or partial ID
        found_repo = None
        for rid, repo in self.repositories.items():
            if rid == repo_id or rid.startswith(repo_id):
                found_repo = repo
                break

        if not found_repo:
            return "Repository not found"

        repo = found_repo

        if not repo.last_commit:
            return "No previous indexing found. Please re-index the full repository."

        if not repo.clone_path or not Path(repo.clone_path).exists():
            return "Repository cache not found. Please re-index."

        try:
            repo_path = Path(repo.clone_path)

            # Update local repository
            logger.info("Updating local repository", repo_id=repo.id)
            self.loader._update_repository(repo_path, repo.branch, None)

            current_commit = self._get_current_commit(repo_path)

            if current_commit == repo.last_commit:
                return "Repository is already up to date."

            added, modified, deleted = self._get_changed_files(
                repo_path, repo.last_commit, current_commit
            )

            logger.info(
                "Changes detected",
                added=len(added),
                modified=len(modified),
                deleted=len(deleted),
            )

            # Delete chunks for deleted/modified files
            for file_path in deleted | modified:
                self.vectorstore.delete_file_chunks(repo.id, file_path)

            # Index new/modified files
            files_to_index = []
            file_filter = FileFilter()
            for file_path in added | modified:
                full_path = repo_path / file_path
                if full_path.exists() and file_filter.should_include(full_path, repo_path):
                    files_to_index.append(full_path)

            new_chunks = 0
            if files_to_index:
                batch_size = self.settings.ingestion.batch_size
                doc_generator = self._document_generator(files_to_index, repo_path, repo.id)
                new_chunks = self._stream_index_repository(doc_generator, repo.id, batch_size)

            # Update metadata
            repo.last_commit = current_commit
            repo.indexed_at = datetime.now()
            repo.chunk_count = self.vectorstore.get_repo_chunk_count(repo.id)
            self._save_repositories()

            return (
                f"Incremental update complete:\n"
                f"- Added/Modified: {len(added | modified)} files\n"
                f"- Deleted: {len(deleted)} files\n"
                f"- New chunks: {new_chunks}\n"
                f"- Total chunks: {repo.chunk_count}"
            )

        except Exception as e:
            logger.error("Incremental indexing failed", error=str(e), exc_info=True)
            return f"Error: {str(e)}"

    def index_repository(
        self,
        url: str,
        branch: str = "",
        include_patterns: str = "",
        exclude_patterns: str = "",
    ) -> Iterator[str]:
        """Index a GitHub repository with progress updates."""
        try:
            # Validate URL (sync version, skip accessibility check for UI)
            yield "Validating repository URL..."
            logger.info("Starting indexing", url=url, branch=branch)
            repo_info = self.validator.parse_url(url)
            branch = branch.strip() or repo_info.branch or "main"

            # Create repository record
            repo = Repository(
                url=repo_info.url,
                branch=branch,
                status=RepositoryStatus.CLONING,
            )
            self.repositories[repo.id] = repo

            # Clone repository
            yield f"Cloning {repo_info.full_name} (branch: {branch})..."
            logger.info("Cloning repository", url=url, branch=branch)
            repo_path = self.loader.clone_repository(repo_info, branch)
            repo.clone_path = repo_path
            repo.status = RepositoryStatus.INDEXING

            # Setup filter with custom patterns
            include = [p.strip() for p in include_patterns.split(",") if p.strip()] or None
            exclude = [p.strip() for p in exclude_patterns.split(",") if p.strip()] or None
            file_filter = FileFilter(include_patterns=include, exclude_patterns=exclude)

            # Process files
            yield "Scanning files..."
            logger.info("Filtering files", repo_path=str(repo_path))
            files = list(file_filter.filter_files(repo_path))
            file_count = len(files)
            logger.info("Files to process", count=file_count)

            # Validate repository size (Nivel 2)
            estimate = self._estimate_repo_size(files)
            can_proceed, message = self._validate_repo_size(estimate)

            if not can_proceed:
                repo.status = RepositoryStatus.ERROR
                repo.error_message = message
                self._save_repositories()
                yield f"Error: {message}"
                return

            if message:
                logger.warning(message)
                yield f"Warning: {message}"

            yield f"Found {file_count} files to index (~{estimate['estimated_chunks']} chunks)"

            # Delete existing chunks for this repo (re-indexing)
            logger.info("Deleting previous chunks for repo", repo_id=repo.id)
            self.vectorstore.delete_repo_chunks(repo.id)

            # Stream indexing with batches and progress updates
            yield f"Indexing... (0/{file_count} files, 0 chunks)"
            logger.info("Starting streaming indexing", file_count=file_count)
            batch_size = self.settings.ingestion.batch_size
            doc_generator = self._document_generator(files, repo_path, repo.id)

            # Process with progress updates
            total_chunks = 0
            batch: list[Chunk] = []
            doc_count = 0

            for doc in doc_generator:
                doc_count += 1
                for chunk in self.chunker.chunk_document(doc):
                    chunk.repo_id = repo.id
                    batch.append(chunk)

                    if len(batch) >= batch_size:
                        total_chunks += self._process_batch(batch)
                        batch = []
                        # Yield progress update
                        yield f"Indexing... ({doc_count}/{file_count} files, {total_chunks} chunks)"

            # Process final batch
            if batch:
                total_chunks += self._process_batch(batch)

            logger.info("Streaming indexing complete", chunk_count=total_chunks)

            # Save current commit for incremental updates (Nivel 3)
            try:
                repo.last_commit = self._get_current_commit(repo_path)
            except Exception:
                repo.last_commit = None

            # Update repository status
            repo.chunk_count = total_chunks
            repo.indexed_at = datetime.now()
            repo.status = RepositoryStatus.READY
            self._save_repositories()

            result = f"Successfully indexed {repo_info.full_name}\n{file_count} files processed\n{total_chunks} chunks indexed"
            logger.info("Indexing complete", result=result)
            yield result

        except ValidationError as e:
            logger.error("Validation error", error=str(e))
            yield f"Validation error: {str(e)}"
        except Exception as e:
            logger.error("Indexing failed", error=str(e), exc_info=True)
            if "repo" in locals():
                repo.status = RepositoryStatus.ERROR
                repo.error_message = str(e)
                self._save_repositories()
            yield f"Error: {str(e)}"

    def ask_question(
        self,
        repo_id: str,
        question: str,
        top_k: int = 5,
    ) -> tuple[str, str, str]:
        """Ask a question about a repository."""
        if not repo_id:
            return "", "", "Please select a repository"

        if not question.strip():
            return "", "", "Please enter a question"

        try:
            # Lazy load generator
            if self.generator is None:
                self.generator = ResponseGenerator()

            query = Query(
                question=question.strip(),
                repo_id=repo_id,
                top_k=int(top_k),
            )

            response = self.generator.generate(query)

            # Format answer
            answer_md = f"## Answer\n\n{response.answer}"
            if response.citations:
                answer_md += "\n\n### Citations\n"
                for citation in response.citations:
                    answer_md += f"- `{citation}`\n"

            # Format evidence
            evidence_md = response.format_evidence()

            status = "Grounded" if response.grounded else "Not grounded (no citations)"

            return answer_md, evidence_md, status

        except Exception as e:
            logger.error("Question failed", error=str(e))
            return "", "", f"Error: {str(e)}"

    def get_repositories(self):
        """Get list of repositories for dropdown."""
        import gradio as gr
        choices = []
        for repo in self.repositories.values():
            if repo.status == RepositoryStatus.READY:
                label = f"{repo.full_name} ({repo.chunk_count} chunks)"
                choices.append((label, repo.id))
        return gr.update(choices=choices)

    def get_repositories_table(self) -> list[list]:
        """Get repositories as table data."""
        rows = []
        for repo in self.repositories.values():
            rows.append([
                repo.id[:8],
                repo.full_name,
                repo.branch,
                repo.chunk_count,
                repo.status.value,
                repo.indexed_at.strftime("%Y-%m-%d %H:%M") if repo.indexed_at else "-",
            ])
        return rows

    def delete_repository(self, repo_id: str) -> tuple[str, list[list]]:
        """Delete a repository."""
        repo_id = repo_id.strip()

        # Find by full or partial ID
        found_repo = None
        for rid, repo in self.repositories.items():
            if rid == repo_id or rid.startswith(repo_id):
                found_repo = repo
                break

        if not found_repo:
            return "Repository not found", self.get_repositories_table()

        try:
            # Delete from vector store
            self.vectorstore.delete_repo_chunks(found_repo.id)

            # Delete cached repo
            self.loader.delete_cache(
                type("RepoInfo", (), {"owner": found_repo.owner, "name": found_repo.name})()
            )

            # Remove from records
            del self.repositories[found_repo.id]
            self._save_repositories()

            return f"Deleted {found_repo.full_name}", self.get_repositories_table()

        except Exception as e:
            logger.error("Delete failed", error=str(e))
            return f"Error: {str(e)}", self.get_repositories_table()