from flask import Flask, request, jsonify, Response, redirect, url_for, render_template, send_from_directory
from flask_cors import CORS
from functools import wraps
from bson import ObjectId
from datetime import datetime
from pymongo import MongoClient
from config import CHATBOT_API_TOKEN, DOCS_USERNAME, DOCS_PASSWORD, DEBUG
from bootstrap import run_workflow
from connection import chat_collection, uri, database
import uuid, secrets, base64, re, json
from flasgger import Swagger
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address

app = Flask(__name__, static_folder='static', template_folder='templates')
CORS(app)

limiter = Limiter(
    get_remote_address,
    app=app,
    default_limits=["200 per day", "50 per hour"],
    storage_uri="memory://",  # "redis://localhost:6379" 
)

swagger = Swagger(app, config={
    "headers": [],
    "specs": [
        {
            "endpoint": 'apispec_1',
            "route": '/apispec_1.json',
            "rule_filter": lambda rule: True,
            "model_filter": lambda tag: True,
        }
    ],
    "static_url_path": "/flasgger_static",
    "swagger_ui": True,
    "specs_route": "/apidocs/",
    # This part defines the "Authorize" button and Bearer token security scheme
    "securityDefinitions": {
        "Bearer": {
            "type": "apiKey",
            "name": "Authorization",
            "in": "header",
            "description": "Enter your Bearer token in the format **Bearer <token>**"
        }
    }
})

def debug(msg):
    if DEBUG:
        print(f"Api => {msg}")

def requires_auth(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        auth = request.authorization
        if not auth or not (secrets.compare_digest(auth.username, DOCS_USERNAME) and secrets.compare_digest(auth.password, DOCS_PASSWORD)):
            return Response("Credenciais inválidas", 401, {"WWW-Authenticate": "Basic realm='Login Required'"})
        return f(*args, **kwargs)
    return decorated

@app.before_request
def protect_swagger_ui():
    """
    Protects Swagger UI endpoints (/apidocs, /flasgger_static, /apispec_1.json)
    with Basic Authentication.
    """
    if request.endpoint and ('flasgger' in request.endpoint or request.endpoint == 'apispec_1'):
        auth = request.authorization
        if not auth or not (secrets.compare_digest(auth.username, DOCS_USERNAME) and secrets.compare_digest(auth.password, DOCS_PASSWORD)):
            return Response("Credenciais inválidas", 401, {"WWW-Authenticate": "Basic realm='Login Required'"})

def verify_token():
    auth_header = request.headers.get("Authorization")
    # The check should be for "Bearer <token>"
    if not auth_header or not auth_header.startswith("Bearer "):
        return jsonify({"detail": "Authorization header must be in Bearer format"}), 401
    
    token = auth_header.split(" ")[1]
    if token != CHATBOT_API_TOKEN:
        return jsonify({"detail": "Unauthorized"}), 401
    # Return None if token is valid
    return None

@app.route("/")
def root():
  return redirect("/apidocs/")

@app.route("/generate-session-id")
def generate_session_id():
    """
    Gera um novo session_id
    ---
    tags:
      - Sessão
    security:
      - Bearer: []
    responses:
      200:
        description: Session ID gerado com sucesso
      401:
        description: Unauthorized
    """
    # UPDATED: Added token verification
    auth = verify_token()
    if auth:
        return auth
    return jsonify({"session_id": str(uuid.uuid4())})



@app.route("/docs")
@requires_auth
def docs():
    return redirect("/apidocs/")

@app.route("/chat", methods=["POST"])
@limiter.limit("15 per minute")
def chat():
    """
    Executa o agente de chatbot
    ---
    tags:
      - Chat
    parameters:
      - name: body
        in: body
        required: true
        schema:
          type: object
          required:
            - message
            - session_id
          properties:
            message:
              type: string
            session_id:
              type: string
            source:
              type: string
    responses:
      200:
        description: Resposta gerada pelo agente
    """

    req = request.json
    result = run_workflow(
        user_message=req.get("message"),
        session_id=req.get("session_id", str(uuid.uuid4())),
        source=req.get("source", "api")
    )
    return jsonify(result)

@app.route("/reset-session", methods=["POST"])
def reset_session():
    """
    Reseta os dados da sessão no MongoDB e LangGraph
    ---
    tags:
      - Sessão
    security:
      - Bearer: []
    parameters:
      - name: session_id
        in: query
        required: true
        type: string
    responses:
      200:
        description: Sessão resetada com sucesso
      401:
        description: Unauthorized
    """
    auth = verify_token()
    if auth:
        return auth

    session_id = request.args.get("session_id")
    client = MongoClient(uri)
    checkpointer = client[database]

    try:
        chat_delete = chat_collection.delete_many({"session_id": session_id})
        checkpoint_delete = checkpointer["checkpoint"].delete_many({"thread_id": session_id})
        writes_delete = checkpointer["checkpoint_writes"].delete_many({"thread_id": session_id})

        debug(f"Chat history deleted: {chat_delete.deleted_count}")
        debug(f"Checkpoints deleted: {checkpoint_delete.deleted_count}")
        debug(f"Checkpoint writes deleted: {writes_delete.deleted_count}")

        return jsonify({
            "status": "✅ Session reset complete.",
            "session_id": session_id,
            "deleted": {
                "chat_history": chat_delete.deleted_count,
                "checkpoints": checkpoint_delete.deleted_count,
                "checkpoint_writes": writes_delete.deleted_count
            },
            "timestamp": datetime.now().isoformat()
        })
    except Exception as e:
        return jsonify({"detail": str(e)}), 500
    finally:
        client.close()

@app.route("/store", methods=["POST"])
def store():
    """
    Armazena uma interação do chat manualmente
    ---
    tags:
      - Chat
    security:
      - Bearer: []
    parameters:
      - name: body
        in: body
        required: true
        schema:
          type: object
          required:
            - session_id
            - user_input
            - agent_response
          properties:
            session_id:
              type: string
            user_input:
              type: string
            agent_response:
              type: string
            source:
              type: string
            history:
              type: array
              items:
                type: object
    responses:
      200:
        description: Interação armazenada com sucesso
      401:
        description: Unauthorized
    """
    auth = verify_token()
    if auth:
        return auth

    data = request.json
    chat_collection.insert_one({
        "session_id": data["session_id"],
        "source": data.get("source", "unknown"),
        "timestamp": datetime.now(),
        "user_input": data["user_input"],
        "agent_response": data["agent_response"],
        "full_history": data.get("history", [])
    })
    return jsonify({"status": "ok"})

@app.route("/sessions")
def get_sessions():
    """
    Lista sessões registradas
    ---
    tags:
      - Sessão
    security:
      - Bearer: []
    parameters:
      - name: search
        in: query
        required: false
        type: string
    responses:
      200:
        description: Lista de sessões
      401:
        description: Unauthorized
    """
    auth = verify_token()
    if auth:
        return auth

    search = request.args.get("search")
    match_stage = {}
    if search:
        search_regex = re.compile(re.escape(search), re.IGNORECASE)
        match_stage = {"$or": [
            {"user_input": {"$regex": search_regex}},
            {"agent_response": {"$regex": search_regex}},
            {"full_history.content": {"$regex": search_regex}}
        ]}

    pipeline = []
    if match_stage:
        pipeline.append({"$match": match_stage})
    pipeline += [
        {"$sort": {"timestamp": -1}},
        {"$group": {
            "_id": "$session_id",
            "session_id": {"$first": "$session_id"},
            "source": {"$first": "$source"},
            "timestamp": {"$first": "$timestamp"}
        }},
        {"$sort": {"timestamp": -1}}
    ]

    sessions = list(chat_collection.aggregate(pipeline))
    for r in sessions:
        r["_id"] = str(r["_id"])
        if hasattr(r["timestamp"], "isoformat"):
            r["timestamp"] = r["timestamp"].isoformat()
    return jsonify(sessions)

@app.route("/sessions/<session_id>/history")
def session_history(session_id):
    """
    Retorna o histórico completo de uma sessão
    ---
    tags:
      - Sessão
    security:
      - Bearer: []
    parameters:
      - name: session_id
        in: path
        required: true
        type: string
    responses:
      200:
        description: Histórico da sessão
      401:
        description: Unauthorized
    """
    auth = verify_token()
    if auth:
        return auth

    results = list(chat_collection.find({"session_id": session_id}).sort("timestamp", 1))
    for r in results:
        r["_id"] = str(r["_id"])
        if hasattr(r["timestamp"], "isoformat"):
            r["timestamp"] = r["timestamp"].isoformat()
    return jsonify(results)

@app.route("/health")
def health():
    """
    Verifica o status da API
    ---
    tags:
      - Health Check
    responses:
      200:
        description: OK
    """
    return jsonify({"status": "ok"})



@app.route('/chatbot')
def chatbot_ui():
    from connection import chat_collection
    session_id = request.args.get('session_id')
    chat_history = []

    if session_id:
        try:
            results = list(chat_collection.find({"session_id": session_id}).sort("timestamp", 1))
            for r in results:
                chat_history.append({"role": "user", "content": r.get("user_input", "")})
                chat_history.append({"role": "agent", "content": r.get("agent_response", "")})
        except Exception as e:
            print(f"Erro ao carregar histórico: {e}")

    return render_template('chatbot.html', session_id=session_id, chat_history=chat_history, history_json=json.dumps(chat_history))



# Servir arquivos estáticos (caso necessário)
@app.route('/static/<path:filename>')
def static_files(filename):
    return send_from_directory(app.static_folder, filename)


if __name__ == '__main__':
    app.run(debug=DEBUG, port=8001)