app.py

from flask import Flask, request, jsonify, render_template, flash, redirect, url_for
from flask_login import LoginManager, UserMixin, login_user, logout_user, login_required, current_user
from werkzeug.security import generate_password_hash, check_password_hash
from transformers import pipeline
import torch
from pydub import AudioSegment
import os
import io
import uuid
from datetime import datetime
import sqlite3
from pathlib import Path
import whisper
from extensions import db, login_manager


instance_path = Path(__file__).parent / 'instance'
instance_path.mkdir(exist_ok=True, mode=0o755)

app = Flask(__name__)
app.secret_key = 'очень_сложный_секретный_ключ_здесь'
db_path = instance_path / 'chats.db'
app.config['SQLALCHEMY_DATABASE_URI'] = f'sqlite:///{db_path}'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

# Инициализация Flask-Login
db.init_app(app)
login_manager.init_app(app)
login_manager.login_view = 'auth_bp.login'

# Инициализация моделей
def init_models():
    try:
        emotion_map = {
            'joy': '😊 Радость',
            'neutral': '😐 Нейтрально',
            'anger': '😠 Злость',
            'sadness': '😢 Грусть',
            'surprise': '😲 Удивление'
        }

        speech_to_text_model = whisper.load_model("base")
        text_classifier = pipeline(
            "text-classification",
            model="cointegrated/rubert-tiny2-cedr-emotion-detection"
        )
        audio_classifier = pipeline(
            "audio-classification",
            model="superb/hubert-large-superb-er"
        )

        return {
            'emotion_map': emotion_map,
            'speech_to_text_model': speech_to_text_model,
            'text_classifier': text_classifier,
            'audio_classifier': audio_classifier
        }
    except Exception as e:
        print(f"Ошибка загрузки моделей: {e}")
        return None

models = init_models()
if not models:
    raise RuntimeError("Не удалось загрузить модели")

# Импорт Blueprint
from auth import auth_bp
from profile import profile_bp

app.register_blueprint(auth_bp)
app.register_blueprint(profile_bp)

# Делаем переменные доступными
emotion_map = models['emotion_map']
speech_to_text_model = models['speech_to_text_model']
text_classifier = models['text_classifier']
audio_classifier = models['audio_classifier']


def transcribe_audio(audio_path):
    """Преобразование аудио в текст с помощью Whisper"""
    if not speech_to_text_model:
        return None
    try:
        result = speech_to_text_model.transcribe(audio_path, language="ru")
        return result["text"]
    except Exception as e:
        print(f"Ошибка преобразования аудио в текст: {e}")
        return None


# Инициализация Flask-Login
login_manager = LoginManager(app)
login_manager.login_view = 'login'


# Модель пользователя для Flask-Login
class User(UserMixin):
    def __init__(self, id, username, email, password_hash):
        self.id = id
        self.username = username
        self.email = email
        self.password_hash = password_hash

    def check_password(self, password):
        return check_password_hash(self.password_hash, password)


@login_manager.user_loader
def load_user(user_id):
    conn = get_db_connection()
    user = conn.execute(
        "SELECT id, username, email, password_hash FROM users WHERE id = ?",
        (user_id,)
    ).fetchone()
    conn.close()
    if user:
        return User(id=user['id'], username=user['username'], email=user['email'], password_hash=user['password_hash'])
    return None


# Инициализация БД
def get_db_connection():
    instance_path = Path('instance')
    instance_path.mkdir(exist_ok=True)
    db_path = instance_path / 'chats.db'
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    return conn


def init_db():
    conn = get_db_connection()
    try:
        conn.execute('''
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                email TEXT UNIQUE NOT NULL,
                password_hash TEXT NOT NULL,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        conn.execute('''
            CREATE TABLE IF NOT EXISTS chats (
                chat_id TEXT PRIMARY KEY,
                user_id INTEGER,
                created_at TEXT,
                title TEXT,
                FOREIGN KEY(user_id) REFERENCES users(id)
            )
        ''')
        conn.execute('''
            CREATE TABLE IF NOT EXISTS messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                chat_id TEXT,
                sender TEXT,
                content TEXT,
                timestamp TEXT,
                FOREIGN KEY(chat_id) REFERENCES chats(chat_id)
            )
        ''')
        conn.execute('''
            CREATE TABLE IF NOT EXISTS analysis_reports (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER,
                content TEXT,
                emotion TEXT,
                confidence REAL,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY(user_id) REFERENCES users(id)
            )
        ''')
        conn.commit()
    finally:
        conn.close()


init_db()


# Маршруты аутентификации
@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        email = request.form.get('email')
        password = request.form.get('password')

        conn = get_db_connection()
        user = conn.execute(
            "SELECT id, username, email, password_hash FROM users WHERE email = ?",
            (email,)
        ).fetchone()
        conn.close()

        if user and check_password_hash(user['password_hash'], password):
            user_obj = User(id=user['id'], username=user['username'],
                            email=user['email'], password_hash=user['password_hash'])
            login_user(user_obj)
            return redirect(url_for('index'))

        flash('Неверный email или пароль', 'danger')

    return render_template('auth/login.html')


@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form.get('username')
        email = request.form.get('email')
        password = request.form.get('password')
        confirm_password = request.form.get('confirm_password')

        if password != confirm_password:
            flash('Пароли не совпадают', 'danger')
            return redirect(url_for('register'))

        conn = get_db_connection()
        try:
            password_hash = generate_password_hash(password)
            conn.execute(
                "INSERT INTO users (username, email, password_hash) VALUES (?, ?, ?)",
                (username, email, password_hash)
            )
            conn.commit()
            flash('Регистрация прошла успешно! Теперь вы можете войти.', 'success')
            return redirect(url_for('login'))
        except sqlite3.IntegrityError:
            flash('Пользователь с таким email или именем уже существует', 'danger')
        finally:
            conn.close()

    return render_template('auth/register.html')


@app.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('login'))


# Основные маршруты
@app.route("/")
@login_required
def index():
    conn = get_db_connection()
    try:
        chats = conn.execute(
            "SELECT chat_id, title FROM chats WHERE user_id = ? ORDER BY created_at DESC",
            (current_user.id,)
        ).fetchall()
        return render_template("index.html", chats=chats)
    finally:
        conn.close()


@app.route("/analyze", methods=["POST"])
@login_required
def analyze_text():
    if not text_classifier:
        return jsonify({"error": "Model not loaded"}), 500

    try:
        data = request.get_json()
        text = data.get("text", "").strip()

        if not text:
            return jsonify({"error": "Empty text"}), 400

        # Получаем предсказания модели
        result = text_classifier(text)

        # Проверяем структуру ответа
        if not result or not isinstance(result, list):
            return jsonify({"error": "Invalid model response"}), 500

        # Берем первый результат (самый вероятный)
        prediction = result[0] if result else {}

        # Проверяем наличие нужных полей
        if not all(key in prediction for key in ['label', 'score']):
            return jsonify({"error": "Invalid prediction format"}), 500

        # Сохраняем в базу данных
        conn = get_db_connection()
        conn.execute(
            "INSERT INTO analysis_reports (user_id, content, emotion, confidence) VALUES (?, ?, ?, ?)",
            (current_user.id, text, prediction['label'], prediction['score'])
        )
        conn.commit()
        conn.close()

        return jsonify({
            "emotion": emotion_map.get(prediction['label'], "❓ Неизвестно"),
            "confidence": float(prediction['score'])
        })

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


@app.route('/analyze_audio', methods=['POST'])
@login_required
def analyze_audio():
    if not audio_classifier or not speech_to_text_model:
        return jsonify({"error": "Model not loaded"}), 500

    if 'audio' not in request.files:
        return jsonify({'error': 'No audio file'}), 400

    try:
        audio_file = request.files['audio']
        temp_path = "temp_audio.wav"

        audio = AudioSegment.from_file(io.BytesIO(audio_file.read()))
        audio = audio.set_frame_rate(16000).set_channels(1)
        audio.export(temp_path, format="wav", codec="pcm_s16le")

        transcribed_text = transcribe_audio(temp_path)
        result = audio_classifier(temp_path)
        os.remove(temp_path)

        emotion_mapping = {
            'hap': 'happy',
            'sad': 'sad',
            'neu': 'neutral',
            'ang': 'angry'
        }
        emotions = {emotion_mapping.get(item['label'].lower(), 'neutral'): item['score']
                    for item in result if item['label'].lower() in emotion_mapping}

        dominant_emotion = max(emotions.items(), key=lambda x: x[1])
        response_map = {
            'happy': '😊 Радость',
            'sad': '😢 Грусть',
            'angry': '😠 Злость',
            'neutral': '😐 Нейтрально'
        }

        conn = get_db_connection()
        conn.execute(
            "INSERT INTO analysis_reports (user_id, content, emotion, confidence) VALUES (?, ?, ?, ?)",
            (current_user.id, transcribed_text, dominant_emotion[0], dominant_emotion[1])
        )
        conn.commit()
        conn.close()

        return jsonify({
            'emotion': response_map.get(dominant_emotion[0], 'неизвестно'),
            'confidence': round(dominant_emotion[1], 2),
            'transcribed_text': transcribed_text if transcribed_text else "Не удалось распознать текст"
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/get_chats')
@login_required
def get_chats():
    conn = get_db_connection()
    try:
        chats = conn.execute(
            "SELECT chat_id, title FROM chats WHERE user_id = ? ORDER BY created_at DESC",
            (current_user.id,)
        ).fetchall()
        return jsonify([dict(chat) for chat in chats])
    finally:
        conn.close()


@app.route('/start_chat', methods=['POST'])
@login_required
def start_chat():
    conn = get_db_connection()
    try:
        chat_id = str(uuid.uuid4())
        conn.execute(
            "INSERT INTO chats (chat_id, user_id, title, created_at) VALUES (?, ?, ?, ?)",
            (chat_id, current_user.id, f"Новый чат {datetime.now().strftime('%d.%m')}", datetime.now())
        )
        conn.commit()
        return jsonify({"chat_id": chat_id, "title": f"Новый чат {datetime.now().strftime('%d.%m')}"})
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    finally:
        conn.close()


@app.route('/load_chat/<chat_id>')
@login_required
def load_chat(chat_id):
    conn = get_db_connection()
    try:
        # Получаем информацию о чате
        chat = conn.execute(
            "SELECT chat_id, title FROM chats WHERE chat_id = ? AND user_id = ?",
            (chat_id, current_user.id)
        ).fetchone()

        if not chat:
            return jsonify({"error": "Чат не найден"}), 404

        # Получаем сообщения чата
        messages = conn.execute(
            "SELECT sender, content FROM messages WHERE chat_id = ? ORDER BY timestamp ASC",
            (chat_id,)
        ).fetchall()

        return jsonify({
            "chat_id": chat["chat_id"],
            "title": chat["title"],
            "messages": [dict(msg) for msg in messages]
        })
    finally:
        conn.close()



@app.route('/save_message', methods=['POST'])
@login_required
def save_message():
    data = request.get_json()
    if not data or 'chat_id' not in data or 'content' not in data or 'sender' not in data:
        return jsonify({"error": "Неверные данные"}), 400

    conn = get_db_connection()
    try:
        # Проверяем, что чат принадлежит текущему пользователю
        chat = conn.execute(
            "SELECT chat_id FROM chats WHERE chat_id = ? AND user_id = ?",
            (data['chat_id'], current_user.id)
        ).fetchone()

        if not chat:
            return jsonify({"error": "Чат не найден"}), 404

        # Анализируем эмоцию в тексте
        emotion = "neutral"
        confidence = 0.0
        if text_classifier and data['content'].strip():
            try:
                predictions = text_classifier(data['content'])[0]
                top_prediction = max(predictions, key=lambda x: x["score"])
                emotion = top_prediction["label"]
                confidence = top_prediction["score"]

                # Сохраняем анализ в базу
                conn.execute(
                    "INSERT INTO analysis_reports (user_id, content, emotion, confidence) VALUES (?, ?, ?, ?)",
                    (current_user.id, data['content'], emotion, confidence)
                )
            except Exception as e:
                print(f"Ошибка анализа эмоции: {e}")

        # Сохраняем сообщение
        conn.execute(
            "INSERT INTO messages (chat_id, sender, content, timestamp) VALUES (?, ?, ?, ?)",
            (data['chat_id'], data['sender'], data['content'], datetime.now())
        )
        conn.commit()

        return jsonify({
            "status": "success",
            "emotion": emotion_map.get(emotion, "❓ Неизвестно"),
            "confidence": round(confidence, 2)
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    finally:
        conn.close()

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