document_processor.py

import os
import re
from typing import List, Dict
from pathlib import Path
import fitz  # PyMuPDF
import docx
import pandas as pd
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document

class DocumentProcessor:
    """복무관리 문서 처리 클래스"""

    def __init__(self, chunk_size: int = 500, chunk_overlap: int = 50):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=["\n\n", "\n", " ", ""]
        )

    def load_documents_from_folder(self, folder_path: str) -> List[Document]:
        """폴더에서 모든 문서 로드"""
        documents = []
        folder = Path(folder_path)

        if not folder.exists():
            print(f"⚠️ 폴더가 존재하지 않습니다: {folder_path}")
            return documents

        # 지원하는 파일 형식
        supported_extensions = ['.pdf', '.docx', '.txt', '.xlsx', '.csv']

        for file_path in folder.rglob('*'):
            if file_path.suffix.lower() in supported_extensions:
                try:
                    print(f"📄 문서 로드: {file_path.name}")
                    docs = self.load_single_document(str(file_path))
                    documents.extend(docs)
                except Exception as e:
                    print(f"❌ 문서 로드 실패 ({file_path.name}): {str(e)}")

        return documents

    def load_single_document(self, file_path: str) -> List[Document]:
        """단일 문서 로드"""
        file_ext = Path(file_path).suffix.lower()

        if file_ext == '.pdf':
            return self._load_pdf(file_path)
        elif file_ext == '.docx':
            return self._load_docx(file_path)
        elif file_ext == '.txt':
            return self._load_txt(file_path)
        elif file_ext in ['.xlsx', '.csv']:
            return self._load_table(file_path)
        else:
            raise ValueError(f"지원하지 않는 파일 형식: {file_ext}")

    def _load_pdf(self, file_path: str) -> List[Document]:
        """PDF 파일 로드"""
        documents = []

        try:
            with fitz.open(file_path) as doc:
                full_text = ""

                for page_num in range(len(doc)):
                    page = doc[page_num]
                    page_text = page.get_text()

                    # 페이지 정제
                    page_text = self._clean_text(page_text)

                    if page_text.strip():
                        full_text += f"\n\n--- 페이지 {page_num + 1} ---\n\n{page_text}"

                if full_text.strip():
                    chunks = self.text_splitter.split_text(full_text)

                    for i, chunk in enumerate(chunks):
                        metadata = {
                            "source": Path(file_path).name,
                            "page": "multiple",
                            "chunk_id": i,
                            "file_type": "pdf"
                        }

                        documents.append(Document(page_content=chunk, metadata=metadata))

        except Exception as e:
            print(f"PDF 로드 중 오류: {str(e)}")
            raise

        return documents

    def _load_docx(self, file_path: str) -> List[Document]:
        """Word 문서 로드"""
        documents = []

        try:
            doc = docx.Document(file_path)
            paragraphs = []

            for para in doc.paragraphs:
                if para.text.strip():
                    paragraphs.append(para.text)

            full_text = "\n\n".join(paragraphs)

            if full_text.strip():
                chunks = self.text_splitter.split_text(full_text)

                for i, chunk in enumerate(chunks):
                    metadata = {
                        "source": Path(file_path).name,
                        "chunk_id": i,
                        "file_type": "docx"
                    }

                    documents.append(Document(page_content=chunk, metadata=metadata))

        except Exception as e:
            print(f"DOCX 로드 중 오류: {str(e)}")
            raise

        return documents

    def _load_txt(self, file_path: str) -> List[Document]:
        """텍스트 파일 로드"""
        documents = []

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                text = f.read()

            text = self._clean_text(text)

            if text.strip():
                chunks = self.text_splitter.split_text(text)

                for i, chunk in enumerate(chunks):
                    metadata = {
                        "source": Path(file_path).name,
                        "chunk_id": i,
                        "file_type": "txt"
                    }

                    documents.append(Document(page_content=chunk, metadata=metadata))

        except Exception as e:
            print(f"TXT 로드 중 오류: {str(e)}")
            raise

        return documents

    def _load_table(self, file_path: str) -> List[Document]:
        """엑셀/CSV 파일 로드"""
        documents = []

        try:
            if file_path.endswith('.xlsx'):
                df = pd.read_excel(file_path)
            else:
                df = pd.read_csv(file_path, encoding='utf-8')

            # 데이터프레임을 텍스트로 변환
            text_parts = []
            text_parts.append(f"파일: {Path(file_path).name}")
            text_parts.append(f"컬럼: {', '.join(df.columns.tolist())}")

            for index, row in df.iterrows():
                row_text = " | ".join([f"{col}: {val}" for col, val in row.items() if pd.notna(val)])
                text_parts.append(f"행 {index + 1}: {row_text}")

            full_text = "\n\n".join(text_parts)

            if full_text.strip():
                chunks = self.text_splitter.split_text(full_text)

                for i, chunk in enumerate(chunks):
                    metadata = {
                        "source": Path(file_path).name,
                        "chunk_id": i,
                        "file_type": "table",
                        "total_rows": len(df)
                    }

                    documents.append(Document(page_content=chunk, metadata=metadata))

        except Exception as e:
            print(f"테이블 로드 중 오류: {str(e)}")
            raise

        return documents

    def _clean_text(self, text: str) -> str:
        """텍스트 정제"""
        # 불필요한 공백 제거
        text = re.sub(r'\s+', ' ', text)

        # 특수문자 정리
        text = re.sub(r'[^\w\s\.\,\?\!\:\;\-\(\)\/\&\@]', ' ', text)

        # 연속된 공백 제거
        text = re.sub(r'\s+', ' ', text).strip()

        return text

    def process_documents(self, documents: List[Document]) -> List[Document]:
        """문서 후처리"""
        processed_docs = []

        for doc in documents:
            content = doc.page_content.strip()

            if content and len(content) > 20:  # 너무 짧은 청크는 제외
                # 복무관리 특화 키워드 강화
                content = self._enhance_fire_service_terms(content)

                processed_doc = Document(
                    page_content=content,
                    metadata=doc.metadata
                )
                processed_docs.append(processed_doc)

        return processed_docs

    def _enhance_fire_service_terms(self, text: str) -> str:
        """소방 용어 강화"""
        # 복무관리 관련 키워드 매핑
        term_mappings = {
            "연차": "연차휴가",
            "연장": "연장근무",
            "당직": "당직근무",
            "파견": "파견근무",
            "인사": "인사평가",
            "승진": "승진시험",
            "교육": "교육훈련",
            "휴가": "휴가사용",
            "상벌": "상벌규정",
            "징계": "징계절차"
        }

        enhanced_text = text
        for standard_term, enhanced_term in term_mappings.items():
            enhanced_text = enhanced_text.replace(standard_term, enhanced_term)

        return enhanced_text

# 테스트용 함수
def test_document_processor():
    """문서 처리기 테스트"""
    processor = DocumentProcessor()

    # 샘플 documents 폴더 생성
    docs_folder = "documents"
    os.makedirs(docs_folder, exist_ok=True)

    # 샘플 문서 생성
    sample_text = """
    복무관리 규정

    제1장 총칙
    제1조 (목적)
    이 규정은 소방공무원의 복무에 관한 사항을 규정하여 직무수행의 효율성을 높이고
    조직의 발전에 기여함을 목적으로 한다.

    제2조 (근무시간)
    1. 정규근무시간은 09:00부터 18:00까지로 한다.
    2. 점심시간은 12:00부터 13:00까지로 한다.
    3. 당직근무는 정규근무시간 외에 수행하는 근무를 말한다.

    제3조 (연차휴가)
    1. 연차휴가는 1년간 정상 근무한 자에게 15일을 부여한다.
    2. 연차휴가 사용 시 3일 전까지 신청서를 제출해야 한다.
    3. 부서장의 승인을 받아 사용한다.
    """

    with open(os.path.join(docs_folder, "sample_policy.txt"), "w", encoding="utf-8") as f:
        f.write(sample_text)

    # 문서 로드 테스트
    documents = processor.load_documents_from_folder(docs_folder)
    print(f"✅ {len(documents)}개 문서 청크 생성 완료")

    return documents

if __name__ == "__main__":
    test_document_processor()