import os
import requests
from datetime import datetime
from banyarag.file_uploader import FileUploader

class banyainstance:
    def __init__(self, server_url="https://train.banya.ai"):
        """
        BanyaInstance 초기화.
        
        :param server_url: 파일 업로드 및 학습 서버 URL
        """
        self.server_url = server_url
        self.drive_path = None

    def setDrivePath(self, drive_path):
        """
        드라이브 경로를 설정합니다.
        
        :param drive_path: 로컬 드라이브 경로
        """
        try:
            if not os.path.isdir(drive_path):
                raise NotADirectoryError(f"Invalid drive path: {drive_path}")
            self.drive_path = drive_path
            print(f"Drive path set to: {self.drive_path}")
            return True
        except Exception as e:
            print(f"Failed to set drive path: {e}")
            return False

    def uploadAllFiles(self):
        """
        설정된 드라이브 경로 내 모든 파일을 서버로 업로드합니다.
        
        :return: 업로드 타임스탬프
        """
        try:
            if not self.drive_path:
                raise ValueError("Drive path is not set. Use setDrivePath() first.")

            uploader = FileUploader(self.server_url)
            responses = uploader.upload_folder(self.drive_path)
            
            # 업로드 성공 여부 출력
            for file_path, status_code, response_text in responses:
                if status_code == 200:
                    print(f"Uploaded: {file_path} - Success")
                else:
                    print(f"Uploaded: {file_path} - Failed ({status_code})")

            # 현재 타임스탬프 반환
            timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
            print(f"Upload completed at: {timestamp}")
            return timestamp
        except Exception as e:
            print(f"Failed to upload files: {e}")
            return False

    def trainRAG(self, timestamp, embedder, llm):
        """
        업로드된 데이터로 학습을 수행합니다.
        
        :param timestamp: 업로드 타임스탬프
        :param embedder: 임베더 모델 이름
        :param llm: LLM 모델 이름
        :return: 학습 완료 후 Chat URL
        """
        try:
            train_endpoint = f"{self.server_url}/llm/lib/trainer"
            payload = {
                "timestamp": timestamp,
                "embedder": embedder,
                "llm": llm,
            }
            
            response = requests.post(train_endpoint, json=payload)
            
            if response.status_code == 200:
                chat_url = response.json().get("chat_url", "")
                print(f"Training completed. Chat URL: {chat_url}")
                return chat_url
            else:
                raise RuntimeError(f"Training failed: {response.status_code} - {response.text}")
        except Exception as e:
            print(f"Failed to train RAG: {e}")
            return False