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)  # 연결된 모델 이름
    is_public = db.Column(db.Boolean, default=False, nullable=False)  # 공개 여부 (기본값: 미공개)
    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,
            'is_public': self.is_public,
            '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)  # 하위 호환성을 위해 유지
    analysis_model = db.Column(db.String(100), nullable=True)  # 질문 분석용 모델
    answer_model = 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,  # 하위 호환성
            'analysis_model': self.analysis_model,
            'answer_model': self.answer_model,
            '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=True)  # 시스템 사용은 NULL 허용
    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)
    # 토큰 사용량 정보 (AI 응답 메시지에만 저장)
    input_tokens = db.Column(db.Integer, nullable=True)  # 입력 토큰 수
    output_tokens = db.Column(db.Integer, nullable=True)  # 출력 토큰 수
    model_name = db.Column(db.String(100), nullable=True)  # 사용된 AI 모델명
    usage_type = db.Column(db.String(20), nullable=True, default='user')  # 'user' or 'system' (시스템 사용 구분)
    
    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,
            'input_tokens': self.input_tokens,
            'output_tokens': self.output_tokens,
            'model_name': self.model_name,
            'usage_type': self.usage_type
        }

# 문서 청크 모델 (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 문자열로 저장)
    chunk_metadata = db.Column(db.Text, nullable=True)  # 메타데이터 (JSON 문자열로 저장: chapter, pov, characters, time_background)
    created_at = db.Column(db.DateTime, default=datetime.utcnow, nullable=False)
    
    # 관계
    file = db.relationship('UploadedFile', backref='chunks')
    
    def to_dict(self):
        import json
        metadata_dict = None
        if self.chunk_metadata:
            try:
                metadata_dict = json.loads(self.chunk_metadata)
            except:
                metadata_dict = None
        
        return {
            'id': self.id,
            'file_id': self.file_id,
            'chunk_index': self.chunk_index,
            'content': self.content,
            'metadata': metadata_dict,
            '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:
                print(f"[SystemConfig.set_config] 기존 설정 업데이트: {key}")
                config.value = value
                if description:
                    config.description = description
                config.updated_at = datetime.utcnow()
            else:
                print(f"[SystemConfig.set_config] 새 설정 생성: {key}")
                config = SystemConfig(key=key, value=value, description=description)
                db.session.add(config)
            
            db.session.commit()
            
            # 저장 확인
            verify_config = SystemConfig.query.filter_by(key=key).first()
            if verify_config and verify_config.value == value:
                print(f"[SystemConfig.set_config] 설정 저장 확인됨: {key} (길이: {len(str(value))}자)")
            else:
                print(f"[SystemConfig.set_config] 경고: 설정 저장 후 확인 실패: {key}")
            
            return config
        except Exception as e:
            db.session.rollback()
            print(f"[SystemConfig.set_config] 오류: {e}")
            import traceback
            traceback.print_exc()
            raise

# 회차별 분석 모델
class EpisodeAnalysis(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    file_id = db.Column(db.Integer, db.ForeignKey('uploaded_file.id'), nullable=False)
    episode_title = db.Column(db.String(100), nullable=False)  # 회차 제목 (예: '1화', '2화')
    analysis_content = db.Column(db.Text, nullable=False)  # 분석 결과 (하나의 텍스트로 이어서 저장)
    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='episode_analyses')
    
    def to_dict(self):
        return {
            'id': self.id,
            'file_id': self.file_id,
            'episode_title': self.episode_title,
            'analysis_content': self.analysis_content,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None
        }

# Graph Extraction 모델 (GraphRAG용)
class GraphEntity(db.Model):
    """Graph Extraction에서 추출된 엔티티 (인물/장소)"""
    id = db.Column(db.Integer, primary_key=True)
    file_id = db.Column(db.Integer, db.ForeignKey('uploaded_file.id'), nullable=False)
    episode_title = db.Column(db.String(100), nullable=False)  # 회차 제목
    entity_name = db.Column(db.String(200), nullable=False)  # 엔티티 이름
    entity_type = db.Column(db.String(50), nullable=False)  # 'character' 또는 'location'
    description = db.Column(db.Text, nullable=True)  # 엔티티 설명
    role = db.Column(db.String(200), nullable=True)  # 인물의 경우 역할
    category = db.Column(db.String(200), nullable=True)  # 장소의 경우 유형
    created_at = db.Column(db.DateTime, default=datetime.utcnow, nullable=False)
    
    # 관계
    file = db.relationship('UploadedFile', backref='graph_entities')
    
    def to_dict(self):
        return {
            'id': self.id,
            'file_id': self.file_id,
            'episode_title': self.episode_title,
            'entity_name': self.entity_name,
            'entity_type': self.entity_type,
            'description': self.description,
            'role': self.role,
            'category': self.category,
            'created_at': self.created_at.isoformat() if self.created_at else None
        }

class GraphRelationship(db.Model):
    """Graph Extraction에서 추출된 관계"""
    id = db.Column(db.Integer, primary_key=True)
    file_id = db.Column(db.Integer, db.ForeignKey('uploaded_file.id'), nullable=False)
    episode_title = db.Column(db.String(100), nullable=False)  # 회차 제목
    source = db.Column(db.String(200), nullable=False)  # 관계의 주체
    target = db.Column(db.String(200), nullable=False)  # 관계의 대상
    relationship_type = db.Column(db.String(200), nullable=False)  # 관계 유형
    description = db.Column(db.Text, nullable=True)  # 관계 설명
    event = db.Column(db.String(500), nullable=True)  # 관계를 형성/변화시킨 사건
    created_at = db.Column(db.DateTime, default=datetime.utcnow, nullable=False)
    
    # 관계
    file = db.relationship('UploadedFile', backref='graph_relationships')
    
    def to_dict(self):
        return {
            'id': self.id,
            'file_id': self.file_id,
            'episode_title': self.episode_title,
            'source': self.source,
            'target': self.target,
            'relationship_type': self.relationship_type,
            'description': self.description,
            'event': self.event,
            'created_at': self.created_at.isoformat() if self.created_at else None
        }

class GraphEvent(db.Model):
    """Graph Extraction에서 추출된 사건"""
    id = db.Column(db.Integer, primary_key=True)
    file_id = db.Column(db.Integer, db.ForeignKey('uploaded_file.id'), nullable=False)
    episode_title = db.Column(db.String(100), nullable=False)  # 회차 제목
    event_name = db.Column(db.String(500), nullable=False)  # 사건 이름
    description = db.Column(db.Text, nullable=False)  # 사건 설명
    participants = db.Column(db.Text, nullable=True)  # 관련 인물들 (JSON 문자열로 저장)
    location = db.Column(db.String(500), nullable=True)  # 사건 발생 장소
    significance = db.Column(db.String(200), nullable=True)  # 사건의 중요도
    created_at = db.Column(db.DateTime, default=datetime.utcnow, nullable=False)
    
    # 관계
    file = db.relationship('UploadedFile', backref='graph_events')
    
    def to_dict(self):
        import json
        participants_list = []
        if self.participants:
            try:
                participants_list = json.loads(self.participants)
            except:
                participants_list = []
        
        return {
            'id': self.id,
            'file_id': self.file_id,
            'episode_title': self.episode_title,
            'event_name': self.event_name,
            'description': self.description,
            'participants': participants_list,
            'location': self.location,
            'significance': self.significance,
            'created_at': self.created_at.isoformat() if self.created_at else None
        }