#!/usr/bin/env python3
"""
간단한 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"
}

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_training_simple, request)
    
    return {
        "message": "Training started",
        "status": "starting",
        "model_name": request.model_name
    }

async def run_training_simple(request: TrainingRequest):
    """간단한 학습 실행 (시뮬레이션)"""
    global training_status
    
    try:
        logger.info("Starting simple training process...")
        training_status["status"] = "running"
        
        # 데이터 파일 확인
        train_file = Path("/app/train.csv")
        val_file = Path("/app/validation.csv")
        
        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
        
        logger.info("Data files found, starting training simulation...")
        
        # 간단한 훈련 시뮬레이션
        for epoch in range(1, 4):
            training_status["current_epoch"] = epoch
            training_status["progress"] = (epoch / 3) * 100
            training_status["loss"] = 2.5 - (epoch * 0.5)  # 시뮬레이션 손실값
            
            logger.info(f"Epoch {epoch}/3 - Loss: {training_status['loss']:.3f}")
            await asyncio.sleep(5)  # 5초 대기 (시뮬레이션)
        
        training_status.update({
            "is_running": False,
            "progress": 100,
            "status": "completed"
        })
        logger.info("Training completed successfully!")
        
    except Exception as e:
        logger.error(f"Training error: {str(e)}")
        training_status.update({
            "is_running": False,
            "status": "error",
            "error": str(e)
        })

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

@app.get("/logs")
async def get_logs():
    """로그 조회"""
    log_file = Path("/app/training.log")
    if log_file.exists():
        with open(log_file, "r", encoding="utf-8") as f:
            logs = f.read()
        return {"logs": logs}
    else:
        return {"logs": "No logs available"}

@app.get("/logs/stream")
async def stream_logs():
    """실시간 로그 스트리밍"""
    def generate_logs():
        log_file = Path("/app/training.log")
        if log_file.exists():
            with open(log_file, "r", encoding="utf-8") as f:
                for line in f:
                    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")
    
    info = {
        "train_file_exists": train_file.exists(),
        "validation_file_exists": val_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
    }
    
    return info

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