#!/usr/bin/env python3
"""
Intelligent Phi-Powered RAG Agent with AWS S3 Integration
===========================================================
Advanced Document Intelligence System with Vector Search and Intelligent Caching
Author: Damilare Lekan Adekeye
"""

import sys
import os
import json
import boto3
import time
import re
from openai import OpenAI
from typing import List
from phi.agent import Agent
from phi.model.openai import OpenAIChat
from phi.vectordb.lancedb import LanceDb, SearchType
from phi.playground import serve_playground_app
from phi.embedder.openai import OpenAIEmbedder
from phi.storage.agent.sqlite import SqlAgentStorage
from phi.knowledge.pdf import PDFUrlKnowledgeBase
from google.colab import drive

# ==========================
# MOUNT GOOGLE DRIVE
# ==========================
drive.mount('/content/drive')

# ==========================
# CACHE MANAGEMENT
# ==========================
cache_path = "/content/cached_knowledge_base"
drive_cache_path = f"/content/drive/MyDrive/Colab Notebooks/cached_knowledge_base"

# Check if cached files exist in Drive
if os.path.exists(drive_cache_path):
    print("Restoring cached knowledge base from Google Drive...")
    os.system('cp -r /content/drive/MyDrive/Colab\ Notebooks/cached_knowledge_base /content/')
    os.system('cp -r /content/drive/MyDrive/Colab\ Notebooks/cached_metadata.json /content/')
    print("Restored.")
else:
    print("No cached knowledge base found in Drive. It will be created from scratch.")
    # Ensure the cache is saved in Drive if not saved in Drive.
    if os.path.exists(cache_path):
        print("Saving updated cached knowledge base to Google Drive...")
        os.system('rm -rf /content/drive/MyDrive/Colab\ Notebooks/cached_knowledge_base')
        os.system('rm -rf /content/drive/MyDrive/Colab\ Notebooks/cached_metadata.json')
        os.system('cp -r /content/cached_knowledge_base /content/drive/MyDrive/Colab\ Notebooks/')
        os.system('cp -r /content/cached_metadata.json /content/drive/MyDrive/Colab\ Notebooks/')
    else:
        print("No cache found to save in colab.")

# ==========================
# CONFIGURATION PARAMETERS
# ==========================

# AWS Configuration
AWS_S3_BUCKET = "lbb-and-dataforseo-pdfs"
AWS_REGION = "eu-north-1"
PDF_FOLDER = "folder-lbb-and-dataforseo-pdfs/"

# Local cache configuration
CACHE_METADATA_FILE = "cached_metadata.json"
LANCEDB_URI = "cached_knowledge_base"  # Store locally
LANCEDB_TABLE = "documents"

# OpenAI API Key (securely stored)
os.environ["OPENAI_API_KEY"] = "sk-vDFgK4fxTRrJp1egIS***************gzL9RnT3BlbkFJ9f4caDIJRogirg-0bsJkhPZ4GfBFAWRHfrAYoucS0A"
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# AWS Credentials (Configure via AWS CLI)
# AWS_ACCESS_KEY_ID: AKIAWX2IFNT4***************
# AWS_SECRET_ACCESS_KEY: BshQLJRsPBKa***************zYmyBDxS
# AWS_DEFAULT_REGION: eu-north-1

# Global variable to store knowledge base
_loaded_knowledge_base = None

# ==========================
# CORE FUNCTIONS
# ==========================

def list_public_pdf_urls(bucket: str, folder: str) -> List[str]:
    """
    Fetch PDF URLs from S3 bucket.
    
    Args:
        bucket: S3 bucket name
        folder: Folder path in bucket
    
    Returns:
        List of public PDF URLs
    """
    s3_client = boto3.client("s3", region_name=AWS_REGION)
    response = s3_client.list_objects_v2(Bucket=bucket, Prefix=folder)

    pdf_urls = []
    if "Contents" in response:
        for obj in response["Contents"]:
            if obj["Key"].endswith(".pdf"):
                public_url = f"https://{bucket}.s3.{AWS_REGION}.amazonaws.com/{obj['Key']}"
                pdf_urls.append(public_url)

    return pdf_urls


def load_cached_metadata():
    """
    Load metadata of previously processed PDFs.
    
    Returns:
        Dictionary with cached metadata or None
    """
    if os.path.exists(CACHE_METADATA_FILE):
        with open(CACHE_METADATA_FILE, "r") as f:
            return json.load(f)
    return None


def save_cached_metadata(pdf_urls):
    """
    Save metadata of processed PDFs.
    
    Args:
        pdf_urls: List of processed PDF URLs
    """
    with open(CACHE_METADATA_FILE, "w") as f:
        json.dump({"pdf_urls": pdf_urls}, f)


def create_knowledge_base(pdf_urls: List[str]) -> PDFUrlKnowledgeBase:
    """
    Creates or loads the knowledge base with intelligent caching.
    
    Args:
        pdf_urls: List of PDF URLs to process
    
    Returns:
        PDFUrlKnowledgeBase object
    """
    global _loaded_knowledge_base

    # If cached, load from local database
    cached_metadata = load_cached_metadata()
    if cached_metadata and set(cached_metadata["pdf_urls"]) == set(pdf_urls):
        print("✅ Using cached knowledge base (Skipping PDF processing)...")
        if _loaded_knowledge_base is None:
            _loaded_knowledge_base = PDFUrlKnowledgeBase(
                urls=[],
                vector_db=LanceDb(
                    table_name=LANCEDB_TABLE,
                    uri=LANCEDB_URI,
                    search_type=SearchType.vector,
                    embedder=OpenAIEmbedder(model="text-embedding-3-small"),
                ),
            )
            _loaded_knowledge_base.load()
        return _loaded_knowledge_base

    print("⚡ Processing PDFs for the first time...")

    # Create new knowledge base and save metadata
    knowledge_base = PDFUrlKnowledgeBase(
        urls=pdf_urls,
        vector_db=LanceDb(
            table_name=LANCEDB_TABLE,
            uri=LANCEDB_URI,
            search_type=SearchType.vector,
            embedder=OpenAIEmbedder(model="text-embedding-3-small"),
        ),
    )

    knowledge_base.load()
    _loaded_knowledge_base = knowledge_base
    save_cached_metadata(pdf_urls)  # Save processed files

    return knowledge_base


def create_ai_agent(knowledge_base: PDFUrlKnowledgeBase) -> Agent:
    """
    Create an AI agent with the knowledge base.
    
    Args:
        knowledge_base: Loaded knowledge base
    
    Returns:
        Configured Phi Agent
    """
    agent = Agent(
        model=OpenAIChat(id="gpt-4o-mini"),
        knowledge=knowledge_base,
        markdown=True,
        show_tool_calls=True
    )
    return agent


def extract_and_save_function(response_text: str, filename: str = "generated_function.py"):
    """
    Extracts Python function from response and saves it.
    
    Args:
        response_text: Agent response containing code
        filename: Output filename
    
    Returns:
        Filename if successful, None otherwise
    """
    match = re.search(r"```python\n(.*?)\n```", response_text, re.DOTALL)

    if match:
        function_code = match.group(1).strip()

        with open(filename, "w", encoding="utf-8") as f:
            f.write(function_code)

        print(f"\n✅ Function saved as {filename}")
        return filename
    else:
        print("\n❌ No function detected in response.")
        return None


# ==========================
# MAIN EXECUTION LOGIC
# ==========================

def main():
    """Main execution function."""
    
    print("=" * 60)
    print("INTELLIGENT PHI-POWERED RAG AGENT")
    print("=" * 60)
    
    # Fetch PDFs from S3
    print("\nFetching PDFs from S3...")
    pdf_urls = list_public_pdf_urls(AWS_S3_BUCKET, PDF_FOLDER)

    if not pdf_urls:
        raise Exception("No PDF files found in S3 bucket!")

    print(f"Loaded {len(pdf_urls)} PDFs for processing.")

    # Initialize knowledge base with caching
    print("\nInitializing knowledge base...")
    start_time = time.time()
    knowledge_base = create_knowledge_base(pdf_urls)
    print(f"Knowledge base ready in {time.time() - start_time:.2f} seconds.")

    # Create AI Agent
    print("\nInitializing AI Agent...")
    agent = create_ai_agent(knowledge_base)

    if not agent.knowledge:
        raise Exception("Agent does not have a knowledge base assigned!")

    print("\n" + "=" * 60)
    print("AGENT READY - Enter 'exit' to quit")
    print("=" * 60)
    
    # Interactive prompt loop
    while True:
        print("\nEnter your prompt (or 'exit' to quit): ")
        user_input = input("▶ ")
        
        if user_input.lower() == "exit":
            break
        
        # Guided message for enhanced context
        guided_msg = ("While you gain massive insights as you scan through all the pdfs "
                     "(for Link Building Blog, LBB and DataForSEO) that was fed in the "
                     "knowledge base; ")
        
        # Process and stream response
        print("\n" + "-" * 60)
        agent.print_response(guided_msg + user_input, stream=True)
        print("-" * 60)
    
    print("\n✅ Session ended. Thank you for using the Intelligent Phi-Powered RAG Agent!")


# ==========================
# CACHE PERSISTENCE
# ==========================

def save_cache_to_drive():
    """Save the cache to Google Drive for persistence."""
    if os.path.exists(cache_path):
        print("Saving updated cached knowledge base to Google Drive...")
        os.system('rm -rf /content/drive/MyDrive/Colab\ Notebooks/cached_knowledge_base')
        os.system('rm -rf /content/drive/MyDrive/Colab\ Notebooks/cached_metadata.json')
        os.system('cp -r /content/cached_knowledge_base /content/drive/MyDrive/Colab\ Notebooks/')
        os.system('cp -r /content/cached_metadata.json /content/drive/MyDrive/Colab\ Notebooks/')
        print("✅ Cache saved to Google Drive")
    else:
        print("No cache found to save.")


# ==========================
# ENTRY POINT
# ==========================

if __name__ == "__main__":
    try:
        main()
    finally:
        # Always save cache when exiting
        save_cache_to_drive()

import sys
import os
import json
import boto3
import time
import re
from openai import OpenAI
from typing import List
from phi.agent import Agent
from phi.model.openai import OpenAIChat
from phi.vectordb.lancedb import LanceDb, SearchType
from phi.playground import serve_playground_app
from phi.embedder.openai import OpenAIEmbedder
from phi.storage.agent.sqlite import SqlAgentStorage
from phi.knowledge.pdf import PDFUrlKnowledgeBase

# AWS Configuration
AWS_S3_BUCKET = "lbb-and-dataforseo-pdfs"
AWS_REGION = "eu-north-1"
PDF_FOLDER = "folder-lbb-and-dataforseo-pdfs/"

# Local cache configuration
CACHE_METADATA_FILE = "cached_metadata.json"
LANCEDB_URI = "cached_knowledge_base"  # Store locally
LANCEDB_TABLE = "documents"

# OpenAI API Key (securely stored)
os.environ["OPENAI_API_KEY"] = "sk-vDFgK4fxTRrJp1egIS***************gzL9RnT3BlbkFJ9f4caDIJRogirg-0bsJkhPZ4GfBFAWRHfrAYoucS0A"
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# AWS Credentials (Redacted for security)
# AWS_ACCESS_KEY_ID: AKIAWX2IFNT4***************
# AWS_SECRET_ACCESS_KEY: BshQLJRsPBKa***************zYmyBDxS
# AWS_DEFAULT_REGION: eu-north-1

def list_public_pdf_urls(bucket: str, folder: str) -> List[str]:
    """Fetch PDF URLs from S3."""
    s3_client = boto3.client("s3", region_name=AWS_REGION)
    response = s3_client.list_objects_v2(Bucket=bucket, Prefix=folder)

    pdf_urls = []
    if "Contents" in response:
        for obj in response["Contents"]:
            if obj["Key"].endswith(".pdf"):
                public_url = f"https://{bucket}.s3.{AWS_REGION}.amazonaws.com/{obj['Key']}"
                pdf_urls.append(public_url)

    return pdf_urls


def load_cached_metadata():
    """Load metadata of previously processed PDFs."""
    if os.path.exists(CACHE_METADATA_FILE):
        with open(CACHE_METADATA_FILE, "r") as f:
            return json.load(f)
    return None


def save_cached_metadata(pdf_urls):
    """Save metadata of processed PDFs."""
    with open(CACHE_METADATA_FILE, "w") as f:
        json.dump({"pdf_urls": pdf_urls}, f)


def create_knowledge_base(pdf_urls: List[str]) -> PDFUrlKnowledgeBase:
    """Creates or loads the knowledge base."""
    global _loaded_knowledge_base

    # If cached, load from local database
    cached_metadata = load_cached_metadata()
    if cached_metadata and set(cached_metadata["pdf_urls"]) == set(pdf_urls):
        print("✅ Using cached knowledge base (Skipping PDF processing)...")
        if _loaded_knowledge_base is None:
            _loaded_knowledge_base = PDFUrlKnowledgeBase(
                urls=[],
                vector_db=LanceDb(
                    table_name=LANCEDB_TABLE,
                    uri=LANCEDB_URI,
                    search_type=SearchType.vector,
                    embedder=OpenAIEmbedder(model="text-embedding-3-small"),
                ),
            )
            _loaded_knowledge_base.load()
        return _loaded_knowledge_base

    print("⚡ Processing PDFs for the first time...")

    # Create new knowledge base and save metadata
    knowledge_base = PDFUrlKnowledgeBase(
        urls=pdf_urls,
        vector_db=LanceDb(
            table_name=LANCEDB_TABLE,
            uri=LANCEDB_URI,
            search_type=SearchType.vector,
            embedder=OpenAIEmbedder(model="text-embedding-3-small"),
        ),
    )

    knowledge_base.load()
    _loaded_knowledge_base = knowledge_base
    save_cached_metadata(pdf_urls)  # Save processed files

    return knowledge_base


def create_ai_agent(knowledge_base: PDFUrlKnowledgeBase) -> Agent:
    """Create an AI agent with the knowledge base."""
    agent = Agent(
        model=OpenAIChat(id="gpt-4o-mini"),
        knowledge=knowledge_base,
        markdown=True,
        show_tool_calls=True
    )
    return agent

# Mount Google Drive Automatically
from google.colab import drive
import os

drive.mount('/content/drive')

# Check if Cached Files Exist in Drive and Restore Them
cache_path = "/content/cached_knowledge_base"
drive_cache_path = f"/content/drive/MyDrive/Colab Notebooks/cached_knowledge_base"

# Check if cached files exist in Drive
if os.path.exists(drive_cache_path):
    print("Restoring cached knowledge base from Google Drive...")
    !cp -r /content/drive/MyDrive/Colab\ Notebooks/cached_knowledge_base /content/
    !cp -r /content/drive/MyDrive/Colab\ Notebooks/cached_metadata.json /content/
    print("Restored.")
else:
    print("No cached knowledge base found in Drive. It will be created from scratch or copied from colab.")
    # Ensure the cache is saved in Drive if not saved in Drive.
    if os.path.exists(cache_path):
        print("Saving updated cached knowledge base to Google Drive...")
        !rm -rf /content/drive/MyDrive/Colab\ Notebooks/cached_knowledge_base
        !rm -rf /content/drive/MyDrive/Colab\ Notebooks/cached_metadata.json
        !cp -r /content/cached_knowledge_base /content/drive/MyDrive/Colab\ Notebooks/
        !cp -r /content/cached_metadata.json /content/drive/MyDrive/Colab\ Notebooks/
    else:
        print("No cache found to save in colab. Please Run all the code from start to end to create cache.")

# Main Execution Logic
print("Fetching PDFs from S3...")
pdf_urls = list_public_pdf_urls(AWS_S3_BUCKET, PDF_FOLDER)

if not pdf_urls:
    raise Exception("No PDF files found in S3 bucket!")

print(f"Loaded {len(pdf_urls)} PDFs for processing.")

print("Initializing knowledge base...")
start_time = time.time()
knowledge_base = create_knowledge_base(pdf_urls)
print(f"Knowledge base ready in {time.time() - start_time:.2f} seconds.")

print("Initializing AI Agent...")
agent = create_ai_agent(knowledge_base)

if not agent.knowledge:
    raise Exception("Agent does not have a knowledge base assigned!")

# Start prompt loop (continuous interaction)
while True:
    print("Enter your prompt (or 'exit' to quit): ")
    user_input = input(" ")
    if user_input.lower() == "exit":
        break
    guided_msg = "While you gain massive insights as you scan through all the pdfs (for Link Building Blog, LBB and DataForSEO) that was fed in the knowledge base; "
    agent.print_response(guided_msg + user_input, stream=True)