#!/usr/bin/env python3
"""
Chat API Server for Ollama Integration

This Flask server provides a secure API endpoint for the chat interface.
It handles rate limiting, input validation, and communicates with Ollama.

Usage:
    python3 chat_api.py

Configuration:
    - Ollama endpoint: 10.30.20.110:11434
    - Model: drjones-posts-to-much
    - Port: 5000 (configurable via PORT env var)
"""

import os
import sys
import time
import re
import yaml
from pathlib import Path
from datetime import datetime, timedelta
from flask import Flask, request, jsonify
from flask_cors import CORS
from functools import wraps
import hashlib

try:
    from ollama import Client
except ImportError:
    print("ERROR: ollama Python library not installed. Run: pip3 install ollama", file=sys.stderr)
    sys.exit(1)

try:
    import pdfplumber
    PDF_SUPPORT = True
except ImportError:
    PDF_SUPPORT = False
    print("WARNING: pdfplumber not installed. PDF text extraction disabled.", file=sys.stderr)

# Configuration
OLLAMA_HOST = os.getenv("OLLAMA_HOST", "http://10.30.20.110:11434")
MODEL = os.getenv("OLLAMA_MODEL", "drjones-posts-to-much")
PORT = int(os.getenv("PORT", 5000))
API_KEY = os.getenv("CHAT_API_KEY", "")  # Optional API key for additional security
LIBRARY_DIR = Path("/root/hydro-sterile/static/library")
LIBRARY_YAML = Path("/root/hydro-sterile/data/library.yaml")
MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB max file size
MAX_TEXT_LENGTH = 5000  # Max characters to extract from PDF

# Rate limiting storage (in-memory, simple implementation)
rate_limits = {}

# System prompt for the chat
SYSTEM_PROMPT = """You are Dr. Jones, a sterile hydroponics expert with 20 years of experience growing cannabis in a legal state.
You provide practical, direct, and knowledgeable advice about sterile hydroponics, cannabis cultivation, and growing techniques.
Keep responses concise (2-4 paragraphs), practical, and actionable. Write in first person (I, me, my) and be conversational."""

app = Flask(__name__)
# Allow CORS from the website domain and local development
CORS(app, origins=[
    "https://hydrolord.thetempleofdoom.com",
    "http://hydrolord.thetempleofdoom.com",
    "http://localhost:1313",
    "http://127.0.0.1:1313",
    "http://10.30.20.243",
    "http://10.30.20.243:1313"
])

def get_client_ip():
    """Get client IP address for rate limiting."""
    if request.headers.get('X-Forwarded-For'):
        return request.headers.get('X-Forwarded-For').split(',')[0].strip()
    return request.remote_addr

def rate_limit(max_per_minute=10, window_minutes=1):
    """Simple rate limiting decorator."""
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            client_ip = get_client_ip()
            now = time.time()
            window_seconds = window_minutes * 60
            
            # Clean old entries
            rate_limits[client_ip] = [
                timestamp for timestamp in rate_limits.get(client_ip, [])
                if now - timestamp < window_seconds
            ]
            
            # Check rate limit
            if len(rate_limits.get(client_ip, [])) >= max_per_minute:
                return jsonify({
                    "error": "Rate limit exceeded. Please wait a moment before sending another message."
                }), 429
            
            # Add current request
            if client_ip not in rate_limits:
                rate_limits[client_ip] = []
            rate_limits[client_ip].append(now)
            
            return f(*args, **kwargs)
        return decorated_function
    return decorator

def sanitize_input(text):
    """Sanitize user input to prevent injection attacks."""
    if not text or not isinstance(text, str):
        return ""
    
    # Remove potentially dangerous characters
    text = re.sub(r'[<>]', '', text)
    
    # Limit length
    text = text[:1000].strip()
    
    return text

def load_library_files():
    """Load library file list from YAML."""
    try:
        with open(LIBRARY_YAML, 'r', encoding='utf-8') as f:
            data = yaml.safe_load(f)
        
        files = []
        for entry in data or []:
            if isinstance(entry, dict) and 'title' in entry and 'file' in entry:
                file_path = entry.get('file', '')
                # Convert /library/ path to actual file path
                if file_path.startswith('/library/'):
                    filename = file_path.replace('/library/', '')
                    full_path = LIBRARY_DIR / filename
                    if full_path.exists():
                        files.append({
                            'title': entry.get('title', ''),
                            'filename': entry.get('filename', ''),
                            'file': file_path,
                            'category': entry.get('category', ''),
                            'size': full_path.stat().st_size if full_path.exists() else 0
                        })
        return files
    except Exception as e:
        print(f"Error loading library: {e}", file=sys.stderr)
        return []

def extract_pdf_text(file_path, max_length=MAX_TEXT_LENGTH):
    """Extract text from PDF file."""
    if not PDF_SUPPORT:
        return None
    
    try:
        full_path = LIBRARY_DIR / file_path.replace('/library/', '')
        
        if not full_path.exists():
            return None
        
        if full_path.stat().st_size > MAX_FILE_SIZE:
            return f"[File too large: {full_path.stat().st_size / 1024 / 1024:.1f}MB. Max size: {MAX_FILE_SIZE / 1024 / 1024}MB]"
        
        text_parts = []
        with pdfplumber.open(str(full_path)) as pdf:
            for page in pdf.pages[:10]:  # Limit to first 10 pages
                page_text = page.extract_text()
                if page_text:
                    text_parts.append(page_text)
                    if len(' '.join(text_parts)) > max_length:
                        break
        
        full_text = ' '.join(text_parts)
        if len(full_text) > max_length:
            full_text = full_text[:max_length] + "... [truncated]"
        
        return full_text.strip()
    except Exception as e:
        print(f"Error extracting PDF text: {e}", file=sys.stderr)
        return None

@app.route('/api/chat', methods=['POST'])
@rate_limit(max_per_minute=10, window_minutes=1)
def chat():
    """Handle chat requests."""
    try:
        data = request.get_json()
        
        if not data or 'message' not in data:
            return jsonify({"error": "Message is required"}), 400
        
        message = sanitize_input(data['message'])
        
        if not message:
            return jsonify({"error": "Message cannot be empty"}), 400
        
        # Optional API key check
        if API_KEY and data.get('api_key') != API_KEY:
            return jsonify({"error": "Invalid API key"}), 401
        
        # Get selected library files
        selected_files = data.get('files', [])
        file_context = ""
        
        if selected_files:
            file_context_parts = []
            for file_ref in selected_files[:3]:  # Limit to 3 files max
                if isinstance(file_ref, str):
                    file_path = file_ref
                elif isinstance(file_ref, dict) and 'file' in file_ref:
                    file_path = file_ref['file']
                else:
                    continue
                
                # Extract text from PDF
                pdf_text = extract_pdf_text(file_path)
                if pdf_text:
                    # Get file title
                    file_title = file_ref.get('title', '') if isinstance(file_ref, dict) else file_path
                    file_context_parts.append(f"\n\n--- Content from: {file_title} ---\n{pdf_text}\n--- End of {file_title} ---")
            
            if file_context_parts:
                file_context = "\n\n[The user has referenced the following library documents. Use this information to provide accurate, specific answers based on these sources:]\n" + "\n".join(file_context_parts)
        
        # Get conversation history (last 5 messages for context)
        history = data.get('history', [])
        history_messages = []
        
        # Build conversation history
        for msg in history[-5:]:  # Last 5 messages
            if isinstance(msg, dict) and 'role' in msg and 'content' in msg:
                history_messages.append({
                    "role": msg['role'],
                    "content": sanitize_input(msg['content'])
                })
        
        # Add current user message with file context
        user_message = message
        if file_context:
            user_message = message + file_context
        
        history_messages.append({
            "role": "user",
            "content": user_message
        })
        
        # Call Ollama
        try:
            client = Client(host=OLLAMA_HOST)
            
            # Build messages for Ollama
            ollama_messages = [
                {"role": "system", "content": SYSTEM_PROMPT}
            ]
            
            # Add conversation history
            for msg in history_messages:
                ollama_messages.append({
                    "role": msg['role'],
                    "content": msg['content']
                })
            
            # Generate response
            response = client.chat(
                model=MODEL,
                messages=ollama_messages,
                options={
                    "temperature": 0.7,
                    "top_p": 0.9,
                }
            )
            
            assistant_response = response.message.content.strip()
            
            return jsonify({
                "response": assistant_response,
                "model": MODEL
            })
            
        except Exception as e:
            print(f"Ollama error: {e}", file=sys.stderr)
            return jsonify({
                "error": "Unable to connect to AI service. Please try again later."
            }), 503
            
    except Exception as e:
        print(f"API error: {e}", file=sys.stderr)
        return jsonify({
            "error": "An error occurred processing your request."
        }), 500

@app.route('/api/library/files', methods=['GET'])
def list_library_files():
    """List available library files."""
    try:
        files = load_library_files()
        # Filter to only PDF files for now
        pdf_files = [f for f in files if f['filename'].lower().endswith('.pdf')]
        return jsonify({
            "files": pdf_files,
            "count": len(pdf_files)
        })
    except Exception as e:
        print(f"Error listing files: {e}", file=sys.stderr)
        return jsonify({
            "error": "Failed to load library files"
        }), 500

@app.route('/api/library/preview', methods=['POST'])
def preview_file():
    """Preview/extract text from a library file."""
    try:
        data = request.get_json()
        file_path = data.get('file')
        
        if not file_path:
            return jsonify({"error": "File path required"}), 400
        
        text = extract_pdf_text(file_path)
        
        if text is None:
            return jsonify({
                "error": "Could not extract text from file. File may not be a PDF or may be corrupted."
            }), 400
        
        return jsonify({
            "text": text,
            "length": len(text)
        })
    except Exception as e:
        print(f"Error previewing file: {e}", file=sys.stderr)
        return jsonify({
            "error": "Failed to preview file"
        }), 500

@app.route('/api/health', methods=['GET'])
def health():
    """Health check endpoint."""
    return jsonify({
        "status": "healthy",
        "model": MODEL,
        "ollama_host": OLLAMA_HOST,
        "pdf_support": PDF_SUPPORT
    })

if __name__ == '__main__':
    print(f"Starting Chat API server on port {PORT}")
    print(f"Ollama endpoint: {OLLAMA_HOST}")
    print(f"Model: {MODEL}")
    app.run(host='0.0.0.0', port=PORT, debug=False)