import base64
from flask import Flask, render_template, request, jsonify, Response, stream_with_context
from huggingface_hub import login
from vismem.core import SectorMemory
from ai_engine import LocalModelHandler, get_embedding

import json
import os
from werkzeug.utils import secure_filename

app = Flask(__name__)

if os.getenv("HF_TOKEN"): login(token=os.getenv("HF_TOKEN"))
repo_id = os.getenv("LOCAL_MODEL_PATH", "google/gemma-3-1b-it") 
local_bot = LocalModelHandler(repo_id=repo_id)

mem_history = SectorMemory("history")
# Future: mem_docs = SectorMemory("documents") 

@app.route('/')
def index():
    return render_template('index.html')

# --- NEW: RAG INGESTION ENDPOINTS ---
@app.route('/ingest/upload', methods=['POST'])
def ingest_upload():
    if 'files[]' not in request.files:
        return jsonify({"error": "No file part"}), 400
    
    files = request.files.getlist('files[]')
    processed_files = []
    
    for file in files:
        if file.filename == '': continue
        filename = secure_filename(file.filename)
        
        # TODO: Implement your novel document parsing logic here
        # content = parse_pdf_or_txt(file)
        # vec = get_embedding(content)
        # mem_docs.write_entry(content, vec)
        
        processed_files.append(filename)
        print(f"[RAG] Received file: {filename}")

    return jsonify({
        "success": True, 
        "message": f"Queued {len(processed_files)} documents for vectorization.",
        "files": processed_files
    })

@app.route('/ingest/url', methods=['POST'])
def ingest_url():
    data = request.json
    target_url = data.get('url')
    
    if not target_url: return jsonify({"error": "No URL provided"}), 400
    
    # TODO: Implement URL scraping and PDF extraction logic here
    print(f"[RAG] Processing URL: {target_url}")
    
    return jsonify({
        "success": True,
        "message": f"URL {target_url} added to processing queue."
    })
# ------------------------------------

@app.route('/visualize/<mem_type>')
def visualize(mem_type):
    target = mem_history
    
    # 1. Generate Image
    img_bytes = target.to_image_bytes()
    b64 = base64.b64encode(img_bytes).decode('utf-8')
    
    # 2. Get Stats
    active_count = target.count
    text_usage = target.cursor_green
    
    return jsonify({
        "image": f"data:image/png;base64,{b64}",
        "stats": f"{active_count} Active Entries",
        "usage": f"{text_usage} Bytes Text"
    })

@app.route('/inspect/<mem_type>')
def inspect(mem_type):
    content = mem_history.dump_heap_content()
    return jsonify({"content": content})

@app.route('/chat', methods=['POST'])
def chat():
    user_msg = request.json.get('message')
    if not user_msg: return jsonify({"error": "Empty"}), 400

    def generate():
        # 1. VECTORIZE USER INPUT
        q_vec = get_embedding(user_msg)
        
        # B. Relevant Past History
        hist_hits = mem_history.search(q_vec, top_k=2)
        long_term_txt = "\n".join([f"[Memory]: {h['text']}" for h in hist_hits if h['score'] > 0.4])
        
        # C. Recent Conversation
        recent_msgs = mem_history.get_recent_entries(n=4)
        recent_txt = "\n".join(recent_msgs)
        
        # 3. BUILD PROMPT
        system_prompt = f"""You are a helpful AI Assistant with an evolving rule system and knowledge base.
### RELEVANT MEMORIES (Context from past):
{long_term_txt}
### CURRENT CONVERSATION:
{recent_txt}
"""
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_msg}
        ]

        # 4. STREAM GENERATION
        full_response = ""
        for chunk in local_bot.chat_stream(messages):
            full_response += chunk
            yield chunk
        
        # 5. BACKGROUND: WRITE & REFLECT
        log_entry = f"User: {user_msg}\nAI: {full_response}"
        mem_history.write_entry(log_entry, q_vec)
        
    return Response(stream_with_context(generate()), mimetype='text/plain')

@app.route('/wipe', methods=['POST'])
def wipe():
    mem_history.wipe()
    return jsonify({"success":True})

if __name__ == '__main__':
    app.run(debug=True, host="0.0.0.0", port=7860)