import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TextClassificationPipeline
import torch

print("Gradio app script started.") # 스크립트 시작 로그

MODEL_PATH = "./" # 모델 파일이 컨테이너 내 app.py와 동일한 디렉토리에 있다고 가정

pipeline_instance = None # 파이프라인 인스턴스를 전역적으로 선언

try:
    print(f"Loading tokenizer from: {MODEL_PATH}")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    print(f"Loading model from: {MODEL_PATH}")
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
    device_to_use = 0 if torch.cuda.is_available() else -1 # GPU 우선 사용, 없으면 CPU
    print(f"Using device: {'cuda:0' if device_to_use == 0 else 'cpu'}")
    pipeline_instance = TextClassificationPipeline(model=model, tokenizer=tokenizer, device=device_to_use)
    print("Pipeline created successfully.")
except Exception as e:
    print(f"Error loading model, tokenizer, or creating pipeline: {e}")
    # 파이프라인 생성 실패 시 앱이 계속 실행되도록 하지만, predict 함수에서 처리

def predict(lyrics: str): # 입력 타입 명시 (Python 3.9+ 에서 권장)
    print(f"--- PREDICT FUNCTION CALLED ---")
    print(f"Received lyrics: '{lyrics}' (Type: {type(lyrics)})")

    if pipeline_instance is None:
        print("Pipeline is not initialized. Cannot predict.")
        return "오류: 모델 파이프라인이 초기화되지 않았습니다. (0.00)"

    if not lyrics or not isinstance(lyrics, str) or lyrics.strip() == "":
        print("Lyrics are empty, not a string, or whitespace only.")
        return "입력 가사가 비어있거나 유효하지 않습니다. (0.00)"
    
    try:
        print("Performing prediction with pipeline...")
        result = pipeline_instance(lyrics) # 전역 pipeline_instance 사용
        print(f"Pipeline raw result: {result}")
    except Exception as e:
        print(f"Error during pipeline prediction: {e}")
        return f"예측 중 오류 발생: {str(e)} (0.00)" # 오류 메시지 포함

    # 결과 추출 로직 (이전과 동일)
    if isinstance(result, list) and len(result) > 0 and isinstance(result[0], dict):
        label = result[0].get("label", "Unknown")
        score = result[0].get("score", 0.0) # 점수가 float인지 확인
        print(f"Extracted prediction: Label='{label}', Score={score}")
        return f"{label} ({score:.2f})"
    else:
        print(f"Pipeline returned no result or unexpected format: {result}")
        return "결과를 추출할 수 없거나 형식이 올바르지 않습니다. (0.00)"

print("Creating Gradio Interface...")
demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Lyrics", placeholder="여기에 가사를 입력하세요..."),
    outputs=gr.Textbox(label="Predicted Genre"),
    title="Lyrics Genre Predictor (Local Docker)",
    description="가사를 입력하면 로컬 Docker 컨테이너에서 실행 중인 모델이 장르를 예측합니다."
    # api_name="predict" # 명시적으로 API 이름을 설정하면 /api/predict 엔드포인트가 확실히 생성됩니다.
)

# ... (이전 답변의 수정된 app.py 코드 전체) ...
if __name__ == "__main__":
    print("Launching Gradio app with queue and explicit API exposure...")
    # queue() 사용 시 /api/predict 가 기본 엔드포인트일 가능성이 높음
    # 또는 api_name을 명시적으로 설정:
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)
    # demo.launch(server_name="0.0.0.0", server_port=7860, api_name="predict") # 이렇게 해도 됨
    print(f"Gradio app should be accessible at http://localhost:7860")
    print(f"Try API endpoint for Next.js: http://localhost:7860/api/predict (or as defined by api_name)")