chat_history.py

from __future__ import annotations

import json
import os

from dotenv import load_dotenv
from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine


load_dotenv()

_engine: Engine | None = None


def get_engine() -> Engine:
    global _engine
    if _engine is None:
        database_url = os.getenv("DATABASE_URL", "").strip("'\"")
        if not database_url:
            raise RuntimeError("DATABASE_URL is missing in .env")
        _engine = create_engine(database_url, pool_pre_ping=True)
    return _engine


def list_users() -> list[dict]:
    query = text("SELECT id, username FROM auth_user ORDER BY username")
    with get_engine().connect() as connection:
        return [dict(row) for row in connection.execute(query).mappings()]


def list_conversations(user_id: int) -> list[dict]:
    query = text("""
        SELECT id, title, created_at, updated_at
        FROM chat_conversation
        WHERE user_id = :user_id
        ORDER BY updated_at DESC
    """)
    with get_engine().connect() as connection:
        return [dict(row) for row in connection.execute(query, {"user_id": user_id}).mappings()]


def create_conversation(user_id: int, title: str) -> dict:
    query = text("""
        INSERT INTO chat_conversation (user_id, title, created_at, updated_at)
        VALUES (:user_id, :title, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
        RETURNING id, title, created_at, updated_at
    """)
    with get_engine().begin() as connection:
        user_exists = connection.execute(
            text("SELECT 1 FROM auth_user WHERE id = :user_id"),
            {"user_id": user_id},
        ).scalar_one_or_none()
        if not user_exists:
            raise ValueError("User does not exist")
        return dict(connection.execute(query, {"user_id": user_id, "title": title[:255]}).mappings().one())


def ensure_conversation_owner(conversation_id: int, user_id: int) -> None:
    query = text("""
        SELECT 1 FROM chat_conversation
        WHERE id = :conversation_id AND user_id = :user_id
    """)
    with get_engine().connect() as connection:
        if connection.execute(query, {"conversation_id": conversation_id, "user_id": user_id}).scalar_one_or_none() is None:
            raise ValueError("Conversation does not exist for this user")


def list_messages(conversation_id: int, user_id: int, limit: int = 100) -> list[dict]:
    ensure_conversation_owner(conversation_id, user_id)
    query = text("""
        SELECT id, role, content, metadata, created_at
        FROM (
            SELECT id, role, content, metadata, created_at
            FROM chat_message
            WHERE conversation_id = :conversation_id
            ORDER BY created_at DESC, id DESC
            LIMIT :limit
        ) recent
        ORDER BY created_at, id
    """)
    with get_engine().connect() as connection:
        return [dict(row) for row in connection.execute(
            query,
            {"conversation_id": conversation_id, "limit": limit},
        ).mappings()]


def add_message(conversation_id: int, role: str, content: str, metadata: dict | None = None) -> dict:
    if role not in {"user", "assistant", "system", "tool"}:
        raise ValueError("Invalid message role")
    query = text("""
        INSERT INTO chat_message (conversation_id, role, content, metadata, created_at)
        VALUES (:conversation_id, :role, :content, CAST(:metadata AS jsonb), CURRENT_TIMESTAMP)
        RETURNING id, role, content, metadata, created_at
    """)
    with get_engine().begin() as connection:
        result = connection.execute(query, {
            "conversation_id": conversation_id,
            "role": role,
            "content": content,
            "metadata": json.dumps(metadata or {}),
        }).mappings().one()
        connection.execute(
            text("UPDATE chat_conversation SET updated_at = CURRENT_TIMESTAMP WHERE id = :id"),
            {"id": conversation_id},
        )
        return dict(result)


def format_history(messages: list[dict]) -> str:
    labels = {"user": "Usuario", "assistant": "Asistente", "system": "Sistema", "tool": "Herramienta"}
    return "\n".join(f"{labels.get(item['role'], item['role'])}: {item['content']}" for item in messages)