import os
import logging
import shutil
import json
import hashlib
from typing import Optional

from rag_components import KnowledgeRAG
from utils import download_and_unzip_gdrive_folder, prepare_rag_source_directory, validate_rag_source_csv
from config import (
    GDRIVE_SOURCES_ENABLED, GDRIVE_FOLDER_ID_OR_URL, RAG_SOURCES_DIR, FAQ_SOURCES_DIR,
    RAG_STORAGE_PARENT_DIR, RAG_FAISS_INDEX_SUBDIR_NAME, RAG_LOAD_INDEX_ON_STARTUP,
    RAG_EMBEDDING_MODEL_NAME, RAG_EMBEDDING_USE_GPU, 
    RAG_CHUNK_SIZE, RAG_CHUNK_OVERLAP,
    RAG_RERANKER_MODEL_NAME, RAG_RERANKER_ENABLED,
    RAG_CHUNKED_SOURCES_FILENAME
)

logger = logging.getLogger(__name__)


def _approved_source_signature(source_dir: str, expected_kind: str) -> dict:
    """Tie each saved FAISS index to the approved CSV contents and filter version."""
    expected_name = "faqs.csv" if expected_kind == "faq" else "quiz.csv"
    file_path = os.path.join(source_dir, expected_name)
    allowed, _ = validate_rag_source_csv(file_path, expected_name, "faq" if expected_kind == "faq" else "quiz")
    signature = {"source_filter_version": 1, "kind": expected_kind, "files": {}}
    if allowed:
        digest = hashlib.sha256()
        with open(file_path, "rb") as handle:
            for block in iter(lambda: handle.read(1024 * 1024), b""):
                digest.update(block)
        signature["files"][expected_name] = digest.hexdigest()
    if expected_kind == "url":
        url_text_path = os.path.join(source_dir, "url_data.txt")
        if os.path.isfile(url_text_path):
            digest = hashlib.sha256()
            with open(url_text_path, "rb") as handle:
                for block in iter(lambda: handle.read(1024 * 1024), b""):
                    digest.update(block)
            signature["files"]["url_data.txt"] = digest.hexdigest()
    return signature


def initialize_and_get_rag_system(force_rebuild: bool = False, source_dir_override: Optional[str] = None, storage_dir_override: Optional[str] = None) -> Optional[KnowledgeRAG]:
    
    logger.info("[RAG_SYSTEM_INIT] Initializing...")
    source_dir_to_use = source_dir_override if source_dir_override else RAG_SOURCES_DIR
    source_kind = ("faq" if source_dir_override is not None and os.path.abspath(source_dir_override) == os.path.abspath(FAQ_SOURCES_DIR)
                   else "url" if source_dir_override is not None else "quiz")
    storage_dir_to_use = storage_dir_override if storage_dir_override else RAG_STORAGE_PARENT_DIR
    
    # GDrive Logic 
    if GDRIVE_SOURCES_ENABLED and not source_dir_override and GDRIVE_FOLDER_ID_OR_URL:
        logger.info("[RAG_SYSTEM_INIT] Downloading sources from GDrive...")
        import tempfile
        staged_sources = tempfile.mkdtemp(prefix="rag_gdrive_download_")
        try:
            if download_and_unzip_gdrive_folder(GDRIVE_FOLDER_ID_OR_URL, staged_sources):
                if os.path.exists(RAG_SOURCES_DIR):
                    shutil.rmtree(RAG_SOURCES_DIR)
                shutil.move(staged_sources, RAG_SOURCES_DIR)
            else:
                logger.error("[RAG_SYSTEM_INIT] Google Drive folder download failed; local sources retained.")
        finally:
            if os.path.isdir(staged_sources):
                shutil.rmtree(staged_sources)

    if source_kind == "quiz":
        prepare_rag_source_directory(RAG_SOURCES_DIR, FAQ_SOURCES_DIR)

    faiss_index_path = os.path.join(storage_dir_to_use, RAG_FAISS_INDEX_SUBDIR_NAME)
    manifest_path = os.path.join(storage_dir_to_use, "approved_sources_manifest.json")
    current_signature = _approved_source_signature(source_dir_to_use, source_kind)
    manifest_matches = False
    try:
        with open(manifest_path, "r", encoding="utf-8") as handle:
            manifest_matches = json.load(handle) == current_signature
    except (OSError, ValueError):
        pass
    if not manifest_matches:
        logger.info("[SOURCE_FILTER] %s source changed or prior index unverified; rebuilding to exclude disallowed files.", source_kind.upper())
        force_rebuild = True
    pre_chunked_file = os.path.join(storage_dir_to_use, RAG_CHUNKED_SOURCES_FILENAME)

    if force_rebuild:
        if os.path.exists(faiss_index_path):
            logger.info("[RAG_SYSTEM_INIT] Force rebuild: deleting old index.")
            shutil.rmtree(faiss_index_path)
        if os.path.exists(pre_chunked_file):
            logger.info("[RAG_SYSTEM_INIT] Force rebuild: deleting old pre-chunked file.")
            os.remove(pre_chunked_file)
        if os.path.exists(manifest_path):
            os.remove(manifest_path)

    try:
        rag = KnowledgeRAG(
            index_storage_dir=storage_dir_to_use, 
            embedding_model_name=RAG_EMBEDDING_MODEL_NAME,
            use_gpu_for_embeddings=RAG_EMBEDDING_USE_GPU,
            chunk_size=RAG_CHUNK_SIZE,
            chunk_overlap=RAG_CHUNK_OVERLAP,
            reranker_model_name=RAG_RERANKER_MODEL_NAME,
            enable_reranker=RAG_RERANKER_ENABLED,
        )

        rag.allowed_source_kind = source_kind
        loaded = False
        if RAG_LOAD_INDEX_ON_STARTUP and not force_rebuild:
            if rag.chunk_config_has_changed():
                logger.warning("[RAG_SYSTEM_INIT] Chunk config changed — forcing index rebuild.")
                if os.path.exists(faiss_index_path):
                    shutil.rmtree(faiss_index_path)
                if os.path.exists(pre_chunked_file):
                    os.remove(pre_chunked_file)
            else:
                try:
                    rag.load_index_from_disk()
                    loaded = True
                except Exception as e:
                    logger.warning(f"[RAG_SYSTEM_INIT] Load failed ({e}). Building new.")

        if not loaded:
            if not current_signature["files"]:
                logger.warning(
                    f"[RAG_SYSTEM_INIT] No supported source files found in {source_dir_to_use}. "
                    "System empty for this index."
                )
            else:
                rag.build_index_from_source_files(source_dir_to_use)
                os.makedirs(storage_dir_to_use, exist_ok=True)
                with open(manifest_path, "w", encoding="utf-8") as handle:
                    json.dump(current_signature, handle, sort_keys=True)

        logger.info("[RAG_SYSTEM_INIT] Complete.")
        return rag

    except Exception as e:
        logger.critical(f"[RAG_SYSTEM_INIT] FATAL: {e}", exc_info=True)
        return None