                                         

import argparse
import os
import sys
import logging
from typing import List, Dict, Any, Optional

                                                  
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

                                                  
                                                                        
                                                                  
try:
    from pydantic import BaseModel, Field
    import uuid
except ImportError:
    print("Error: pydantic library not found. Please run 'pip install pydantic'")
    sys.exit(1)

                                                                
class MockPydanticModel(BaseModel):
    class Config:
        extra = 'ignore'
class NLPFeatures(MockPydanticModel): pass
class InputPromptRecord(MockPydanticModel): pass
class DatabasePromptRecord(MockPydanticModel):
    prompt_id: uuid.UUID = Field(default_factory=uuid.uuid4)
    original_prompt: str = ""
    label: int = 0
    embedding: Optional[List[float]] = None

                                            
                                                                         
                                                           
                                             

                          
try:
                                                                
    from fortress.core.vector_store_interface import ChromaVectorStore
    from fortress.config import get_config
    
                                                                                                  
                                                               
    import fortress.common.data_models
    fortress.common.data_models.DatabasePromptRecord = DatabasePromptRecord
    fortress.common.data_models.NLPFeatures = NLPFeatures
    fortress.common.data_models.InputPromptRecord = InputPromptRecord

except ImportError as e:
    print("Error: A required module could not be imported.")
    print("\nPlease ensure you are running this script from the project's root directory.")
    print(f"Original error: {e}")
    sys.exit(1)

def get_directory_size(path: str) -> int:
    """Calculates the total size of a directory in bytes."""
    total_size = 0
    for dirpath, dirnames, filenames in os.walk(path):
        for f in filenames:
            fp = os.path.join(dirpath, f)
            if not os.path.islink(fp):
                total_size += os.path.getsize(fp)
    return total_size

def format_size(size_bytes: int) -> str:
    """Formats size in bytes to a human-readable string (KB, MB, GB)."""
    if size_bytes < 1024:
        return f"{size_bytes} Bytes"
    if size_bytes < 1024**2:
        return f"{size_bytes / 1024:.2f} KB"
    if size_bytes < 1024**3:
        return f"{size_bytes / (1024**2):.2f} MB"
    return f"{size_bytes / (1024**3):.2f} GB"

def main():
    """
    Main function to parse arguments and inspect the ChromaDB.
    """
    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
    parser = argparse.ArgumentParser(
        description="Inspect a ChromaDB folder to get the entry count and total size.",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument("db_path", type=str, help="The target folder path where the ChromaDB is located.")
    parser.add_argument("-c", "--collection-name", type=str, default=None, help="Specify the collection name to inspect.")
    args = parser.parse_args()

    if not os.path.isdir(args.db_path):
        print(f"Error: The specified path is not a valid directory: {args.db_path}")
        sys.exit(1)

    print(f"Inspecting database at: {os.path.abspath(args.db_path)}")
    if args.collection_name:
        print(f"Targeting collection: {args.collection_name}")

    try:
        print("Connecting to vector store...")
        vector_store = ChromaVectorStore(db_path=args.db_path, collection_name=args.collection_name)
        print("Fetching collection size...")
        num_entries = vector_store.get_collection_size()
        print("Calculating directory size...")
        total_size_bytes = get_directory_size(args.db_path)
        readable_size = format_size(total_size_bytes)
        print("\n--- Database Inspection Report ---")
        print(f"  Collection Name: {vector_store.collection_name}")
        print(f"  Total Entries:   {num_entries:,}")
        print(f"  Database Size:   {readable_size}")
        print("----------------------------------\n")
    except Exception as e:
        print(f"\nAn error occurred during inspection: {e}")
        print("Please ensure the path points to a valid ChromaDB directory.")
        sys.exit(1)

if __name__ == "__main__":
    main()
