app/database.py

from flask_sqlalchemy import SQLAlchemy
from flask_login import UserMixin
from datetime import datetime
from werkzeug.security import generate_password_hash, check_password_hash

db = SQLAlchemy()

# 사용자 모델
class User(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    nickname = db.Column(db.String(80), nullable=True)  # 닉네임 필드 추가
    password_hash = db.Column(db.String(255), nullable=False)
    is_admin = db.Column(db.Boolean, default=False, nullable=False)
    is_active = db.Column(db.Boolean, default=True, nullable=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow, nullable=False)
    last_login = db.Column(db.DateTime, nullable=True)
    
    def set_password(self, password):
        self.password_hash = generate_password_hash(password)
    
    def check_password(self, password):
        return check_password_hash(self.password_hash, password)
    
    def to_dict(self):
        return {
            'id': self.id,
            'username': self.username,
            'nickname': self.nickname,
            'is_admin': self.is_admin,
            'is_active': self.is_active,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'last_login': self.last_login.isoformat() if self.last_login else None
        }

# 업로드된 파일 정보 모델
class UploadedFile(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    filename = db.Column(db.String(255), nullable=False)
    original_filename = db.Column(db.String(255), nullable=False)
    file_path = db.Column(db.String(500), nullable=False)
    file_size = db.Column(db.Integer, nullable=False)
    model_name = db.Column(db.String(100), nullable=True)  # 연결된 모델 이름
    uploaded_at = db.Column(db.DateTime, default=datetime.utcnow, nullable=False)
    uploaded_by = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=True)
    parent_file_id = db.Column(db.Integer, db.ForeignKey('uploaded_file.id'), nullable=True)  # 이어서 업로드한 경우 원본 파일 ID
    
    # 관계
    parent_file = db.relationship('UploadedFile', remote_side=[id], backref='child_files')
    
    def to_dict(self):
        # 청크 개수 계산
        chunk_count = len(self.chunks) if hasattr(self, 'chunks') else 0
        
        return {
            'id': self.id,
            'filename': self.filename,
            'original_filename': self.original_filename,
            'file_size': self.file_size,
            'model_name': self.model_name,
            'uploaded_at': self.uploaded_at.isoformat() if self.uploaded_at else None,
            'uploaded_by': self.uploaded_by,
            'parent_file_id': self.parent_file_id,
            'chunk_count': chunk_count,
            'child_count': len(self.child_files) if self.child_files else 0
        }

# 대화 세션 모델
class ChatSession(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=False)
    title = db.Column(db.String(255), nullable=True)
    model_name = db.Column(db.String(100), nullable=True)
    created_at = db.Column(db.DateTime, default=datetime.utcnow, nullable=False)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)
    
    # 관계
    user = db.relationship('User', backref='chat_sessions')
    messages = db.relationship('ChatMessage', backref='session', lazy=True, cascade='all, delete-orphan', order_by='ChatMessage.created_at')
    
    def to_dict(self):
        return {
            'id': self.id,
            'user_id': self.user_id,
            'title': self.title,
            'model_name': self.model_name,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None,
            'message_count': len(self.messages)
        }

# 대화 메시지 모델
class ChatMessage(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    session_id = db.Column(db.Integer, db.ForeignKey('chat_session.id'), nullable=False)
    role = db.Column(db.String(20), nullable=False)  # 'user' or 'ai'
    content = db.Column(db.Text, nullable=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow, nullable=False)
    
    def to_dict(self):
        return {
            'id': self.id,
            'session_id': self.session_id,
            'role': self.role,
            'content': self.content,
            'created_at': self.created_at.isoformat() if self.created_at else None
        }

# 문서 청크 모델 (RAG용)
class DocumentChunk(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    file_id = db.Column(db.Integer, db.ForeignKey('uploaded_file.id'), nullable=False)
    chunk_index = db.Column(db.Integer, nullable=False)  # 청크 순서
    content = db.Column(db.Text, nullable=False)  # 청크 내용
    embedding = db.Column(db.Text, nullable=True)  # 임베딩 벡터 (JSON 문자열로 저장)
    created_at = db.Column(db.DateTime, default=datetime.utcnow, nullable=False)
    
    # 관계
    file = db.relationship('UploadedFile', backref='chunks')
    
    def to_dict(self):
        return {
            'id': self.id,
            'file_id': self.file_id,
            'chunk_index': self.chunk_index,
            'content': self.content,
            'created_at': self.created_at.isoformat() if self.created_at else None
        }

# Parent Chunk 모델 (AI 분석 결과 저장)
class ParentChunk(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    file_id = db.Column(db.Integer, db.ForeignKey('uploaded_file.id'), nullable=False, unique=True)
    world_view = db.Column(db.Text, nullable=True)  # 세계관 설명
    characters = db.Column(db.Text, nullable=True)  # 주요 캐릭터 분석
    story = db.Column(db.Text, nullable=True)  # 주요 스토리 분석
    episodes = db.Column(db.Text, nullable=True)  # 주요 에피소드 분석
    others = db.Column(db.Text, nullable=True)  # 기타
    created_at = db.Column(db.DateTime, default=datetime.utcnow, nullable=False)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)
    
    # 관계
    file = db.relationship('UploadedFile', backref='parent_chunk')
    
    def to_dict(self):
        return {
            'id': self.id,
            'file_id': self.file_id,
            'world_view': self.world_view,
            'characters': self.characters,
            'story': self.story,
            'episodes': self.episodes,
            'others': self.others,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None
        }

# 시스템 설정 모델 (API 키 등)
class SystemConfig(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    key = db.Column(db.String(100), unique=True, nullable=False)  # 설정 키 (예: 'gemini_api_key')
    value = db.Column(db.Text, nullable=True)  # 설정 값 (암호화된 API 키)
    description = db.Column(db.String(255), nullable=True)  # 설정 설명
    created_at = db.Column(db.DateTime, default=datetime.utcnow, nullable=False)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)
    
    def to_dict(self):
        return {
            'id': self.id,
            'key': self.key,
            'value': self.value,  # 실제 값 반환 (보안을 위해 마스킹 필요할 수 있음)
            'description': self.description,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None
        }
    
    @staticmethod
    def get_config(key, default=None):
        """설정 값 가져오기"""
        try:
            config = SystemConfig.query.filter_by(key=key).first()
            return config.value if config else default
        except Exception as e:
            # 테이블이 없거나 오류 발생 시 기본값 반환
            print(f"[SystemConfig.get_config] 오류: {e}")
            return default
    
    @staticmethod
    def set_config(key, value, description=None):
        """설정 값 저장/업데이트"""
        try:
            # 테이블이 없으면 생성 시도
            from sqlalchemy import inspect
            inspector = inspect(db.engine)
            if 'system_config' not in inspector.get_table_names():
                print("[SystemConfig.set_config] SystemConfig 테이블 생성 중...")
                db.create_all()
            
            config = SystemConfig.query.filter_by(key=key).first()
            if config:
                config.value = value
                if description:
                    config.description = description
                config.updated_at = datetime.utcnow()
            else:
                config = SystemConfig(key=key, value=value, description=description)
                db.session.add(config)
            db.session.commit()
            return config
        except Exception as e:
            db.session.rollback()
            print(f"[SystemConfig.set_config] 오류: {e}")
            import traceback
            traceback.print_exc()
            raise