import json
import os
import shutil
import threading
import uuid
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel

from paddle_backend import train, training_status, paddle_status, predict, model_ready

ROOT = Path(__file__).resolve().parent

if os.getenv('DATA_DIR'):
    BASE = Path(os.environ['DATA_DIR'])
elif Path('/data').exists() and os.access('/data', os.W_OK):
    BASE = Path('/data') / 'icecream_counter'
else:
    BASE = Path('./data')
DATASET = BASE / 'generated_dataset'
MODEL = BASE / 'model'
CLASSES = BASE / 'classes.txt'
API_OUTPUT = BASE / 'api_outputs'
API_OUTPUT.mkdir(parents=True, exist_ok=True)

app = FastAPI(title='Ice Cream RT-DETR API', version='1.0.0')

class TrainRequest(BaseModel):
    epochs: int = 30
    batch_size: int = 2
    learning_rate: float = 1e-5

@app.get('/dataset/status')
def dataset_status():
    return {
        'train_annotations': (DATASET / 'train' / 'annotations.json').exists(),
        'val_annotations': (DATASET / 'val' / 'annotations.json').exists(),
        'classes': [x.strip() for x in CLASSES.read_text(encoding='utf-8').splitlines() if x.strip()] if CLASSES.exists() else [],
    }

@app.get('/health')
def health():
    p = paddle_status()
    return {
        'ok': True,
        'engine': 'PaddleDetection RT-DETR',
        'huggingface': False,
        'transformers': False,
        'torch_training': False,
        'paddle': p,
        'model_ready': model_ready(MODEL)
    }

@app.get('/train/status')
def status():
    return training_status()

@app.post('/train')
def start_train(req: TrainRequest):
    s = training_status()
    if s['running']:
        return s
    if not (DATASET / 'train' / 'annotations.json').exists():
        raise HTTPException(400, 'Training dataset is not prepared.')
    if not CLASSES.exists():
        raise HTTPException(400, 'classes.txt is missing.')
    def job():
        train(DATASET, CLASSES, MODEL, max(1, req.epochs), max(1, req.batch_size), req.learning_rate)
    threading.Thread(target=job, daemon=True).start()
    return {'running': True, 'message': 'training started', 'engine': 'PaddleDetection RT-DETR'}

@app.post('/predict')
async def predict_image(file: UploadFile = File(...), threshold: float = 0.35):
    if not model_ready(MODEL):
        raise HTTPException(409, 'No trained model yet.')
    job_id = uuid.uuid4().hex
    job_dir = API_OUTPUT / job_id
    job_dir.mkdir(parents=True, exist_ok=True)
    suffix = Path(file.filename or 'image.jpg').suffix or '.jpg'
    src = job_dir / ('input' + suffix)
    src.write_bytes(await file.read())
    try:
        result = predict(src, MODEL, max(0.0, min(1.0, threshold)), CLASSES, job_dir)
        result['job_id'] = job_id
        return result
    except Exception as e:
        shutil.rmtree(job_dir, ignore_errors=True)
        raise HTTPException(500, str(e))

@app.get('/predict/{job_id}/image')
def predicted_image(job_id: str):
    d = API_OUTPUT / job_id
    for p in d.glob('*'):
        if p.suffix.lower() in {'.jpg','.jpeg','.png'} and p.name != 'input.jpg':
            return FileResponse(p)
    raise HTTPException(404, 'Prediction image not found.')