import os
from dotenv import load_dotenv
from utils.document_management import load_document, transcript_chunker, merge_chunks
from utils.vector_db import VectorDB
from langchain.schema import Document
from uuid import uuid4
from pathlib import Path
from typing import Optional

def build_vector_db_for_dir(
    data_dir: str,
    collection_name: str = "custom_collection",
    embedding_length: int = 1536,
    file_suffix: Optional[str] = None,
    openai_api_key: Optional[str] = None,
    verbose: bool = True,
):
    """
    Build a new vector database for all files in a directory.
    Args:
        data_dir: Directory containing files to process.
        collection_name: Name for the vector DB collection.
        embedding_length: Embedding vector length.
        file_suffix: Only process files with this suffix (e.g. '.srt', '.md'). If None, process all files.
        openai_api_key: API key for embedding (if required by backend).
        verbose: Print progress.
    """
    load_dotenv()
    working_dir = os.path.abspath(data_dir)
    if not os.path.isdir(working_dir):
        raise ValueError(f"Directory not found: {working_dir}")

    all_documents = []
    all_files = []
    for root, dirs, files in os.walk(working_dir):
        for file in files:
            if file_suffix is None or file.endswith(file_suffix):
                file_path = os.path.join(root, file)
                relative_path = os.path.relpath(file_path, working_dir)
                all_files.append(relative_path)

    for file_path in all_files:
        path = os.path.join(working_dir, file_path)
        doc = load_document(path)
        if doc is None:
            if verbose:
                print(f"Skipping file {file_path} as it could not be loaded.")
            continue
        # Try chunking if transcript_chunker is available, else treat as single doc
        try:
            chunks = transcript_chunker(doc, filename=file_path)
            merged_chunks = merge_chunks(chunks)
            documents = [
                Document(
                    page_content=chunk["text"],
                    metadata={**chunk, "chunk_id": i, "id": str(uuid4()), "filename": file_path}
                )
                for i, chunk in enumerate(merged_chunks)
            ]
        except Exception:
            # Fallback: treat as single document
            documents = [
                Document(
                    page_content=doc if isinstance(doc, str) else str(doc),
                    metadata={"id": str(uuid4()), "filename": file_path}
                )
            ]
        if verbose:
            print(f"Processing {file_path} with {len(documents)} chunks")
        all_documents.extend(documents)

    # Remove existing vector DB if present
    chroma_db_path = os.path.join(working_dir, "chroma_db")
    if all_documents:
        if os.path.exists(chroma_db_path):
            import shutil
            shutil.rmtree(chroma_db_path)
            if verbose:
                print("Removed existing vector database directory.")
        vector_db = VectorDB(
            collection_name=collection_name,
            embedding_length=embedding_length,
            open_ai_api_key=openai_api_key or os.getenv("OPENAI_API_KEY"),
            working_dir=working_dir,
            documents=all_documents
        )
        if verbose:
            print(f"Created new vector database with {len(all_documents)} documents")
    else:
        if verbose:
            print("No documents found to process")
    if verbose:
        print("Vector database reset and data upload complete with ChromaDB.")