main.py

# Import necessary libraries for the PolySensor application
# Core system libraries
import os  # For file system operations and environment variables
import requests  # For making HTTP requests to external APIs (RAG functionality)
import json  # For parsing JSON responses from LLM

# Environment and configuration management
from dotenv import load_dotenv, find_dotenv, set_key  # For managing environment variables securely

# AI and language model components
from langchain_google_genai import ChatGoogleGenerativeAI  # Google's AI model for text generation and analysis
from langchain.prompts import PromptTemplate  # For creating structured prompts for AI
from langchain.chains import LLMChain  # For chaining AI operations
from langchain_core.messages import HumanMessage, AIMessage  # For chat message formatting

# Data encoding and processing
import base64  # For encoding binary data (like images/audio) to text for AI processing
import getpass  # For securely getting user input (like API keys)
import uuid  # For generating unique identifiers for temporary files
from datetime import datetime  # For timestamp operations and metadata

# Custom application modules
from prompts import DOCUMENT_PROMPT, AUDIO_PROMPT, VIDEO_PROMPT, IMAGE_PROMPT  # Custom prompts for different media types
from data_handling import (
   unstructured_doc_extraction,  # Function to extract text from documents (PDF, Word, etc.)
   audio,  # Function to validate and process audio files (duration limits, format checks)
   image,  # Function to validate and process image files (existence checks)
   video   # Function to validate and process video files (duration limits, format checks)
   )

# PDF generation and document processing
from io import BytesIO  # For handling in-memory binary data streams
from reportlab.pdfgen import canvas  # For PDF generation (basic canvas operations)
from reportlab.lib.pagesizes import letter  # Standard letter page size for PDFs
from reportlab.platypus import Table, TableStyle, Paragraph, SimpleDocTemplate, Spacer  # PDF layout components
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle  # PDF text styling
from reportlab.lib import colors  # Color definitions for PDFs
from reportlab.lib.units import inch  # Unit conversions for PDF layout

# Web framework and API components
from flask import Flask, request, jsonify, session, send_from_directory, send_file  # Flask web framework components
from flask_cors import CORS  # For handling Cross-Origin Resource Sharing (frontend-backend communication)

# Text processing and HTML manipulation
import markdown  # For converting markdown text to HTML (PDF generation)
from bs4 import BeautifulSoup  # For parsing and manipulating HTML (PDF generation)

# Vector database for RAG (Retrieval-Augmented Generation)
from vector_db import get_vector_db  # Function to get vector database instance for storing/retrieving analyses
from api_endpoints import vector_api  # Blueprint for vector database API endpoints

# Load environment variables from .env file for secure configuration
load_dotenv(find_dotenv())

# Get Google API key from environment variables or prompt user to enter it
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
if not GOOGLE_API_KEY:
    GOOGLE_API_KEY = getpass.getpass("Enter API Key:")
    if GOOGLE_API_KEY:
        dotenv_path = find_dotenv()
        if not dotenv_path:
            dotenv_path = '.env'
        set_key(dotenv_path, 'GOOGLE_API_KEY', GOOGLE_API_KEY)
        os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY

# Initialize the Google Generative AI model for text generation and analysis
llm = ChatGoogleGenerativeAI(model='gemini-2.5-pro', google_api_key=GOOGLE_API_KEY)

# Create prompt templates for different types of content analysis
document_prompt = PromptTemplate(input_variables=['doc_json_data'], template=DOCUMENT_PROMPT)
# audio_prompt = PromptTemplate(input_variables=['audio_data'], template=AUDIO_PROMPT)  # Unused
# image_prompt = PromptTemplate(input_variables=['image_data'], template=IMAGE_PROMPT)  # Unused
# video_prompt = PromptTemplate(input_variables=['video_data'], template=VIDEO_PROMPT)  # Unused

# Create AI processing chains for different content types
document_chain = LLMChain(llm=llm, prompt=document_prompt)
# image_chain = LLMChain(llm=llm, prompt = image_prompt)  # Unused
# audio_chain = LLMChain(llm=llm, prompt = audio_prompt)  # Unused
# video_chain = LLMChain(llm=llm, prompt = video_prompt)  # Unused

# Define supported file extensions for different media types
# These tuples contain file extensions that the application can process
doc_extensions = (
    '.cwk', '.mcw', '.csv', '.dif', '.dbf', '.eml', '.msg', '.p7s', '.epub',
    '.htm', '.html', '.bmp', '.heic', '.prn', '.tiff',
    '.md', '.odt', '.org', '.eth', '.pbd', '.sdp', '.pdf', '.txt', '.pot',
    '.ppt', '.pptm', '.pptx', '.rst', '.rtf', '.et', '.fods', '.mw', '.xls',
    '.xlsx', '.sxg', '.tsv', '.abw', '.doc', '.docm', '.docx', '.dot', '.dotm',
    '.hwp', '.zabw', '.xml'
)

image_extensions = ('.jpg', '.png', '.jpeg', '.webp', '.heif')
audio_extensions = ('.mp3', '.wav', '.aiff', '.aac', '.ogg', '.flac')

# MIME types mapping for audio files - used when sending audio data to AI models
audio_mime_types = {
    '.mp3': 'audio/mpeg',
    '.wav': 'audio/wav',
    '.flac': 'audio/flac',
    '.aac': 'audio/aac',
    '.ogg': 'audio/ogg',
    '.aiff': 'audio/aiff'
}
video_extensions = ('.mp4', '.mpeg', '.mov', '.avi', '.x-flav', '.mpg', '.webm', '.wmv', '.3gpp')

# Create Flask application instance
polysensor = Flask(__name__)
# Set secret key for session management (used for security in web applications)
polysensor.secret_key = os.getenv("FLASK_SECRET_KEY", "supersecretkey")
# Enable Cross-Origin Resource Sharing (CORS) to allow frontend requests from specific origins
CORS(polysensor, resources={r"/*": {"origins": ["http://localhost:5173", "http://localhost:5174"], "supports_credentials": True}})

# Initialize the vector database for storing and retrieving analysis results
vector_db = get_vector_db()

# Register the vector API blueprint to add vector database endpoints to the Flask app
polysensor.register_blueprint(vector_api)


# Health check endpoint for Docker healthcheck
@polysensor.route('/health', methods=['GET'])
def health():
    return jsonify({'status': 'healthy'}), 200


# Chat endpoint: Handles conversational AI interactions with Retrieval-Augmented Generation (RAG)
# RAG combines chat history with relevant past analysis results from the vector database for better responses
@polysensor.route('/chat', methods=['POST'])
def chat_llm():
    try:
        # Parse the incoming JSON request containing the conversation history
        data = request.get_json()
        history = data.get('history', [])
        
        # Convert the chat history into LangChain-compatible message objects
        # HumanMessage for user inputs, AIMessage for AI responses
        messages = []
        for msg in history:
            if msg['role'] == 'user':
                messages.append(HumanMessage(content=msg['content']))
            elif msg['role'] == 'assistant':
                messages.append(AIMessage(content=msg['content']))

        # Implement RAG: Retrieve relevant context from vector DB to augment the latest user query
        if messages and isinstance(messages[-1], HumanMessage):
            query = messages[-1].content
            # Search the vector database for similar past analysis content
            try:
                search_payload = {
                    "query": query,
                    "n_results": 3,  # Limit to top 3 most relevant results
                    "content_type": "analysis"  # Filter to only analysis results
                }
                # Send POST request to the internal vector search API
                search_response = requests.post("http://localhost:5000/search", json=search_payload)
                if search_response.status_code == 200:
                    search_results = search_response.json().get('results', [])
                    if search_results:
                        # Combine retrieved contexts into a single string and append to the query
                        context = "\n\nRelevant analysis from previous files:\n" + "\n".join([r['content'] for r in search_results])
                        messages[-1] = HumanMessage(content=f"{query}\n{context}")
            except Exception as rag_error:
                # Log RAG failure but continue without context
                print(f"RAG search failed: {rag_error}")

        # Invoke the LLM with the full conversation history (and optional RAG context) to generate response
        llm_reply = llm.invoke(messages)
        return jsonify({'answer': llm_reply.content})
    except Exception as e:
        # Handle any errors during chat processing
        return jsonify({'error': str(e)}), 500



# File analysis endpoint: Processes uploaded files of various types (documents, images, audio, video)
# Uses AI models to analyze content and stores results in vector database for future retrieval
@polysensor.route('/analyze0', methods=['POST'])
def analyze_file():
    file_path = None
    try:
        # Validate file upload
        if 'file' not in request.files:
            return jsonify({'error': 'No file provided'}), 400

        file = request.files['file']
        if file.filename == '':
            return jsonify({'error': 'No file selected'}), 400

        # Create temporary directory and save uploaded file with unique name
        temp_dir = os.path.join(os.getcwd(), 'temp')
        if not os.path.exists(temp_dir):
            os.makedirs(temp_dir, exist_ok=True)
        unique_filename = f"{uuid.uuid4()}_{file.filename}"
        file_path = os.path.join(temp_dir, unique_filename)
        file.save(file_path)

        result = None

        # Process document files (PDF, Word, etc.) using unstructured data extraction
        if file_path.lower().endswith(doc_extensions):
            doc_json_data = unstructured_doc_extraction(file_path)
            if isinstance(doc_json_data, str):
                os.remove(file_path)
                return jsonify({'error': doc_json_data}), 400
            if doc_json_data:
                # Use AI to analyze the extracted document content
                llm_doc_output = document_chain.run(doc_json_data=doc_json_data)
                # Try to parse as JSON for structured data, fallback to text
                try:
                    result = json.loads(llm_doc_output)
                    result['data_type'] = 'structured'
                except json.JSONDecodeError:
                    result = llm_doc_output
                    result = {'data_type': 'text', 'content': result}

        # Process image files by sending them to AI vision model
        elif file_path.lower().endswith(image_extensions):
            image_result = image(file_path)
            if not os.path.exists(image_result):
                os.remove(file_path)
                return jsonify({'error': 'Image file not found after processing'}), 400
            image_path = image_result
            if image_path:
                # Read image as binary data and encode as base64 for AI processing
                with open(image_path, "rb") as img:
                    image_bytes = img.read()
                    prompt_with_image = HumanMessage(
                        content=[
                            {
                                'type': 'text',
                                'text': IMAGE_PROMPT
                            },
                            {
                                'type': 'image_url',
                                'image_url': f"data:image/png;base64,{base64.b64encode(image_bytes).decode('utf-8')}"
                            }
                        ]
                    )
                    # Invoke AI model with image and prompt
                    final_image_prompt = llm.invoke([prompt_with_image])
                    result = final_image_prompt.content

        # Process audio files by sending them to AI audio analysis model
        elif file_path.lower().endswith(audio_extensions):
            audio_result = audio(file_path)
            if isinstance(audio_result, str):
                os.remove(file_path)
                return jsonify({'error': audio_result}), 400
            audio_path = audio_result
            if audio_path:
                # Read audio as binary data and encode as base64
                with open(audio_path, "rb") as f:
                    audio_bytes = f.read()
                # Determine correct MIME type for the audio file
                ext = file_path.lower().rsplit('.', 1)[-1]
                mime_type = audio_mime_types.get('.' + ext, 'audio/mpeg')

                prompt_with_audio = HumanMessage([
                    {
                        'type': 'text',
                        'text': AUDIO_PROMPT
                    },
                    {
                        'type': 'media',
                        'mime_type': mime_type,
                        'data': base64.b64encode(audio_bytes).decode('utf-8')
                    }
                ])
                # Invoke AI model with audio and prompt
                final_audio_prompt = llm.invoke([prompt_with_audio])
                result = final_audio_prompt.content

        # Process video files by sending them to AI video analysis model
        elif file_path.lower().endswith(video_extensions):
            video_result = video(file_path)
            if not os.path.exists(video_result):
                os.remove(file_path)

            if isinstance(video_result, str):
                return jsonify({'error': video_result}), 400
            video_path = video_result
            if video_path:
                # Read video as binary data and encode as base64
                with open(video_path, "rb") as vid:
                    video_bytes = vid.read()
                    mime_type = 'video/mp4' if file_path.lower().endswith('.mp4') else 'video/mpeg'
                    prompt_with_video = HumanMessage([
                        {
                            'type': 'text',
                            'text': VIDEO_PROMPT
                        },
                        {
                            'type': 'media',
                            'mime_type': mime_type,
                            'data': base64.b64encode(video_bytes).decode('utf-8')
                        }
                    ])
                    # Invoke AI model with video and prompt
                    final_video_prompt = llm.invoke([prompt_with_video])
                    result = final_video_prompt.content

        # If analysis was successful, store result in vector database
        if result:
            # Ensure result is a string for frontend rendering and storage
            if isinstance(result, dict):
                result_str = result.get('content', json.dumps(result))
            else:
                result_str = str(result)
            
            try:
                # Create metadata for the analysis result
                metadata = {
                    'filename': file.filename,
                    'file_type': file_path.lower().rsplit('.', 1)[-1] if '.' in file_path else 'unknown',
                    'analysis_timestamp': datetime.now().isoformat(),
                    'file_size': os.path.getsize(file_path) if os.path.exists(file_path) else 0
                }
                # Store analysis as string in vector database for future RAG retrieval
                doc_id = vector_db.add_content(result_str, metadata=metadata, content_type='analysis')
                print(f"Analysis stored in vector database with ID: {doc_id}")
            except Exception as db_error:
                print(f"Failed to store analysis in vector database: {db_error}")

            # Clean up temporary file and return analysis result as string
            os.remove(file_path)
            return jsonify({'result': result_str})
        else:
            # Clean up and return error for unsupported file types
            os.remove(file_path)
            return jsonify({'error': 'Unsupported file type or processing failed'}), 400

    except Exception as e:
        # Handle any unexpected errors during file processing
        return jsonify({'error': str(e)}), 500
    finally:
        # Ensure temporary file is always cleaned up
        if file_path and os.path.exists(file_path):
            os.remove(file_path)

# PDF export endpoint: Converts markdown-formatted analysis content to professional PDF reports
# Features custom styling, branding, and support for tables, code blocks, and various markdown elements
@polysensor.route('/export-pdf', methods=['POST'])
def export_pdf():
    try:
        # Extract request parameters for PDF generation
        data = request.get_json()
        content = data.get('content', '')  # Markdown content to convert
        filename = data.get('filename', 'analysis')  # Base filename for the PDF
        count = data.get('count', 1)  # Report number for unique naming

        # Import ReportLab utilities for PDF styling and layout
        from reportlab.lib.colors import HexColor
        from reportlab.pdfbase.pdfmetrics import stringWidth

        def clean_content(content):
            """Remove boilerplate phrases from analysis content to clean up the PDF output."""
            unwanted_phrases = [
                "Analysis on textual information has finished and here are the results:",
                "Analysis on audio information has finished and here are the results:",
                "Analysis on image information has finished and here are the results:",
                "Analysis on video information has finished and here are the results:"
            ]
            # Filter out lines containing unwanted boilerplate text
            lines = content.split('\n')
            cleaned_lines = []
            for line in lines:
                if not any(phrase.lower() in line.strip().lower() for phrase in unwanted_phrases):
                    cleaned_lines.append(line)
            return '\n'.join(cleaned_lines)

        # Define custom PDF styling with modern, professional appearance
        def create_custom_styles():
            styles = getSampleStyleSheet()
            # Title style with modern blue color and center alignment
            title_style = ParagraphStyle(
                'ModernTitle',
                parent=styles['Heading1'],
                fontSize=26,
                spaceAfter=40,
                textColor=HexColor('#3498db'),  # Professional blue
                alignment=1,  # Center alignment
                fontName='Helvetica-Bold'
            )
            # Heading style for section headers
            heading_style = ParagraphStyle(
                'ModernHeading',
                parent=styles['Heading2'],
                fontSize=18,
                spaceAfter=15,
                spaceBefore=15,
                textColor=HexColor('#2c3e50'),  # Dark grey-blue
                fontName='Helvetica-Bold'
            )
            # Standard paragraph style for body text
            normal_style = ParagraphStyle(
                'ModernNormal',
                parent=styles['Normal'],
                fontSize=11,
                spaceAfter=8,
                leading=14,  # Line height
                textColor=HexColor('#34495e'),  # Medium grey
                fontName='Helvetica'
            )
            # Style for table data cells
            table_cell_style = ParagraphStyle(
                'TableCell',
                parent=styles['Normal'],
                fontSize=10,
                spaceAfter=4,
                leading=12,
                textColor=HexColor('#2c3e50'),
                fontName='Helvetica',
                alignment=0,  # Left alignment
                leftIndent=0,
                rightIndent=0
            )
            # Style for table header cells with white text
            table_header_style = ParagraphStyle(
                'TableHeader',
                parent=styles['Normal'],
                fontSize=11,
                textColor=colors.white,  # White text for blue background
                fontName='Helvetica-Bold',
                alignment=0,  # Left alignment
                leftIndent=0,
                rightIndent=0
            )
            # Style for code blocks with monospace font
            code_style = ParagraphStyle(
                'ModernCode',
                parent=styles['Normal'],
                fontName='Courier',  # Monospace font for code
                fontSize=9,
                leftIndent=20,  # Indent code blocks
                spaceAfter=12
            )
            # Add all custom styles to the stylesheet
            styles.add(title_style)
            styles.add(heading_style)
            styles.add(normal_style)
            styles.add(table_cell_style)
            styles.add(table_header_style)
            styles.add(code_style)
            return styles

        # Function to add branded footer to each page
        def add_footer(canvas, doc):
            canvas.saveState()
            canvas.setFont('Helvetica-Oblique', 8)
            canvas.setFillColor(HexColor('#bdc3c7'))  # Light grey color
            footer_text = "Generated by PolySensor - Advanced AI Analysis Platform"
            text_width = stringWidth(footer_text, 'Helvetica-Oblique', 8)
            # Center the footer text on the page
            x = (612 - text_width) / 2  # Letter width is 612 points
            canvas.drawString(x, 30, footer_text)
            canvas.restoreState()

        # Convert cleaned markdown content to HTML for parsing
        clean_content_str = clean_content(content)
        html = markdown.markdown(clean_content_str, extensions=['tables', 'fenced_code', 'codehilite', 'nl2br'])

        # Parse HTML structure using BeautifulSoup for element-by-element processing
        soup = BeautifulSoup(html, 'html.parser')

        # Set up PDF document with letter size and margins
        buffer = BytesIO()
        doc = SimpleDocTemplate(buffer, pagesize=letter, rightMargin=72, leftMargin=72, topMargin=72, bottomMargin=72)
        styles = create_custom_styles()
        story = []  # List to hold PDF elements

        # Add branded title section to the PDF
        if filename:
            title_text = f"<b>Analysis Report: {filename}</b><br/><font size=10 color='#7f8c8d'>Generated by PolySensor</font>"
            title = Paragraph(title_text, styles['ModernTitle'])
            story.append(title)
            story.append(Spacer(1, 30))  # Add space after title

        # Process each top-level HTML element and convert to PDF elements
        for element in soup.find_all(recursive=False):
            if element.name == 'p':
                # Handle paragraph elements
                p_text = element.get_text().strip()
                if p_text:
                    para = Paragraph(p_text, styles['ModernNormal'])
                    story.append(para)
                    story.append(Spacer(1, 6))
            elif element.name in ['h1', 'h2', 'h3', 'h4', 'h5', 'h6']:
                # Handle heading elements
                heading_text = element.get_text().strip()
                if heading_text:
                    para = Paragraph(heading_text, styles['ModernHeading'])
                    story.append(para)
                    story.append(Spacer(1, 12))
            elif element.name == 'ul':
                # Handle unordered lists
                for li in element.find_all('li'):
                    li_text = li.get_text().strip()
                    if li_text:
                        para = Paragraph('• ' + li_text, styles['ModernNormal'])
                        story.append(para)
                        story.append(Spacer(1, 6))
                story.append(Spacer(1, 12))
            elif element.name == 'ol':
                # Handle ordered lists
                lis = element.find_all('li')
                for idx, li in enumerate(lis, 1):
                    li_text = li.get_text().strip()
                    if li_text:
                        para = Paragraph(f'{idx}. ' + li_text, styles['ModernNormal'])
                        story.append(para)
                        story.append(Spacer(1, 6))
                story.append(Spacer(1, 12))
            elif element.name == 'pre':
                # Handle code blocks
                code_text = element.get_text().strip()
                if code_text:
                    code_para = Paragraph(code_text, styles['ModernCode'])
                    story.append(code_para)
                    story.append(Spacer(1, 12))
            elif element.name == 'table':
                # Handle table elements with custom styling
                rows = []
                for row in element.find_all('tr'):
                    cells = []
                    for cell in row.find_all(['th', 'td']):
                        cell_text = cell.get_text().strip()
                        if cell.name == 'th':
                            # Header cells with special styling
                            para = Paragraph(cell_text, styles['TableHeader'])
                        else:
                            # Data cells with standard styling
                            para = Paragraph(cell_text, styles['TableCell'])
                        cells.append(para)
                    if cells:
                        rows.append(cells)
                if rows:
                    # Calculate column widths dynamically
                    num_cols = len(rows[0])
                    content_width = 468  # Available width within margins
                    col_widths = [content_width / num_cols] * num_cols
                    table = Table(rows, colWidths=col_widths)
                    # Apply professional table styling
                    table.setStyle(TableStyle([
                        ('BACKGROUND', (0, 0), (-1, 0), HexColor('#3498db')),  # Blue header background
                        ('TEXTCOLOR', (0, 0), (-1, 0), colors.white),  # White header text
                        ('ALIGN', (0, 0), (-1, -1), 'LEFT'),  # Left align all content
                        ('VALIGN', (0, 0), (-1, -1), 'TOP'),  # Top align content
                        ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),  # Bold headers
                        ('FONTNAME', (0, 1), (-1, -1), 'Helvetica'),  # Regular font for data
                        ('FONTSIZE', (0, 0), (-1, 0), 11),  # Header font size
                        ('FONTSIZE', (0, 1), (-1, -1), 10),  # Data font size
                        ('BOTTOMPADDING', (0, 0), (-1, 0), 10),  # Header padding
                        ('TOPPADDING', (0, 0), (-1, 0), 10),
                        ('BACKGROUND', (0, 1), (-1, -1), colors.white),  # White background for data
                        ('GRID', (0, 0), (-1, -1), 0.5, HexColor('#bdc3c7')),  # Light grey grid lines
                        ('LEFTPADDING', (0, 0), (-1, -1), 8),  # Cell padding
                        ('RIGHTPADDING', (0, 0), (-1, -1), 8),
                        ('TOPPADDING', (0, 1), (-1, -1), 6),
                        ('BOTTOMPADDING', (0, 1), (-1, -1), 6),
                    ]))
                    story.append(table)
                    story.append(Spacer(1, 15))  # Space after table

        # Build the PDF document with footer on all pages
        doc.build(story, onFirstPage=add_footer, onLaterPages=add_footer)
        buffer.seek(0)  # Reset buffer position for reading

        # Return the generated PDF as a downloadable file
        return send_file(buffer, as_attachment=True, download_name=f'{filename}_analysis_report_{count}.pdf', mimetype='application/pdf')
    except Exception as e:
        # Log and return any errors during PDF generation
        print(f"PDF generation error: {str(e)}")
        return jsonify({'error': str(e)}), 500

if __name__ == "__main__":
   port = int(os.environ.get('PORT', 5000))
   polysensor.run(host='0.0.0.0', port=port, debug=True)

   # if file_path.lower().endswith(('.mp4', '.mkv')):
   # video_text_data, video_audio_data = video_audio_text(file_path, interval_sec = 3)


   # if video_text_data is not None and video_audio_data is not None:
   #    video_data = f"Frame OCR Text: {video_text_data}\n Audio Transcribed: {video_audio_data}"
   #    final_video_prompt = VIDEO_TEXT.format(video_audio_data = video_data)
   # elif video_text_data is not None:
   #    final_video_prompt = VIDEO_TEXT.format(video_audio_data = video_text_data)
   # elif video_audio_data is not None:
   #    final_video_prompt = VIDEO_TEXT.format(video_audio_data=video_audio_data)