#!/usr/bin/env python3
"""
올바른 AutoTrain 명령어를 사용한 EXAONE Fine-tuning Space FastAPI 애플리케이션
"""

import os
import json
import subprocess
import asyncio
from pathlib import Path
from typing import Dict, Any
import logging

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import uvicorn

# 로깅 설정
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="EXAONE Fine-tuning",
    description="EXAONE 4.0 1.2B 모델 파인튜닝 API",
    version="1.0.0"
)

# 전역 변수
training_status = {
    "is_running": False,
    "progress": 0,
    "current_epoch": 0,
    "total_epochs": 3,
    "loss": 0.0,
    "status": "idle",
    "log_file": "/tmp/training.log"
}

class TrainingRequest(BaseModel):
    model_name: str = "amis5895/exaone-1p2b-nutrition-kdri"

@app.get("/")
async def root():
    """루트 엔드포인트"""
    return {
        "message": "EXAONE Fine-tuning API",
        "status": "running",
        "version": "1.0.0"
    }

@app.post("/start_training")
async def start_training(request: TrainingRequest, background_tasks: BackgroundTasks):
    """학습 시작"""
    global training_status
    
    if training_status["is_running"]:
        raise HTTPException(status_code=400, detail="Training is already running")
    
    training_status.update({
        "is_running": True,
        "progress": 0,
        "current_epoch": 0,
        "status": "starting"
    })
    
    # 백그라운드에서 학습 시작
    background_tasks.add_task(run_corrected_training, request)
    
    return {
        "message": "Training started",
        "status": "starting",
        "model_name": request.model_name
    }

async def run_corrected_training(request: TrainingRequest):
    """수정된 AutoTrain을 사용한 학습 실행"""
    global training_status
    
    try:
        logger.info("Starting corrected AutoTrain training process...")
        training_status["status"] = "running"
        
        # 데이터 파일 확인
        train_file = Path("/app/train.csv")
        val_file = Path("/app/validation.csv")
        config_file = Path("/app/autotrain_ultra_low_final.yaml")
        
        if not train_file.exists():
            logger.error(f"Training file not found: {train_file}")
            training_status.update({
                "is_running": False,
                "status": "failed",
                "error": "Training file not found"
            })
            return
        
        if not val_file.exists():
            logger.error(f"Validation file not found: {val_file}")
            training_status.update({
                "is_running": False,
                "status": "failed",
                "error": "Validation file not found"
            })
            return
        
        if not config_file.exists():
            logger.error(f"Config file not found: {config_file}")
            training_status.update({
                "is_running": False,
                "status": "failed",
                "error": "Config file not found"
            })
            return
        
        logger.info("All files found, starting corrected AutoTrain training...")
        
        # 로그 파일 초기화
        log_file = Path(training_status["log_file"])
        try:
            log_file.write_text("Starting corrected AutoTrain training...\n", encoding="utf-8")
        except Exception as e:
            logger.warning(f"Could not write to log file: {e}")
            training_status["log_content"] = "Starting corrected AutoTrain training...\n"
        
        # 환경변수 설정
        env = os.environ.copy()
        env["TRANSFORMERS_CACHE"] = "/tmp/huggingface_cache"
        env["HF_HOME"] = "/tmp/huggingface"
        env["OMP_NUM_THREADS"] = "1"
        
        # 수정된 AutoTrain 명령어 (올바른 형식 사용)
        cmd = [
            "autotrain", "llm",
            "--train",
            "--project_name", "exaone-finetuning",
            "--model", "LGAI-EXAONE/EXAONE-4.0-1.2B",
            "--data_path", "/app",
            "--text_column", "text",
            "--use-peft",  # --use_peft 대신 --use-peft
            "--quantization", "int4",
            "--lora-r", "16",  # --lora_r 대신 --lora-r
            "--lora-alpha", "32",  # --lora_alpha 대신 --lora-alpha
            "--lora-dropout", "0.05",  # --lora_dropout 대신 --lora-dropout
            "--target-modules", "all-linear",  # --target_modules 대신 --target-modules
            "--epochs", "3",
            "--batch-size", "4",  # --batch_size 대신 --batch-size
            "--gradient-accumulation", "4",  # --gradient_accumulation 대신 --gradient-accumulation
            "--learning-rate", "2e-4",  # --learning_rate 대신 --learning-rate
            "--warmup-ratio", "0.03",  # --warmup_ratio 대신 --warmup-ratio
            "--mixed-precision", "fp16",  # --mixed_precision 대신 --mixed-precision
            "--push-to-hub",  # --push_to_hub 대신 --push-to-hub
            "--hub-model-id", request.model_name,  # --hub_model_id 대신 --hub-model-id
            "--username", "amis5895"
        ]
        
        logger.info(f"Running corrected command: {' '.join(cmd)}")
        
        # 로그 파일에 명령어 기록
        try:
            with open(log_file, "a", encoding="utf-8") as f:
                f.write(f"Corrected Command: {' '.join(cmd)}\n")
                f.write("=" * 50 + "\n")
        except:
            if "log_content" not in training_status:
                training_status["log_content"] = ""
            training_status["log_content"] += f"Corrected Command: {' '.join(cmd)}\n" + "=" * 50 + "\n"
        
        # AutoTrain 프로세스 실행
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
            universal_newlines=True,
            cwd="/app",
            env=env
        )
        
        # 학습 진행 상황 모니터링
        for line in process.stdout:
            logger.info(line.strip())
            
            # 로그 파일에 기록
            try:
                with open(log_file, "a", encoding="utf-8") as f:
                    f.write(line)
            except:
                if "log_content" not in training_status:
                    training_status["log_content"] = ""
                training_status["log_content"] += line
            
            # 진행률 파싱
            if "epoch" in line.lower() and "/" in line:
                try:
                    # "Epoch 1/3" 형태에서 진행률 추출
                    parts = line.split()
                    for i, part in enumerate(parts):
                        if part.lower() == "epoch" and i + 1 < len(parts):
                            epoch_info = parts[i + 1]
                            if "/" in epoch_info:
                                current, total = epoch_info.split("/")
                                training_status["current_epoch"] = int(current)
                                training_status["total_epochs"] = int(total)
                                training_status["progress"] = (int(current) / int(total)) * 100
                                break
                except:
                    pass
            
            # 손실값 파싱
            if "loss" in line.lower():
                try:
                    parts = line.split()
                    for i, part in enumerate(parts):
                        if part.lower() == "loss" and i + 1 < len(parts):
                            loss_value = float(parts[i + 1])
                            training_status["loss"] = loss_value
                            break
                except:
                    pass
        
        process.wait()
        
        if process.returncode == 0:
            training_status.update({
                "is_running": False,
                "progress": 100,
                "status": "completed"
            })
            logger.info("Training completed successfully!")
            
            # 완료 로그 기록
            try:
                with open(log_file, "a", encoding="utf-8") as f:
                    f.write("\n" + "=" * 50 + "\n")
                    f.write("Training completed successfully!\n")
            except:
                if "log_content" not in training_status:
                    training_status["log_content"] = ""
                training_status["log_content"] += "\n" + "=" * 50 + "\nTraining completed successfully!\n"
        else:
            training_status.update({
                "is_running": False,
                "status": "failed"
            })
            logger.error("Training failed!")
            
            # 실패 로그 기록
            try:
                with open(log_file, "a", encoding="utf-8") as f:
                    f.write("\n" + "=" * 50 + "\n")
                    f.write(f"Training failed with return code: {process.returncode}\n")
            except:
                if "log_content" not in training_status:
                    training_status["log_content"] = ""
                training_status["log_content"] += "\n" + "=" * 50 + f"\nTraining failed with return code: {process.returncode}\n"
            
    except Exception as e:
        logger.error(f"Training error: {str(e)}")
        training_status.update({
            "is_running": False,
            "status": "error",
            "error": str(e)
        })
        
        # 오류 로그 기록
        try:
            with open(log_file, "a", encoding="utf-8") as f:
                f.write(f"\nError: {str(e)}\n")
        except:
            if "log_content" not in training_status:
                training_status["log_content"] = ""
            training_status["log_content"] += f"\nError: {str(e)}\n"

@app.get("/status")
async def get_status():
    """학습 상태 조회"""
    return training_status

@app.get("/logs")
async def get_logs():
    """로그 조회"""
    log_file = Path(training_status["log_file"])
    if log_file.exists():
        try:
            with open(log_file, "r", encoding="utf-8") as f:
                logs = f.read()
            return {"logs": logs}
        except:
            pass
    
    # 파일을 읽을 수 없으면 메모리에서 가져오기
    if "log_content" in training_status:
        return {"logs": training_status["log_content"]}
    else:
        return {"logs": "No logs available"}

@app.get("/logs/stream")
async def stream_logs():
    """실시간 로그 스트리밍"""
    def generate_logs():
        log_file = Path(training_status["log_file"])
        if log_file.exists():
            try:
                with open(log_file, "r", encoding="utf-8") as f:
                    for line in f:
                        yield f"data: {line}\\n\\n"
            except:
                pass
        
        # 파일을 읽을 수 없으면 메모리에서 가져오기
        if "log_content" in training_status:
            for line in training_status["log_content"].split('\n'):
                yield f"data: {line}\\n\\n"
        else:
            yield "data: No logs available\\n\\n"
    
    return StreamingResponse(generate_logs(), media_type="text/plain")

@app.post("/stop_training")
async def stop_training():
    """학습 중지"""
    global training_status
    
    if not training_status["is_running"]:
        raise HTTPException(status_code=400, detail="No training is running")
    
    training_status.update({
        "is_running": False,
        "status": "stopped"
    })
    
    return {"message": "Training stopped"}

@app.get("/health")
async def health_check():
    """헬스 체크"""
    return {"status": "healthy", "timestamp": "2024-01-01T00:00:00Z"}

@app.get("/data_info")
async def get_data_info():
    """데이터 정보 조회"""
    train_file = Path("/app/train.csv")
    val_file = Path("/app/validation.csv")
    config_file = Path("/app/autotrain_ultra_low_final.yaml")
    
    info = {
        "train_file_exists": train_file.exists(),
        "validation_file_exists": val_file.exists(),
        "config_file_exists": config_file.exists(),
        "train_file_size": train_file.stat().st_size if train_file.exists() else 0,
        "validation_file_size": val_file.stat().st_size if val_file.exists() else 0,
        "config_file_size": config_file.stat().st_size if config_file.exists() else 0
    }
    
    return info

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)