utils.py

import os
import logging
import re
import shutil
import tempfile
import zipfile
import requests
import csv
from typing import Optional, List, Tuple, Dict, Any
from bs4 import BeautifulSoup

import gdown
from pypdf import PdfReader
import docx as python_docx

logger = logging.getLogger(__name__)

def extract_text_from_file(file_path: str, file_type: str) -> Optional[str]:
    logger.info(f"[TEXT_EXTRACTION] Starting extraction from {file_type.upper()} file: {file_path}")
    text_content = None
    try:
        if file_type == 'pdf':
            reader = PdfReader(file_path)
            text_content = "".join(page.extract_text() + "\n" for page in reader.pages if page.extract_text())
        elif file_type == 'docx':
            doc = python_docx.Document(file_path)
            text_content = "\n".join(para.text for para in doc.paragraphs if para.text)
        elif file_type == 'txt':
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                text_content = f.read()
        else:
            logger.warning(f"[TEXT_EXTRACTION] Unsupported file type: {file_type}")
            return None

        if not text_content or not text_content.strip():
            logger.warning(f"[TEXT_EXTRACTION] No text content extracted from {file_path}")
            return None
        
        return text_content.strip()
    except Exception as e:
        logger.error(f"[TEXT_EXTRACTION] Error extracting text: {e}", exc_info=True)
        return None

FAISS_RAG_SUPPORTED_EXTENSIONS = {
    'pdf': lambda path: extract_text_from_file(path, 'pdf'),
    'docx': lambda path: extract_text_from_file(path, 'docx'),
    'txt': lambda path: extract_text_from_file(path, 'txt'),
    'csv': lambda path: "CSV_HANDLED_NATIVELY" # Bypassed directly in components to allow row-by-row chunks
}

def fetch_and_clean_url(url: str, output_txt_path: str) -> bool:
    """Fetches HTML from a URL, cleans it, and saves it as pure text."""
    try:
        logger.info(f"[URL_FETCH] Fetching and cleaning data from: {url}")
        headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64)'}
        response = requests.get(url, headers=headers, timeout=15)
        response.raise_for_status()
        
        soup = BeautifulSoup(response.text, 'html.parser')
        
        # Remove non-content tags completely
        for tag in soup(["script", "style", "nav", "footer", "header", "noscript"]):
            tag.decompose()
        
        # Attempt to find specific content divs if they exist (generalizes rentry and others)
        entry = soup.find('div', class_='entry-text')
        if entry:
            text = entry.get_text(separator='\n', strip=True)
        else:
            # Fallback for other sites or layout changes
            text = soup.body.get_text(separator='\n', strip=True) if soup.body else soup.get_text(separator='\n', strip=True)
        
        # Clean up excessive blank lines
        text = re.sub(r'\n\s*\n', '\n\n', text)
        
        # Ensure directory exists
        os.makedirs(os.path.dirname(output_txt_path), exist_ok=True)
        
        with open(output_txt_path, 'w', encoding='utf-8') as f:
            f.write(text)
            
        logger.info(f"[URL_FETCH] Success! Text saved to {output_txt_path}")
        return True
    except Exception as e:
        logger.error(f"[URL_FETCH] Error fetching/cleaning URL {url}: {e}", exc_info=True)
        return False

def get_id_from_gdrive_input(url_or_id: str) -> Optional[str]:
    if not url_or_id: return None
    match_folder = re.search(r"/folders/([a-zA-Z0-9_-]+)", url_or_id)
    if match_folder: return match_folder.group(1)
    match_file_d = re.search(r"/d/([a-zA-Z0-9_-]+)", url_or_id)
    if match_file_d: return match_file_d.group(1)
    match_uc = re.search(r"id=([a-zA-Z0-9_-]+)", url_or_id)
    if match_uc: return match_uc.group(1)
    return url_or_id if len(url_or_id) > 10 else None

def download_gdrive_file(file_id_or_url: str, target_path: str) -> bool:
    logger.info(f"[GDRIVE_SINGLE] Downloading file. Input: {file_id_or_url}")
    file_id = get_id_from_gdrive_input(file_id_or_url)
    if not file_id: return False

    try:
        os.makedirs(os.path.dirname(target_path), exist_ok=True)
        gdown.download(id=file_id, output=target_path, quiet=False, fuzzy=True)
        
        if os.path.exists(target_path) and os.path.getsize(target_path) > 0:
            return True
        return False
    except Exception as e:
        logger.error(f"[GDRIVE_SINGLE] Error: {e}", exc_info=True)
        return False

def download_and_unzip_gdrive_folder(folder_id_or_url: str, target_dir_for_contents: str) -> bool:
    logger.info(f"[GDRIVE] Downloading folder. Input: {folder_id_or_url}")
    folder_id = get_id_from_gdrive_input(folder_id_or_url)
    if not folder_id: return False

    temp_dir = tempfile.mkdtemp()
    try:
        downloaded = gdown.download_folder(id=folder_id, output=temp_dir, quiet=False, use_cookies=False)
        if not downloaded:
            logger.error("[GDRIVE] Folder download did not return any completed files.")
            return False
        if not os.path.exists(target_dir_for_contents):
            os.makedirs(target_dir_for_contents)

        downloaded_items = os.listdir(temp_dir)
        if not downloaded_items:
            logger.error("[GDRIVE] Folder download returned no files. Existing sources will be retained.")
            return False
        src_root = temp_dir

        # gdown often creates a single wrapper directory for the downloaded folder.
        # Flatten that wrapper for normal source folders, but preserve app_faqs/ so
        # sources/app_faqs/faqs.csv stays separate from the main knowledgebase.
        if len(downloaded_items) == 1:
            single_item = downloaded_items[0]
            single_path = os.path.join(temp_dir, single_item)
            if os.path.isdir(single_path) and single_item != "app_faqs":
                src_root = single_path

        for item in os.listdir(src_root):
            shutil.move(os.path.join(src_root, item), os.path.join(target_dir_for_contents, item))
        return True
    except Exception as e:
        logger.error(f"[GDRIVE] Error: {e}", exc_info=True)
        return False
    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)

def download_and_unzip_gdrive_file(file_id_or_url: str, target_extraction_dir: str) -> bool:
    logger.info(f"[GDRIVE_ZIP] Downloading ZIP. Input: {file_id_or_url}")
    file_id = get_id_from_gdrive_input(file_id_or_url)
    if not file_id: return False

    temp_zip = os.path.join(tempfile.gettempdir(), "temp_download.zip")
    try:
        gdown.download(id=file_id, output=temp_zip, quiet=False)
        with zipfile.ZipFile(temp_zip, 'r') as zip_ref:
            zip_ref.extractall(target_extraction_dir)
        return True
    except Exception as e:
        logger.error(f"[GDRIVE_ZIP] Error: {e}", exc_info=True)
        return False

# The shared Google Drive folder may contain other datasets (including tf.csv).
# Only these two named CSV datasets are permitted in the RAG indexes.
RAG_APPROVED_CSV_FILES = {"quiz.csv": "quiz", "faqs.csv": "faq"}


def validate_rag_source_csv(file_path: str, filename: str, expected_kind: str) -> Tuple[bool, str]:
    """Validate both filename and CSV schema before it can enter a RAG index."""
    if not os.path.isfile(file_path):
        return False, "not a regular file"
    expected_filename = "quiz.csv" if expected_kind == "quiz" else "faqs.csv"
    if filename != expected_filename:
        return False, f"only {expected_filename} is permitted in the {expected_kind} index"
    try:
        with open(file_path, "r", encoding="utf-8-sig", newline="") as handle:
            reader = csv.reader(handle)
            headers = next(reader, [])
            normalized = {str(h).strip().lower() for h in headers}
            if expected_kind == "quiz":
                if not {"question", "answer"}.issubset(normalized):
                    return False, "quiz CSV requires question and answer columns"
                options = {"opt_1", "opt_2", "opt_3", "opt_4"}
                if options.intersection(normalized) and not options.issubset(normalized):
                    return False, "quiz option columns must include all opt_1 through opt_4"
            else:
                if not ({"q", "a"}.issubset(normalized) or {"question", "answer"}.issubset(normalized)):
                    return False, "FAQ CSV requires Q,A or question,answer columns"
            if next(reader, None) is None:
                return False, "CSV has no data rows"
    except (OSError, UnicodeError, csv.Error) as exc:
        return False, f"unreadable CSV: {exc}"
    return True, "valid CSV schema"


def prepare_rag_source_directory(source_dir: str, faq_dir: str) -> None:
    """Route root-level faqs.csv to the existing separate FAQ index; log every candidate."""
    os.makedirs(source_dir, exist_ok=True)
    os.makedirs(faq_dir, exist_ok=True)
    faq_in_root = os.path.join(source_dir, "faqs.csv")
    faq_target = os.path.join(faq_dir, "faqs.csv")
    if os.path.abspath(faq_in_root) != os.path.abspath(faq_target) and os.path.isfile(faq_in_root):
        allowed, reason = validate_rag_source_csv(faq_in_root, "faqs.csv", "faq")
        if allowed:
            try:
                os.replace(faq_in_root, faq_target)
            except OSError as exc:
                logger.error("[SOURCE_FILTER] REJECTED faqs.csv: failed to route FAQ source: %s", exc)
        else:
            logger.warning("[SOURCE_FILTER] REJECTED faqs.csv: %s", reason)

    for name in sorted(os.listdir(source_dir)):
        path = os.path.join(source_dir, name)
        if os.path.abspath(path) == os.path.abspath(faq_dir) and os.path.isdir(path):
            continue
        allowed, reason = validate_rag_source_csv(path, name, "quiz")
        if allowed:
            logger.info("[SOURCE_FILTER] ACCEPTED %s -> QUIZ index (%s)", name, reason)
        else:
            logger.warning("[SOURCE_FILTER] REJECTED %s -> skipped (%s)", name, reason)

    if os.path.abspath(faq_dir) != os.path.abspath(source_dir):
        for name in sorted(os.listdir(faq_dir)):
            path = os.path.join(faq_dir, name)
            allowed, reason = validate_rag_source_csv(path, name, "faq")
            if allowed:
                logger.info("[SOURCE_FILTER] ACCEPTED %s -> FAQ index (%s)", name, reason)
            else:
                logger.warning("[SOURCE_FILTER] REJECTED FAQ/%s -> skipped (%s)", name, reason)

def extract_csv_chunks(file_path: str, filename: str, rows_per_chunk: int) -> List[Tuple[str, Dict[str, Any]]]:
    """Centralized logic to extract chunk contents and metadata from a CSV file."""
    rows_per_chunk = max(1, int(rows_per_chunk))
    chunks: List[Tuple[str, Dict[str, Any]]] = []
    row_buffer: List[str] = []
    row_start = 1
    chunk_index = 0

    def flush_buffer() -> None:
        nonlocal row_buffer, row_start, chunk_index

        if not row_buffer:
            return

        row_end = row_start + len(row_buffer) - 1
        if len(row_buffer) == 1:
            full_location = f"{filename}, Row {row_start}"
        else:
            full_location = f"{filename}, Rows {row_start}-{row_end}"

        metadata = {
            "source_document_name": filename,
            "chunk_index": chunk_index,
            "source_type": "csv",
            "full_location": full_location,
            "row_start": row_start,
            "row_end": row_end,
            "rows_in_chunk": len(row_buffer),
        }

        chunks.append(("\n\n".join(row_buffer), metadata))
        row_buffer = []
        chunk_index += 1

    with open(file_path, mode='r', encoding='utf-8-sig') as f:
        reader = csv.reader(f)
        
        try:
            first_row = next(reader)
        except StopIteration:
            first_row = None
        
        if first_row:
            is_header = all(len(str(h).strip()) <= 50 and not str(h).strip().startswith("Q.") for h in first_row)
            headers = first_row if is_header else None
            
            data_row_number = 1
            
            normalized_headers = [str(header).strip().lower() for header in headers] if headers else []
            question_index = normalized_headers.index("question") if "question" in normalized_headers else None
            answer_index = normalized_headers.index("answer") if "answer" in normalized_headers else None
            is_question_answer_csv = question_index is not None and answer_index is not None
            # Quiz CSVs use a 1-based answer index into the four opt_N columns.
            # FAQ question/answer CSVs retain their existing plain-text behavior.
            quiz_option_headers = ("opt_1", "opt_2", "opt_3", "opt_4")
            is_quiz_csv = is_question_answer_csv and all(
                name in normalized_headers for name in quiz_option_headers
            )
            quiz_option_indices = (
                [normalized_headers.index(name) for name in quiz_option_headers]
                if is_quiz_csv else []
            )

            def process_row_list(r):
                if is_quiz_csv:
                    question = str(r[question_index]).strip() if question_index < len(r) else ""
                    answer_key = str(r[answer_index]).strip() if answer_index < len(r) else ""
                    if not question or answer_key not in ("1", "2", "3", "4"):
                        logger.warning("[CSV_QUIZ] Skipping row with missing question or invalid answer key in %s", filename)
                        return ""
                    correct_option_index = quiz_option_indices[int(answer_key) - 1]
                    correct_answer = (
                        str(r[correct_option_index]).strip()
                        if correct_option_index < len(r) else ""
                    )
                    if not correct_answer:
                        logger.warning("[CSV_QUIZ] Skipping row with empty correct option in %s", filename)
                        return ""
                    return f"question: {question}\nanswer: {correct_answer}"

                if is_question_answer_csv:
                    # Quiz RAG CSVs are indexed row by row using only question and answer.
                    # Columns such as id are intentionally excluded from the chunk content.
                    question = str(r[question_index]).strip() if question_index < len(r) else ""
                    answer = str(r[answer_index]).strip() if answer_index < len(r) else ""
                    parts = []
                    if question:
                        parts.append(f"question: {question}")
                    if answer:
                        parts.append(f"answer: {answer}")
                    return "\n".join(parts)

                if is_header:
                    return "\n".join([
                        f"{headers[i]}: {v}"
                        for i, v in enumerate(r)
                        if i < len(headers) and headers[i] and v and str(v).strip()
                    ])

                return "\n".join([str(v) for v in r if v and str(v).strip()])

            if not is_header:
                row_text = process_row_list(first_row)
                if row_text:
                    row_buffer.append(row_text)
                    if len(row_buffer) >= rows_per_chunk:
                        flush_buffer()
                data_row_number += 1

            for row in reader:
                row_text = process_row_list(row)
                if not row_text:
                    data_row_number += 1
                    continue

                if not row_buffer:
                    row_start = data_row_number

                row_buffer.append(row_text)

                if len(row_buffer) >= rows_per_chunk:
                    flush_buffer()
                    
                data_row_number += 1

            flush_buffer()

    return chunks