code_chatbot/universal_ingestor.py

"""Universal ingestor that handles multiple input types: ZIP files, GitHub URLs, local directories, etc."""

import logging
import os
import zipfile
import requests
import tempfile
import shutil
from abc import ABC, abstractmethod
from typing import Any, Dict, Generator, Tuple, Optional
from urllib.parse import urlparse
from pathlib import Path

from langchain_core.documents import Document

logger = logging.getLogger(__name__)


class DataManager(ABC):
    """Abstract base class for data managers."""
    
    def __init__(self, dataset_id: str):
        self.dataset_id = dataset_id
    
    @abstractmethod
    def download(self) -> bool:
        """Downloads/prepares the data."""
        pass
    
    @abstractmethod
    def walk(self, get_content: bool = True) -> Generator[Tuple[Any, Dict], None, None]:
        """Yields (content, metadata) tuples for each file."""
        pass


class UniversalIngestor(DataManager):
    """Factory class to ingest data from various sources."""
    
    def __init__(self, source: str, local_dir: Optional[str] = None, **kwargs):
        """
        Args:
            source: Can be:
                - GitHub URL (e.g., "https://github.com/owner/repo")
                - GitHub repo ID (e.g., "owner/repo")
                - Local directory path
                - ZIP file path
                - Web URL
            local_dir: Directory to store/clone/download data
            **kwargs: Additional arguments for specific managers
        """
        super().__init__(dataset_id=source)
        self.source = source
        self.kwargs = kwargs
        self.local_dir = local_dir or os.path.join(tempfile.gettempdir(), "code_chatbot")
        self.delegate = self._detect_handler()
    
    def _detect_handler(self) -> DataManager:
        """Detects the type of input and returns the appropriate handler."""
        source = self.source.strip()
        
        # Check if it's a URL
        if self._is_url(source):
            if "github.com" in source or source.count("/") == 1 and "/" in source:
                # GitHub URL or repo ID (owner/repo)
                if "github.com" in source:
                    # Extract repo_id from URL
                    parts = urlparse(source).path.strip("/").split("/")
                    if len(parts) >= 2:
                        repo_id = f"{parts[0]}/{parts[1]}"
                    else:
                        raise ValueError(f"Invalid GitHub URL: {source}")
                else:
                    # Assume it's owner/repo format
                    repo_id = source
                
                return GitHubRepoManager(
                    repo_id=repo_id,
                    local_dir=self.local_dir,
                    **self.kwargs
                )
            
            # Other web URLs
            return WebDocManager(source, local_dir=self.local_dir)
        
        # Check if it's a ZIP file
        if source.lower().endswith('.zip') and os.path.isfile(source):
            return ZIPFileManager(source, local_dir=self.local_dir)
        
        # Check if it's a local directory
        if os.path.isdir(source):
            return LocalDirectoryManager(source)
        
        # Check if it's a local file
        if os.path.isfile(source):
            return LocalFileManager(source)
        
        raise ValueError(f"Unable to determine source type for: {source}")
    
    def _is_url(self, s: str) -> bool:
        """Checks if a string is a URL."""
        try:
            result = urlparse(s)
            return bool(result.scheme and result.netloc)
        except Exception:
            # Check if it looks like owner/repo (GitHub format)
            if "/" in s and s.count("/") == 1 and not os.path.exists(s):
                return True
            return False
    
    @property
    def local_path(self) -> str:
        """Returns the local path where data is stored."""
        if hasattr(self.delegate, "local_path"):
            return self.delegate.local_path
        if hasattr(self.delegate, "path"):
            return self.delegate.path
        return self.local_dir
    
    def download(self) -> bool:
        """Downloads/prepares the data."""
        return self.delegate.download()
    
    def walk(self, get_content: bool = True) -> Generator[Tuple[Any, Dict], None, None]:
        """Yields (content, metadata) tuples."""
        yield from self.delegate.walk(get_content)


class ZIPFileManager(DataManager):
    """Handles ZIP file ingestion."""
    
    def __init__(self, zip_path: str, local_dir: str):
        super().__init__(dataset_id=zip_path)
        self.zip_path = zip_path
        self.local_dir = local_dir
        self.path = os.path.join(local_dir, "extracted", os.path.basename(zip_path).replace('.zip', ''))
    
    def download(self) -> bool:
        """Extracts the ZIP file."""
        if os.path.exists(self.path):
            logger.info(f"ZIP already extracted to {self.path}")
            return True
        
        os.makedirs(self.path, exist_ok=True)
        
        try:
            with zipfile.ZipFile(self.zip_path, 'r') as zip_ref:
                zip_ref.extractall(self.path)
            logger.info(f"Extracted {self.zip_path} to {self.path}")
            return True
        except Exception as e:
            logger.error(f"Failed to extract ZIP: {e}")
            return False
    
    def walk(self, get_content: bool = True) -> Generator[Tuple[Any, Dict], None, None]:
        """Walks extracted files."""
        if not os.path.exists(self.path):
            return
        
        IGNORE_DIRS = {'__pycache__', '.git', 'node_modules', 'venv', '.venv', '.env'}
        IGNORE_EXTENSIONS = {
            '.pyc', '.png', '.jpg', '.jpeg', '.gif', '.ico', '.svg', '.mp4', '.mov', 
            '.zip', '.tar', '.gz', '.pdf', '.exe', '.bin', '.pkl', '.npy', '.pt', '.pth'
        }
        
        for root, dirs, files in os.walk(self.path):
            dirs[:] = [d for d in dirs if d not in IGNORE_DIRS and not d.startswith('.')]
            
            for file in files:
                if file.startswith('.'):
                    continue
                
                file_path = os.path.join(root, file)
                _, ext = os.path.splitext(file)
                if ext.lower() in IGNORE_EXTENSIONS:
                    continue
                
                rel_path = os.path.relpath(file_path, self.path)
                
                if get_content:
                    try:
                        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                            content = f.read()
                        yield content, {
                            "file_path": file_path,
                            "source": rel_path,
                            "file_name": file
                        }
                    except Exception as e:
                        logger.warning(f"Failed to read {file_path}: {e}")
                else:
                    yield {"file_path": file_path, "source": rel_path, "file_name": file}


class LocalDirectoryManager(DataManager):
    """Handles local directory ingestion."""
    
    def __init__(self, path: str):
        super().__init__(dataset_id=path)
        self.path = path
        self.local_dir = path
    
    def download(self) -> bool:
        return os.path.isdir(self.path)
    
    def walk(self, get_content: bool = True) -> Generator[Tuple[Any, Dict], None, None]:
        """Walks local directory."""
        IGNORE_DIRS = {'__pycache__', '.git', 'node_modules', 'venv', '.venv', '.env'}
        IGNORE_EXTENSIONS = {
            '.pyc', '.png', '.jpg', '.jpeg', '.gif', '.ico', '.svg', '.mp4', '.mov', 
            '.zip', '.tar', '.gz', '.pdf', '.exe', '.bin', '.pkl', '.npy', '.pt', '.pth'
        }
        
        for root, dirs, files in os.walk(self.path):
            dirs[:] = [d for d in dirs if d not in IGNORE_DIRS and not d.startswith('.')]
            
            for file in files:
                if file.startswith('.'):
                    continue
                
                file_path = os.path.join(root, file)
                _, ext = os.path.splitext(file)
                if ext.lower() in IGNORE_EXTENSIONS:
                    continue
                
                rel_path = os.path.relpath(file_path, self.path)
                
                if get_content:
                    try:
                        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                            content = f.read()
                        yield content, {
                            "file_path": file_path,
                            "source": rel_path,
                            "url": f"file://{file_path}"
                        }
                    except Exception as e:
                        logger.warning(f"Skipping {file_path}: {e}")
                else:
                    yield {"file_path": file_path, "source": rel_path}


class LocalFileManager(DataManager):
    """Handles single file ingestion."""
    
    def __init__(self, path: str):
        super().__init__(dataset_id=path)
        self.path = path
    
    def download(self) -> bool:
        return os.path.exists(self.path)
    
    def walk(self, get_content: bool = True) -> Generator[Tuple[Any, Dict], None, None]:
        """Yields the single file."""
        if get_content:
            try:
                with open(self.path, 'r', encoding='utf-8', errors='ignore') as f:
                    content = f.read()
                yield content, {"file_path": self.path, "source": os.path.basename(self.path)}
            except Exception as e:
                logger.error(f"Failed to read {self.path}: {e}")
        else:
            yield {"file_path": self.path, "source": os.path.basename(self.path)}


class GitHubRepoManager(DataManager):
    """Handles GitHub repository cloning and ingestion."""
    
    def __init__(self, repo_id: str, local_dir: str, access_token: Optional[str] = None, commit_hash: Optional[str] = None):
        """
        Args:
            repo_id: GitHub repo in format "owner/repo"
            local_dir: Directory to clone to
            access_token: GitHub token for private repos
            commit_hash: Optional commit hash to checkout
        """
        super().__init__(dataset_id=repo_id)
        self.repo_id = repo_id
        self.local_dir = local_dir
        self.access_token = access_token or os.getenv("GITHUB_TOKEN")
        self.commit_hash = commit_hash
        self.path = os.path.join(local_dir, repo_id.replace("/", "_"))
    
    def download(self) -> bool:
        """Clones the GitHub repository."""
        if os.path.exists(self.path) and os.listdir(self.path):
            logger.info(f"Repo already cloned at {self.path}")
            return True
        
        try:
            from git import Repo, GitCommandError
            
            if self.access_token:
                clone_url = f"https://{self.access_token}@github.com/{self.repo_id}.git"
            else:
                clone_url = f"https://github.com/{self.repo_id}.git"
            
            os.makedirs(self.local_dir, exist_ok=True)
            
            if self.commit_hash:
                repo = Repo.clone_from(clone_url, self.path)
                repo.git.checkout(self.commit_hash)
            else:
                Repo.clone_from(clone_url, self.path, depth=1, single_branch=True)
            
            logger.info(f"Cloned {self.repo_id} to {self.path}")
            return True
        except ImportError:
            logger.error("GitPython not installed. Install with: pip install gitpython")
            raise
        except Exception as e:
            logger.error(f"Failed to clone {self.repo_id}: {e}")
            return False
    
    def walk(self, get_content: bool = True) -> Generator[Tuple[Any, Dict], None, None]:
        """Walks cloned repository."""
        if not os.path.exists(self.path):
            return
        
        # Use LocalDirectoryManager logic
        manager = LocalDirectoryManager(self.path)
        yield from manager.walk(get_content)


class WebDocManager(DataManager):
    """Handles web page/document ingestion."""
    
    def __init__(self, url: str, local_dir: str):
        super().__init__(dataset_id=url)
        self.url = url
        self.local_dir = local_dir
    
    def download(self) -> bool:
        """Checks if URL is accessible."""
        try:
            response = requests.get(self.url, timeout=10)
            return response.status_code == 200
        except Exception as e:
            logger.error(f"Could not reach {self.url}: {e}")
            return False
    
    def walk(self, get_content: bool = True) -> Generator[Tuple[Any, Dict], None, None]:
        """Fetches web page content."""
        try:
            response = requests.get(self.url, timeout=10)
            if get_content:
                from bs4 import BeautifulSoup
                soup = BeautifulSoup(response.content, 'html.parser')
                text = soup.get_text(separator='\n')
                yield text, {"file_path": self.url, "url": self.url, "source": "web"}
            else:
                yield {"file_path": self.url, "url": self.url, "source": "web"}
        except Exception as e:
            logger.error(f"Failed to fetch {self.url}: {e}")


def process_source(source: str, extract_to: str) -> Tuple[list, str]:
    """
    Convenience function to process any source type and return documents + local path.
    
    Returns:
        Tuple of (documents, local_path)
    """
    ingestor = UniversalIngestor(source, local_dir=extract_to)
    
    if not ingestor.download():
        raise ValueError(f"Failed to download/prepare source: {source}")
    
    documents = []
    for content, metadata in ingestor.walk(get_content=True):
        documents.append(Document(
            page_content=content,
            metadata=metadata
        ))
    
    return documents, ingestor.local_path