import os, json, uuid
import requests
from flask import Flask, Response, request, jsonify, session, redirect, url_for, send_file
from authlib.integrations.flask_client import OAuth
from functools import wraps
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from markdown import markdown

from config import Config
from models import Base, Conversation, Message, Attachment

from flask_sock import Sock

app = Flask(__name__)
app.config.from_object(Config)
app.secret_key = app.config["SECRET_KEY"]

engine = create_engine(os.getenv("DATABASE_URL"))
SessionLocal = sessionmaker(bind=engine)
Base.metadata.create_all(engine)



sock = Sock(app)

@sock.route("/ws/notifications")
def notifications(ws):
    # All comments are in English.
    # Simple example: send a ping every time a message is added.
    while True:
        data = ws.receive()
        # You can interpret data or just echo.
        ws.send(json.dumps({"type": "pong", "payload": data}))


# ================
#   DECORATORS
# ================

import hashlib

def make_cache_key(model, prompt):
    raw = f"{model}:{prompt}"
    return hashlib.sha256(raw.encode()).hexdigest()

def get_cache(model, prompt):
    key = make_cache_key(model, prompt)
    db = SessionLocal()
    entry = db.query(Cache).filter_by(key=key).first()
    db.close()
    return entry.response if entry else None

def set_cache(model, prompt, response):
    key = make_cache_key(model, prompt)
    db = SessionLocal()
    entry = Cache(key=key, model=model, response=response)
    db.add(entry)
    db.commit()
    db.close()


import structlog
logger = structlog.get_logger()
@app.before_request
def log_request():
    logger.info(
        "request",
        path=request.path,
        method=request.method,
        remote=request.remote_addr,
    )

@app.after_request
def log_response(response):
    logger.info(
        "response",
        path=request.path,
        status=response.status_code,
    )
    return response


from limits import RateLimitItemPerMinute
from time import time

rate_store = {}

def rate_limited(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        if app.config["AUTH_MODE"] == "oauth":
            email = session["user"]["email"]
        else:
            email = session["db_user"]["email"]

        key = f"{email}:{f.__name__}"
        limit = RateLimitItemPerMinute(30)  # 30 req/min por endpoint

        now = int(time())
        window = now // 60

        used = rate_store.get((key, window), 0)
        if used >= limit.amount:
            return jsonify({"error": "Rate limit exceeded"}), 429

        rate_store[(key, window)] = used + 1
        return f(*args, **kwargs)
    return wrapper

def call_ollama_with_fallback(payload):
    for model in FALLBACK_MODELS:
        payload["model"] = model
        try:
            r = requests.post(f"{app.config['OLLAMA_BASE_URL']}/api/generate", json=payload)
            r.raise_for_status()
            return r.json(), model
        except Exception:
            continue
    raise Exception("All models failed")


def admin_required(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        db = SessionLocal()

        if app.config["AUTH_MODE"] == "oauth":
            email = session["user"]["email"]
        else:
            email = session["db_user"]["email"]

        u = db.query(User).filter_by(email=email).first()
        db.close()

        if not u or u.role != "admin":
            return jsonify({"error": "Admin only"}), 403

        return f(*args, **kwargs)
    return wrapper



def require_capability(cap):
    def decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            db = SessionLocal()
            email = session["user"]["email"]
            u = db.query(User).filter_by(email=email).first()
            db.close()

            if not u or not getattr(u, cap, False):
                return jsonify({"error": "Forbidden", "missing_capability": cap}), 403
            return f(*args, **kwargs)
        return wrapper
    return decorator


# ================
#       AUTH
# ================


import smtplib
from email.message import EmailMessage
import uuid

@app.route("/auth/forgot", methods=["POST"])
def forgot():
    email = request.form.get("email")

    db = SessionLocal()
    user = db.query(User).filter_by(email=email).first()
    if not user:
        db.close()
        return jsonify({"status": "ok"})  # não revela nada

    token = str(uuid.uuid4())
    user.reset_token = token
    db.commit()
    db.close()

    msg = EmailMessage()
    msg["Subject"] = "Password reset"
    msg["From"] = Config.SMTP_USER
    msg["To"] = email
    msg.set_content(f"Reset link: https://{Config.APP_DOMAIN}/auth/reset?token={token}")

    with smtplib.SMTP(Config.SMTP_HOST, Config.SMTP_PORT) as s:
        s.starttls()
        s.login(Config.SMTP_USER, Config.SMTP_PASS)
        s.send_message(msg)

    return jsonify({"status": "ok"})


@app.route("/auth/reset", methods=["POST"])
def reset():
    token = request.form.get("token")
    new_password = request.form.get("password")

    db = SessionLocal()
    user = db.query(User).filter_by(reset_token=token).first()
    if not user:
        db.close()
        return jsonify({"error": "Invalid token"}), 400

    user.password = new_password
    user.reset_token = None
    db.commit()
    db.close()

    return jsonify({"status": "ok"})


import pyotp, base64, os

@app.route("/auth/enable_2fa", methods=["POST"])
@login_required
def enable_2fa():
    db = SessionLocal()
    if app.config["AUTH_MODE"] == "oauth":
        email = session["user"]["email"]
    else:
        email = session["db_user"]["email"]

    u = db.query(User).filter_by(email=email).first()
    secret = base64.b32encode(os.urandom(10)).decode("utf-8")
    u.totp_secret = secret
    u.twofa_enabled = True
    db.commit()
    db.close()

    # Você mostra esse secret como QR code no frontend (otpauth:// URI)
    return jsonify({"secret": secret})


@app.route("/auth_mode")
def auth_mode():
    return jsonify({"mode": app.config["AUTH_MODE"]})

@app.route("/auth/login", methods=["GET", "POST"])
def db_login():
    if request.method == "GET":
        return send_file("/srv/auth/login.html")

    email = request.form.get("email")
    password = request.form.get("password")
    token = request.form.get("token")  # código 2FA

    db = SessionLocal()
    user = db.query(User).filter_by(email=email).first()

    if not user or user.password != password:
        db.close()
        return jsonify({"error": "Invalid credentials"}), 401

    if user.twofa_enabled:
        totp = pyotp.TOTP(user.totp_secret)
        if not totp.verify(token):
            db.close()
            return jsonify({"error": "Invalid 2FA token"}), 401

    db.close()
    session["db_user"] = {"email": user.email, "name": user.name, "role": user.role}
    return redirect("/")



oauth = OAuth(app)
google = oauth.register(
    name="google",
    client_id=app.config["GOOGLE_CLIENT_ID"],
    client_secret=app.config["GOOGLE_CLIENT_SECRET"],
    access_token_url="https://oauth2.googleapis.com/token",
    authorize_url="https://accounts.google.com/o/oauth2/v2/auth",
    api_base_url="https://www.googleapis.com/oauth2/v2/",
    client_kwargs={"scope": "openid email profile"},
)

def login_required(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        # If using OAuth
        if app.config["AUTH_MODE"] == "oauth":
            if "user" not in session:
                return redirect(url_for("login"))
            return f(*args, **kwargs)

        # If using DB authentication
        if "db_user" not in session:
            return redirect("/auth/login")
        return f(*args, **kwargs)
    return wrapper


@app.route("/auth/login", methods=["GET", "POST"])
def db_login():
    if request.method == "GET":
        return send_file("/srv/auth/login.html")

    email = request.form.get("email")
    password = request.form.get("password")

    db = SessionLocal()
    user = db.query(User).filter_by(email=email).first()
    db.close()

    if not user or user.password != password:
        return jsonify({"error": "Invalid credentials"}), 401

    session["db_user"] = {
        "email": user.email,
        "name": user.name,
        "role": user.role
    }

    return redirect("/")

@app.route("/auth/logout")
def db_logout():
    session.pop("db_user", None)
    return redirect("/auth/login")


@app.route("/auth/register", methods=["POST"])
def db_register():
    email = request.form.get("email")
    password = request.form.get("password")
    name = request.form.get("name")

    db = SessionLocal()
    if db.query(User).filter_by(email=email).first():
        return jsonify({"error": "User exists"}), 400

    u = User(
        email=email,
        name=name,
        password=password,
        role="user",
        can_use_audio=False,
        can_use_vision=False,
        can_share=True,
        can_import_docs=True,
        can_export_pdf=True,
        can_export_md=True,
    )
    db.add(u)
    db.commit()
    db.close()

    return jsonify({"status": "ok"})


@app.route("/login")
def login():
    return google.authorize_redirect(app.config["GOOGLE_REDIRECT_URI"])

@app.route("/oauth/callback")
def oauth_callback():
    token = google.authorize_access_token()
    user_info = google.get("userinfo").json()
    #session["user"] = user_info
    db = SessionLocal()
    u = db.query(User).filter_by(email=user_info["email"]).first()
    if not u:
        u = User(
            email=user_info["email"],
            name=user_info.get("name", ""),
            role="user",
            can_use_audio=False,
            can_use_vision=False,
            can_share=True,
        )
        db.add(u)
        db.commit()
    db.close()
    return redirect("/")

# ================
#      ADMIN
# ================

@app.route("/admin/users")
@login_required
@admin_required
def admin_users():
    db = SessionLocal()
    users = db.query(User).all()
    data = [{
        "email": u.email,
        "name": u.name,
        "role": u.role,
        "can_use_audio": u.can_use_audio,
        "can_use_vision": u.can_use_vision,
        "can_share": u.can_share,
        "can_import_docs": u.can_import_docs,
        "can_export_pdf": u.can_export_pdf,
        "can_export_md": u.can_export_md,
    } for u in users]
    db.close()
    return jsonify(data)

@app.route("/admin/update_acl/<email>", methods=["POST"])
@login_required
@admin_required
def update_acl(email):
    db = SessionLocal()
    u = db.query(User).filter_by(email=email).first()
    if not u:
        return jsonify({"error": "User not found"}), 404

    data = request.json
    for key, value in data.items():
        if hasattr(u, key):
            setattr(u, key, value)

    db.commit()
    db.close()
    return jsonify({"status": "ok"})


@app.route("/admin/promote/<email>", methods=["POST"])
@login_required
@admin_required
def promote(email):
    db = SessionLocal()
    u = db.query(User).filter_by(email=email).first()
    if not u:
        return jsonify({"error": "User not found"}), 404

    u.role = "admin"
    db.commit()
    db.close()
    return jsonify({"status": "ok"})

@app.route("/admin/conversations/<email>")
@login_required
@admin_required
def admin_conversations(email):
    db = SessionLocal()
    convs = db.query(Conversation).filter_by(user_email=email).all()
    data = [{"id": c.id, "title": c.title, "tags": c.tags} for c in convs]
    db.close()
    return jsonify(data)


# ================
#       API
# ================

@app.route("/api/chat", methods=["POST"])
@login_required
@rate_limited
def chat():
    data = request.json
    model = data.get("model", "llama3.2")
    messages = data.get("messages", [])

    url = f"{app.config['OLLAMA_BASE_URL']}/api/chat"
    payload = {"model": model, "messages": messages}

    r = requests.post(url, json=payload)
    r.raise_for_status()

    return jsonify(r.json())

@app.route("/api/chat/stream")
@login_required
@rate_limited
def chat_stream():
    model = request.args.get("model", "llama3.2")
    cid = request.args.get("conversation_id")

    db = SessionLocal()
    msgs = db.query(Message).filter_by(conversation_id=cid).order_by(Message.created_at).all()
    db.close()

    chatml = [{"role": m.role, "content": m.content} for m in msgs]

    def generate():
        url = f"{app.config['OLLAMA_BASE_URL']}/api/chat"
        payload = {"model": model, "messages": chatml, "stream": True}

        with requests.post(url, json=payload, stream=True) as r:
            for line in r.iter_lines():
                if not line:
                    continue
                data = json.loads(line.decode())
                token = data.get("message", {}).get("content", "")
                yield f"data: {json.dumps({'token': token})}\n\n"

    return Response(generate(), mimetype="text/event-stream")




@app.route("/api/new_conversation", methods=["POST"])
@login_required
def new_conversation():
    db = SessionLocal()
    cid = str(uuid.uuid4())
    title = request.json.get("title", "New conversation")
    conv = Conversation(id=cid, title=title)
    db.add(conv)
    db.commit()
    db.close()
    return jsonify({"conversation_id": cid, "title": title})

@app.route("/api/conversations")
@login_required
def list_conversations():
    db = SessionLocal()
    user_email = session["user"]["email"]
    convs = db.query(Conversation).filter(
        (Conversation.user_email == user_email) |
        (Conversation.shared == user_email)
    ).all()
    data = [{"id": c.id, "title": c.title, "tags": c.tags} for c in convs]
    db.close()
    return jsonify(data)

@app.route("/api/export/md/<cid>")
@login_required
@rate_limited
def export_md(cid):
    db = SessionLocal()
    conv = db.query(Conversation).filter_by(id=cid).first()
    msgs = conv.messages

    md = f"# {conv.title}\n\n"
    if conv.tags:
        md += f"**Tags:** {conv.tags}\n\n"

    for m in msgs:
        md += f"### {m.role.capitalize()}\n\n{m.content}\n\n"

    db.close()
    return Response(md, mimetype="text/markdown")

from weasyprint import HTML

import uuid
import pypandoc
from pdfminer.high_level import extract_text
from docx import Document

@app.route("/api/import/<cid>", methods=["POST"])
@login_required
@rate_limited
def import_file(cid):
    file = request.files.get("file")
    if not file:
        return jsonify({"error": "No file"}), 400

    ext = file.filename.lower().split(".")[-1]
    temp_path = f"/data/{uuid.uuid4()}.{ext}"
    file.save(temp_path)

    if ext == "pdf":
        text = extract_text(temp_path)
    elif ext == "docx":
        doc = Document(temp_path)
        text = "\n".join([p.text for p in doc.paragraphs])
    elif ext == "doc":
        text = pypandoc.convert_file(temp_path, "md")
    else:
        return jsonify({"error": "Unsupported format"}), 400

    md = pypandoc.convert_text(text, "md", format="plain")

    db = SessionLocal()
    msg = Message(
        id=str(uuid.uuid4()),
        conversation_id=cid,
        role="user",
        content=md
    )
    db.add(msg)
    db.commit()
    db.close()

    return jsonify({"status": "ok", "imported_as_markdown": True})


@app.route("/api/export/pdf/<cid>")
@login_required
@rate_limited
def export_pdf(cid):
    db = SessionLocal()
    conv = db.query(Conversation).filter_by(id=cid).first()
    msgs = conv.messages

    html = "<h1>{}</h1>".format(conv.title)
    if conv.tags:
        html += "<p><strong>Tags:</strong> {}</p>".format(conv.tags)

    for m in msgs:
        html += "<h3>{}</h3><p>{}</p>".format(m.role.capitalize(), m.content)

    db.close()

    pdf = HTML(string=html).write_pdf()
    return Response(pdf, mimetype="application/pdf")



@app.route("/api/history/<cid>")
@login_required
@rate_limited
def history(cid):
    db = SessionLocal()
    msgs = db.query(Message).filter_by(conversation_id=cid).order_by(Message.created_at).all()
    data = [{"role": m.role, "content": markdown(m.content)} for m in msgs]
    db.close()
    return jsonify(data)

@app.route("/api/attachment/<cid>", methods=["POST"])
@login_required
@rate_limited
def upload_attachment(cid):
    db = SessionLocal()
    file = request.files.get("file")
    if not file:
        return jsonify({"error": "No file"}), 400
    content = file.read().decode("latin1")
    att = Attachment(
        id=str(uuid.uuid4()),
        conversation_id=cid,
        filename=file.filename,
        mime_type=file.mimetype,
        data=content,
    )
    db.add(att)
    db.commit()
    db.close()
    return jsonify({"status": "ok"})

@app.route("/api/attachments/<cid>")
@login_required
@rate_limited
def list_attachments(cid):
    db = SessionLocal()
    atts = db.query(Attachment).filter_by(conversation_id=cid).all()
    data = [{"id": a.id, "filename": a.filename} for a in atts]
    db.close()
    return jsonify(data)


@app.route("/api/download/<att_id>")
@login_required
@rate_limited
def download(att_id):
    db = SessionLocal()
    att = db.query(Attachment).filter_by(id=att_id).first()
    db.close()

    if not att:
        return "Not found", 404

    return send_file(att.path, as_attachment=True, download_name=att.filename)


@app.route("/api/stream")
@login_required
@rate_limited
def stream():
    cid = request.args.get("conversation_id")
    model = request.args.get("model", "llama3.2")
    prompt = request.args.get("prompt")

    db = SessionLocal()
    db.add(Message(id=str(uuid.uuid4()), conversation_id=cid, role="user", content=prompt))
    db.commit()

    def generate():
        url = f"{app.config['OLLAMA_BASE_URL']}/api/generate"
        payload = {"model": model, "prompt": prompt, "stream": True}

        with requests.post(url, json=payload, stream=True) as r:
            buffer = ""
            for line in r.iter_lines():
                if not line:
                    continue
                try:
                    data = json.loads(line.decode())
                    token = data.get("response", "")
                    buffer += token
                    # streaming animation: send partial buffer
                    yield f"data: {json.dumps({'token': token})}\n\n"
                except:
                    continue
            db.add(Message(id=str(uuid.uuid4()), conversation_id=cid, role="assistant", content=buffer))
            db.commit()
            db.close()

    return Response(generate(), mimetype="text/event-stream")


@app.route("/api/share/<cid>", methods=["POST"])
@login_required
@require_capability("can_share")
def share(cid):
    target_email = request.json.get("email")
    db = SessionLocal()
    conv = db.query(Conversation).filter_by(id=cid).first()
    conv.shared = target_email
    db.commit()
    db.close()
    return jsonify({"status": "ok"})



@app.route("/api/tags/<cid>", methods=["POST"])
@login_required
def update_tags(cid):
    tags = request.json.get("tags", [])
    db = SessionLocal()
    conv = db.query(Conversation).filter_by(id=cid).first()
    conv.tags = ",".join(tags)
    db.commit()
    db.close()
    return jsonify({"status": "ok"})


@app.route("/api/search")
@login_required
@rate_limited
def search():
    q = request.args.get("q", "").lower()
    user_email = session["user"]["email"]

    db = SessionLocal()
    convs = db.query(Conversation).filter_by(user_email=user_email).all()

    results = []
    for c in convs:
        for m in c.messages:
            if q in m.content.lower():
                results.append({
                    "conversation_id": c.id,
                    "title": c.title,
                    "role": m.role,
                    "content": m.content,
                    "created_at": m.created_at.isoformat()
                })

    db.close()
    return jsonify(results)

from pywhispercpp.model import Model
@app.route("/api/audio/transcribe/<cid>", methods=["POST"])
@login_required
@require_capability("can_use_audio")
def transcribe(cid):
    # Initialize the model (automatically downloads 'base.en' if not present)
    model = Model('base.en', print_realtime=False, print_progress=False)
    
   
    file = request.files.get("file")
    if not file:
        return jsonify({"error": "No audio file"}), 400

    path = f"/data/audio-{uuid.uuid4()}.wav"
    file.save(path)

    # Example using whisper (pseudo-code)
    # result = whisper.transcribe(path)
    # text = result["text"]

    # Transcribe your audio file (must be 16kHz WAV format)
    segments = model.transcribe(path)
    
    # Print the text results
    for segment in segments:
        print(f"[{segment.t0} -> {segment.t1}]: {segment.text}")

    text = "Transcribed text placeholder"

    db = SessionLocal()
    msg = Message(
        id=str(uuid.uuid4()),
        conversation_id=cid,
        role="user",
        content=text,
    )
    db.add(msg)
    db.commit()
    db.close()

    return jsonify({"status": "ok", "text": text})

from gtts import gTTS

@app.route("/api/audio/synthesize/<cid>", methods=["POST"])
@login_required
@require_capability("can_use_audio")
def synthesize(cid):
    text = request.json.get("text", "")
    if not text:
        return jsonify({"error": "No text"}), 400

    tts = gTTS(text=text, lang="en")
    path = f"/data/tts-{uuid.uuid4()}.mp3"
    tts.save(path)

    return send_file(path, mimetype="audio/mpeg", as_attachment=False)


from PIL import Image

@app.route("/api/vision/<cid>", methods=["POST"])
@login_required
@require_capability("can_use_vision")
def vision(cid):
    file = request.files.get("file")
    if not file:
        return jsonify({"error": "No image"}), 400

    img_id = str(uuid.uuid4())
    path = f"/data/img-{img_id}.png"
    file.save(path)

    # Here you would call a vision-capable model (pseudo-code).
    # description = call_vision_model(path)

    description = "Image description placeholder"

    db = SessionLocal()
    msg = Message(
        id=str(uuid.uuid4()),
        conversation_id=cid,
        role="assistant",
        content=description,
    )
    db.add(msg)
    db.commit()
    db.close()

    return jsonify({"status": "ok", "description": description})



from tasks import generate_task

@app.route("/api/queue/generate", methods=["POST"])
@login_required
def queue_generate():
    data = request.json
    cid = data["conversation_id"]
    model = data["model"]
    prompt = data["prompt"]

    task = generate_task.delay(cid, model, prompt)

    return jsonify({"task_id": task.id})

from celery_app import celery

@app.route("/api/queue/status/<task_id>")
@login_required
def queue_status(task_id):
    result = celery.AsyncResult(task_id)
    return jsonify({
        "task_id": task_id,
        "state": result.state,
        "result": result.result if result.ready() else None
    })



# ***

import json
import requests
from flask import Flask, Response, request, jsonify, session, redirect, url_for
from authlib.integrations.flask_client import OAuth
from functools import wraps

from config import Config
from chat_store import ChatStore

app = Flask(__name__)
app.config.from_object(Config)
app.secret_key = app.config["SECRET_KEY"]

store = ChatStore()

oauth = OAuth(app)
google = oauth.register(
    name="google",
    client_id=app.config["GOOGLE_CLIENT_ID"],
    client_secret=app.config["GOOGLE_CLIENT_SECRET"],
    access_token_url="https://oauth2.googleapis.com/token",
    authorize_url="https://accounts.google.com/o/oauth2/v2/auth",
    api_base_url="https://www.googleapis.com/oauth2/v2/",
    client_kwargs={"scope": "openid email profile"},
)

def login_required(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        if "user" not in session:
            return redirect(url_for("login"))
        return f(*args, **kwargs)
    return wrapper

@app.route("/login")
def login():
    return google.authorize_redirect(app.config["GOOGLE_REDIRECT_URI"])

@app.route("/oauth/callback")
def oauth_callback():
    token = google.authorize_access_token()
    user_info = google.get("userinfo").json()
    session["user"] = user_info
    return redirect("/")

@app.route("/api/new_conversation")
@login_required
def new_conversation():
    cid = store.new_conversation()
    return jsonify({"conversation_id": cid})

@app.route("/api/history/<cid>")
@login_required
def history(cid):
    return jsonify(store.get_messages(cid))

@app.route("/api/stream")
@login_required
def stream():
    cid = request.args.get("conversation_id")
    model = request.args.get("model", "llama3.2")
    prompt = request.args.get("prompt")

    store.add_message(cid, "user", prompt)

    def generate():
        url = f"{app.config['OLLAMA_BASE_URL']}/api/generate"
        payload = {"model": model, "prompt": prompt, "stream": True}

        with requests.post(url, json=payload, stream=True) as r:
            for line in r.iter_lines():
                if not line:
                    continue
                try:
                    data = json.loads(line.decode())
                    token = data.get("response", "")
                    store.add_message(cid, "assistant", token)
                    yield f"data: {json.dumps({'token': token})}\n\n"
                except:
                    continue

    return Response(generate(), mimetype="text/event-stream")


# ***

import os
import json
import requests
from functools import wraps

from flask import Flask, request, jsonify, redirect, session, url_for
from authlib.integrations.flask_client import OAuth

from config import Config

# All comments are in English.

app = Flask(__name__)
app.config.from_object(Config)
app.secret_key = app.config["SECRET_KEY"]

oauth = OAuth(app)
google = oauth.register(
    name="google",
    client_id=app.config["GOOGLE_CLIENT_ID"],
    client_secret=app.config["GOOGLE_CLIENT_SECRET"],
    access_token_url="https://oauth2.googleapis.com/token",
    authorize_url="https://accounts.google.com/o/oauth2/v2/auth",
    authorize_params={"access_type": "offline", "prompt": "consent"},
    api_base_url="https://www.googleapis.com/oauth2/v2/",
    client_kwargs={"scope": "openid email profile"},
)


def login_required(f):
    """Decorator to ensure the user is authenticated via Google OAuth."""
    @wraps(f)
    def wrapper(*args, **kwargs):
        if "user" not in session:
            return redirect(url_for("login"))
        return f(*args, **kwargs)
    return wrapper


@app.route("/login")
def login():
    """Start Google OAuth login flow."""
    return google.authorize_redirect(app.config["Google_REDIRECT_URI"])


@app.route("/oauth/callback")
def oauth_callback():
    """Handle Google OAuth callback and store user session."""
    token = google.authorize_access_token()
    user_info = google.get("userinfo").json()

    session["user"] = {
        "email": user_info.get("email"),
        "name": user_info.get("name"),
    }

    return redirect(url_for("home"))


@app.route("/")
def home():
    return jsonify({
        "message": "Ollama Gateway is running.",
        "domain": "https://safira.renatorosa.com",
        "auth": "Google OAuth required for /api/* endpoints.",
        "endpoints": {
            "health": "/health",
            "generate": "/api/generate"
        }
    })


@app.route("/health")
def health():
    return jsonify({"status": "ok"})


@app.route("/api/generate", methods=["POST"])
@login_required
def generate():
    data = request.get_json(silent=True) or {}

    model_name = data.get("model", "llama3.2")
    prompt = data.get("prompt")

    if not prompt:
        return jsonify({"error": "Missing 'prompt' field."}), 400

    model_id = app.config["MODEL_MAP"].get(model_name)
    if not model_id:
        return jsonify({"error": f"Unsupported model '{model_name}'."}), 400

    # ollama_url = f"{app.config['OLLAMA_BASE_URL']}/api/generate"
    # payload = {"model": model_id, "prompt": prompt}
    # try:
    #     response = requests.post(ollama_url, json=payload, timeout=60)
    #     response.raise_for_status()
    # except requests.RequestException as e:
    #     return jsonify({"error": "Failed to call Ollama API.", "details": str(e)}), 502
    # try:
    #     ollama_data = response.json()
    # except json.JSONDecodeError:
    #     return jsonify({"error": "Invalid JSON response from Ollama."}), 502
    #output_text = ollama_data.get("response") or ollama_data.get("output") or ""
    #return jsonify({"model": model_name, "output": output_text})
   
    cached = get_cache(model, prompt)
    if cached:
        return jsonify({"model": model, "output": cached, "cached": True})
    
    result, used_model = call_ollama_with_fallback({"prompt": prompt})
    output = result.get("response")
    
    set_cache(used_model, prompt, output)
    
    return jsonify({"model": used_model, "output": output, "cached": False})

@app.route("/api/multi", methods=["POST"])
@login_required
@rate_limited
def multi():
    data = request.json
    models = data.get("models", ["llama3.2", "gemma4"])
    prompt = data.get("prompt")

    results = {}

    for m in models:
        try:
            r = requests.post(
                f"{app.config['OLLAMA_BASE_URL']}/api/generate",
                json={"model": m, "prompt": prompt}
            )
            r.raise_for_status()
            results[m] = r.json().get("response")
        except Exception as e:
            results[m] = f"Error: {str(e)}"

    return jsonify(results)

@app.route("/api/multi/stream")
@login_required
def multi_stream():
    models = request.args.get("models", "llama3.2,gemma4").split(",")
    prompt = request.args.get("prompt")

    def generate():
        for m in models:
            yield f"event: model\n"
            yield f"data: {json.dumps({'model': m})}\n\n"

            url = f"{app.config['OLLAMA_BASE_URL']}/api/generate"
            payload = {"model": m, "prompt": prompt, "stream": True}

            with requests.post(url, json=payload, stream=True) as r:
                for line in r.iter_lines():
                    if not line:
                        continue
                    data = json.loads(line.decode())
                    token = data.get("response", "")
                    yield f"event: token\n"
                    yield f"data: {json.dumps({'model': m, 'token': token})}\n\n"

    return Response(generate(), mimetype="text/event-stream")



if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8000)

# ***

import os
import json
import requests
from functools import wraps

from flask import Flask, request, jsonify, redirect, session, url_for
from authlib.integrations.flask_client import OAuth

from config import Config

# All comments are in English.

app = Flask(__name__)
app.config.from_object(Config)

# Configure session secret key.
app.secret_key = app.config["SECRET_KEY"]

# Configure OAuth with Google.
oauth = OAuth(app)
google = oauth.register(
    name="google",
    client_id=app.config["GOOGLE_CLIENT_ID"],
    client_secret=app.config["GOOGLE_CLIENT_SECRET"],
    access_token_url="https://oauth2.googleapis.com/token",
    access_token_params=None,
    authorize_url="https://accounts.google.com/o/oauth2/v2/auth",
    authorize_params={
        "access_type": "offline",
        "prompt": "consent",
    },
    api_base_url="https://www.googleapis.com/oauth2/v2/",
    client_kwargs={"scope": "openid email profile"},
)


def login_required(f):
    """Decorator to ensure the user is authenticated via Google OAuth."""

    @wraps(f)
    def wrapper(*args, **kwargs):
        # Check if user info is stored in session.
        if "user" not in session:
            # Redirect to login if not authenticated.
            return redirect(url_for("login"))
        return f(*args, **kwargs)

    return wrapper


@app.route("/login")
def login():
    """Start Google OAuth login flow."""
    redirect_uri = app.config["GOOGLE_REDIRECT_URI"]
    # Use the configured redirect URI for OAuth callback.
    return google.authorize_redirect(redirect_uri)


@app.route("/oauth/callback")
def oauth_callback():
    """Handle Google OAuth callback and store user session."""
    # Exchange authorization code for tokens.
    token = google.authorize_access_token()
    # Fetch user info from Google.
    resp = google.get("userinfo")
    user_info = resp.json()

    # Store minimal user info in session.
    session["user"] = {
        "email": user_info.get("email"),
        "name": user_info.get("name"),
    }

    # Redirect to a simple home or API docs page.
    return redirect(url_for("home"))


@app.route("/")
def home():
    """Simple home endpoint explaining the service."""
    # This endpoint is intentionally simple and easy to understand.
    return jsonify(
        {
            "message": "Ollama Gateway is running.",
            "auth": "Google OAuth required for /api/* endpoints.",
            "endpoints": {
                "health": "/health",
                "generate": "/api/generate",
            },
        }
    )


@app.route("/health")
def health():
    """Health check endpoint."""
    return jsonify({"status": "ok"})


@app.route("/api/generate", methods=["POST"])
@login_required
def generate():
    """Generate text using a selected Ollama model.

    Expected JSON body:
    {
        "model": "llama3.2",
        "prompt": "Your prompt here"
    }
    """
    data = request.get_json(silent=True) or {}

    # Extract model and prompt from request.
    model_name = data.get("model", "llama3.2")
    prompt = data.get("prompt")

    if not prompt:
        return jsonify({"error": "Missing 'prompt' field."}), 400

    # Map user-facing model name to Ollama model identifier.
    model_id = app.config["MODEL_MAP"].get(model_name)
    if not model_id:
        return jsonify({"error": f"Unsupported model '{model_name}'."}), 400

    # Prepare request to Ollama API.
    ollama_url = f"{app.config['OLLAMA_BASE_URL']}/api/generate"
    payload = {
        "model": model_id,
        "prompt": prompt,
        # You can add more Ollama-specific parameters here if needed.
    }

    try:
        # Call Ollama HTTP API.
        response = requests.post(ollama_url, json=payload, timeout=60)
        response.raise_for_status()
    except requests.RequestException as e:
        # Handle network or API errors gracefully.
        return jsonify({"error": "Failed to call Ollama API.", "details": str(e)}), 502

    # Parse Ollama response.
    try:
        ollama_data = response.json()
    except json.JSONDecodeError:
        return jsonify({"error": "Invalid JSON response from Ollama."}), 502

    # Extract output text (this depends on Ollama's response format).
    output_text = ollama_data.get("response") or ollama_data.get("output") or ""

    return jsonify(
        {
            "model": model_name,
            "output": output_text,
        }
    )


if __name__ == "__main__":
    # This block is mainly for local debugging.
    # In production, Gunicorn will run the app.
    app.run(host="0.0.0.0", port=8000)


# ***


from flask import Flask, request, jsonify
from flask_cors import CORS
from db import Base, engine
from models import User
from services.intent_service import IntentService
from services.mode_service import ModeService
from services.preference_service import PreferenceService
from services.suggestion_service import SuggestionService
from services.document_service import DocumentService
from services.search_service import SearchService
from services.monitoring_service import MonitoringService
from services.export_service import ExportService
from services.audit_service import AuditService
from services.security_service import SecurityService

from services.ollama_service import OllamaService
ollama_service = OllamaService()

from services.model_router import ModelRouter
router = ModelRouter()
generated = router.generate(intent, mode, prompt)


app = Flask(__name__)
CORS(app)

Base.metadata.create_all(bind=engine)

intent_service = IntentService()
mode_service = ModeService()
preference_service = PreferenceService()
suggestion_service = SuggestionService()
document_service = DocumentService()
search_service = SearchService()
monitoring_service = MonitoringService()
export_service = ExportService()
audit_service = AuditService()
security_service = SecurityService()

USER_ID = 1

@app.route("/api/conversation", methods=["POST"])
def conversation():
    data = request.json
    message = data.get("message", "")

    mode = mode_service.get_current_mode(USER_ID)
    if not mode:
        reply = (
            "{user.name}, como você quer trabalhar hoje?\n\n"
            "- Modo Advogado\n- Modo Corporativo\n- Apenas Conversar\n\n"
            "Me diga qual modo você prefere agora."
        )
        return jsonify({"reply": reply, "suggestions": [], "quick_actions": []})

    intent = intent_service.detect_intent(message)
    prefs = preference_service.get_preferences(USER_ID)

    if intent == "set_mode":
        mode_service.set_mode(USER_ID, intent_service.extract_mode(message))
        reply = f"Perfeito, vou atuar no modo {mode_service.get_current_mode(USER_ID)}. O que vamos fazer agora?"
        return jsonify({"reply": reply, "suggestions": [], "quick_actions": []})

    if intent == "document":
        response = document_service.handle(message, mode, prefs)
    elif intent == "search":
        response = search_service.handle(message, mode, prefs)
    elif intent == "monitoring":
        response = monitoring_service.handle(message, mode, prefs)
    else:
        response = document_service.chat_like(message, mode, prefs)

    suggestions = suggestion_service.generate_suggestions(response, intent, mode, prefs)
    quick_actions = suggestion_service.generate_quick_actions(response, intent, mode, prefs)

    return jsonify({
        "reply": response["text"],
        "suggestions": suggestions,
        "quick_actions": quick_actions
    })

if __name__ == "__main__":
    app.run(debug=True)