from flask import Flask, render_template, request, jsonify, session, stream_with_context, Response
from dotenv import load_dotenv
from threading import Lock
import os
import secrets
import logging as py_logging
from Lawverse.pipeline.rag_pipeline import rag_components
from Lawverse.pipeline.llm_loader import llm
from Lawverse.memory.langchain_memory import ChatMemory
from Lawverse.logger import logging
from Lawverse.monitoring.dashboard import monitor_bp
from Lawverse.agents.graph import create_agentic_chain
from Lawverse.storage.factory import get_chat_store
from api.auth import auth_bp, login_required
for logger_name in [
    "httpcore", "httpx", "hpack", "filelock", "sentence_transformers",
    "urllib3", "openai", "openai._base_client", "faiss", "datasets",
]:
    py_logging.getLogger(logger_name).setLevel(py_logging.WARNING)


load_dotenv()
app = Flask(__name__, template_folder="../templates")
app.secret_key = os.getenv("SECRET_KEY") or secrets.token_hex(32)

app.register_blueprint(auth_bp)
app.register_blueprint(monitor_bp)

BASE_COMPONENTS = None
BASE_COMPONENTS_LOCK = Lock()
AGENT_CHAIN = None
AGENT_CHAIN_LOCK = Lock()
active_chains = {}

def get_base_components():
    global BASE_COMPONENTS
    if BASE_COMPONENTS is not None:
        return BASE_COMPONENTS

    with BASE_COMPONENTS_LOCK:
        if BASE_COMPONENTS is None:
            logging.info("Loading Lawverse RAG base components...")
            BASE_COMPONENTS = rag_components()
            logging.info("Lawverse RAG base components loaded successfully.")
    return BASE_COMPONENTS


def get_agent_chain():
    global AGENT_CHAIN
    if AGENT_CHAIN is not None:
        return AGENT_CHAIN

    with AGENT_CHAIN_LOCK:
        if AGENT_CHAIN is None:
            components = get_base_components()
            AGENT_CHAIN = create_agentic_chain(components, llm)
    return AGENT_CHAIN


def create_agent_session(chat_id=None, save_on_create=True):
    chain = get_agent_chain()
    memory_manager = ChatMemory(chat_id=chat_id)
    active_chains[memory_manager.chat_id] = (chain, memory_manager)
    session["chat_id"] = memory_manager.chat_id
    if save_on_create:
        memory_manager.save_memory()
    return chain, memory_manager

@app.route("/", methods=["GET"])
def home():
    return render_template("index.html")

@app.route("/chat", methods=["GET"])
@login_required
def chat():
    chat_id = session.get("chat_id")
    if not chat_id or chat_id not in active_chains:
        create_agent_session()
    return render_template("chat.html")

@app.route("/new_chat", methods=["POST"])
@login_required
def new_chat():
    _, memory_manager = create_agent_session()
    return jsonify({"chat_id": memory_manager.chat_id, "title": memory_manager._get_title()})


@app.route("/response", methods=["POST"])
@login_required
def rag_response():
    try:
        chat_id = session.get("chat_id")
        if not chat_id or chat_id not in active_chains:
            _, memory_manager = create_agent_session()
            chat_id = memory_manager.chat_id

        qa, memory_manager = active_chains[chat_id]
        data = request.get_json(silent=True) or {}
        query = data.get("message", "").strip()
        if not query:
            return jsonify({"error": "Empty message"}), 400

        def generate():
            answer_parts = []
            generation_failed = False
            try:
                for chunk in qa.stream({
                    "input": query,
                    "chat_history": memory_manager.memory.chat_memory.messages,
                }):
                    text = chunk if isinstance(chunk, str) else str(chunk)
                    answer_parts.append(text)
                    yield text
            except Exception as e:
                generation_failed = True
                logging.error(f"Error during answer generation: {e}")

            if not generation_failed:
                full_answer = "".join(answer_parts).strip()
                memory_manager.append_exchange(query, full_answer)
                memory_manager.save_memory()

        return Response(stream_with_context(generate()), mimetype="text/plain")

    except Exception as e:
        logging.error(f"Chat error: {e}")
        return jsonify({"error": str(e)}), 500


@app.route("/get_chats", methods=["GET"])
@login_required
def get_chats():
    try:
        user_id = str(session.get("user_id"))
        chats = get_chat_store().list_chats(user_id)
        return jsonify(chats), 200
    except Exception as e:
        logging.error(f"Failed to list chats from cloud store: {e}")
        return jsonify([]), 200


@app.route("/load_chat/<chat_id>", methods=["POST"])
@login_required
def load_chat(chat_id):
    try:
        user_id = str(session.get("user_id"))
        data = get_chat_store().load_chat(user_id, chat_id)
        if not data:
            return jsonify({"error": "Chat not found"}), 404

        _, memory_manager = create_agent_session(chat_id=chat_id, save_on_create=False)
        messages_list = memory_manager.memory.chat_memory.messages
        messages = []
        for i in range(0, len(messages_list), 2):
            user_msg = messages_list[i].content if i < len(messages_list) else None
            ai_msg = messages_list[i + 1].content if i + 1 < len(messages_list) else ""
            if user_msg:
                messages.append({"user": user_msg, "ai": ai_msg})

        return jsonify({
            "chat_id": chat_id,
            "title": memory_manager._get_title(),
            "messages": messages,
        }), 200
    except Exception as e:
        logging.error(f"Failed to load chat from cloud store: {e}")
        return jsonify({"error": "Internal Server Error"}), 500


@app.route("/delete_chat/<chat_id>", methods=["DELETE"])
@login_required
def delete_chat(chat_id):
    try:
        user_id = str(session.get("user_id"))
        deleted = get_chat_store().delete_chat(user_id, chat_id)
        was_active = chat_id in active_chains
        if was_active:
            del active_chains[chat_id]
        if session.get("chat_id") == chat_id:
            session.pop("chat_id", None)
        return jsonify({"success": deleted, "was_active": was_active}), 200
    except Exception as e:
        logging.error(f"Error deleting cloud chat {chat_id}: {e}")
        return jsonify({"error": "Internal Server Error"}), 500

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=int(os.getenv("PORT", 7860)))