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

ROOT = Path(__file__).resolve().parent
PD_DIR = ROOT / '.paddledetection'
PD_REPO = 'https://github.com/PaddlePaddle/PaddleDetection.git'
PD_BRANCH = os.getenv('PADDLEDETECTION_BRANCH', 'release/2.9')

_backend_lock = threading.RLock()
_training = {"running": False, "message": "not started", "error": None, "epoch": None, "log": ""}


def _run(cmd, cwd=None, timeout=None):
    p = subprocess.run(cmd, cwd=cwd, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=timeout)
    if p.returncode != 0:
        raise RuntimeError(p.stdout[-16000:] or f'command failed: {p.returncode}')
    return p.stdout


def ensure_paddledetection():
    """Clone official PaddleDetection only when first needed; no Hugging Face dependency."""
    marker = PD_DIR / 'tools' / 'train.py'
    if marker.exists():
        return PD_DIR
    PD_DIR.parent.mkdir(parents=True, exist_ok=True)
    tmp = ROOT / '.paddledetection_tmp'
    if tmp.exists():
        shutil.rmtree(tmp, ignore_errors=True)
    _run(['git', 'clone', '--depth', '1', '--branch', PD_BRANCH, PD_REPO, str(tmp)])
    if marker.exists():
        return PD_DIR
    if PD_DIR.exists():
        shutil.rmtree(PD_DIR, ignore_errors=True)
    tmp.rename(PD_DIR)
    return PD_DIR


def paddle_status():
    try:
        import paddle
        return {"installed": True, "version": paddle.__version__, "device": paddle.device.get_device(), "cuda": bool(paddle.is_compiled_with_cuda())}
    except Exception as e:
        return {"installed": False, "error": str(e)}


def prepare_config(dataset_root: Path, classes_file: Path, output_dir: Path,
                   epochs: int, batch_size: int, learning_rate: float):
    """Create a native PaddleDetection RT-DETR R50vd config."""
    pd = ensure_paddledetection()
    cfg_dir = pd / 'configs' / 'rtdetr'
    ds_cfg = pd / 'configs' / 'datasets' / 'icecream_coco.yml'
    model_cfg = cfg_dir / 'icecream_rtdetr_r50vd.yml'
    cfg_dir.mkdir(parents=True, exist_ok=True)

    classes = [x.strip() for x in classes_file.read_text(encoding='utf-8').splitlines() if x.strip()]
    if not classes:
        raise RuntimeError('No classes configured.')

    train_dir = (dataset_root / 'train').resolve().as_posix()
    val_dir = (dataset_root / 'val').resolve().as_posix()
    save_dir = output_dir.resolve().as_posix()

    ds_cfg.write_text(f"""metric: COCO
num_classes: {len(classes)}

TrainDataset:
  !COCODataSet
    image_dir: images
    anno_path: annotations.json
    dataset_dir: {train_dir}
    data_fields: ['image', 'gt_bbox', 'gt_class', 'is_crowd']

EvalDataset:
  !COCODataSet
    image_dir: images
    anno_path: annotations.json
    dataset_dir: {val_dir}

TestDataset:
  !ImageFolder
    anno_path: annotations.json
    dataset_dir: {val_dir}
""", encoding='utf-8')

    model_cfg.write_text(f"""_BASE_: [
  '../datasets/icecream_coco.yml',
  '../runtime.yml',
  '_base_/optimizer_6x.yml',
  '_base_/rtdetr_r50vd.yml',
  '_base_/rtdetr_reader.yml',
]

architecture: DETR
num_classes: {len(classes)}
use_gpu: true
find_unused_parameters: True
log_iter: 1
save_dir: {save_dir}
weights: {save_dir}/model_final
pretrain_weights: https://paddledet.bj.bcebos.com/models/pretrained/ResNet50_vd_ssld_v2_pretrained.pdparams
epoch: {int(epochs)}

LearningRate:
  base_lr: {float(learning_rate)}

TrainReader:
  batch_size: {int(batch_size)}
  worker_num: 0
  use_shared_memory: false

EvalReader:
  batch_size: 1
  worker_num: 0

TestReader:
  batch_size: 1
  worker_num: 0
""", encoding='utf-8')
    return pd, model_cfg, classes

def train(dataset_root: Path, classes_file: Path, output_dir: Path, epochs: int, batch_size: int, learning_rate: float):
    if 'transformers' in sys.modules or 'torch' in sys.modules:
        raise RuntimeError(
            'Wrong training engine loaded: this package uses PaddleDetection RT-DETR, '
            'not PyTorch or Hugging Face Transformers. Restart with this package.'
        )

    global _training
    with _backend_lock:
        if _training['running']:
            return _training
        _training = {"running": True, "message": "preparing PaddleDetection", "error": None, "epoch": None, "log": ""}
    try:
        pd, cfg, classes = prepare_config(dataset_root, classes_file, output_dir, epochs, batch_size, learning_rate)
        output_dir.mkdir(parents=True, exist_ok=True)
        status = paddle_status()
        if not status.get('installed'):
            raise RuntimeError(f'PaddlePaddle is not installed: {status.get("error")}')
        if not status.get('cuda'):
            raise RuntimeError(
                'GPU PaddlePaddle is not available. '
                f'Paddle device={status.get("device")}. Install the GPU PaddlePaddle build '
                'matching the cloud CUDA runtime.'
            )
        env = os.environ.copy()
        env['CUDA_VISIBLE_DEVICES'] = env.get('CUDA_VISIBLE_DEVICES', '0')
        env['FLAGS_selected_gpus'] = env.get('FLAGS_selected_gpus', '0')
        cmd = [os.sys.executable, '-u', 'tools/train.py', '-c', str(cfg), '--eval', '-o', 'use_gpu=true']
        _training['message'] = 'training RT-DETR with PaddleDetection'
        proc = subprocess.Popen(cmd, cwd=pd, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1)
        lines = []
        for line in proc.stdout:
            lines.append(line.rstrip())
            if len(lines) > 250:
                lines = lines[-250:]
            _training['log'] = '\n'.join(lines[-80:])
            low = line.lower()
            if 'epoch' in low:
                import re
                m = re.search(r'epoch[^0-9]*(\d+)', low)
                if m:
                    _training['epoch'] = int(m.group(1))
        code = proc.wait()
        if code != 0:
            raise RuntimeError('\n'.join(lines[-120:]) or f'PaddleDetection training exited with {code}')
        final = output_dir / 'model_final.pdparams'
        if not final.exists():
            # Some PaddleDetection versions save model_final as a prefix with .pdparams.
            candidates = list(output_dir.glob('model_final*'))
            if not candidates:
                raise RuntimeError('Training finished but model_final was not produced.')
        (output_dir / 'classes.json').write_text(json.dumps({str(i): c for i, c in enumerate(classes)}, indent=2), encoding='utf-8')
        _training = {"running": False, "message": "training complete", "error": None, "epoch": int(epochs), "log": '\n'.join(lines[-80:])}
    except Exception as e:
        _training = {"running": False, "message": "training failed", "error": str(e), "epoch": _training.get('epoch'), "log": _training.get('log', '')}
    return _training


def training_status():
    return dict(_training)


def model_ready(model_dir: Path):
    return any((model_dir / name).exists() for name in ('model_final.pdparams', 'model_final.pdiparams'))


def predict(image_path: Path, model_dir: Path, threshold: float, classes_file: Path, output_dir: Path):
    pd = ensure_paddledetection()
    if not model_ready(model_dir):
        raise RuntimeError('No trained model yet. Train the model first.')
    cfg = pd / 'configs' / 'rtdetr' / 'icecream_rtdetr_r50vd.yml'
    if not cfg.exists():
        raise RuntimeError('PaddleDetection model config is missing; train once to generate it.')
    output_dir.mkdir(parents=True, exist_ok=True)
    for p in output_dir.iterdir():
        if p.is_file():
            p.unlink()
    cmd = [os.sys.executable, '-u', 'tools/infer.py', '-c', str(cfg), '--infer_img', str(image_path), '--output_dir', str(output_dir), '--draw_threshold', str(threshold), '--save_threshold', str(threshold), '--save_results=True', '--visualize=True', '-o', f'use_gpu=true', f'weights={str(model_dir / "model_final")}']
    log = _run(cmd, cwd=pd, timeout=int(os.getenv('INFER_TIMEOUT', '300')))
    bbox_file = output_dir / 'bbox.json'
    rows = json.loads(bbox_file.read_text(encoding='utf-8')) if bbox_file.exists() else []
    classes = [x.strip() for x in classes_file.read_text(encoding='utf-8').splitlines() if x.strip()]
    detections = []
    for r in rows:
        cid = int(r.get('category_id', 0))
        detections.append({
            'class': classes[cid] if 0 <= cid < len(classes) else str(cid),
            'class_id': cid,
            'confidence': round(float(r.get('score', 0.0)), 4),
            'box': [round(float(r['bbox'][0]), 2), round(float(r['bbox'][1]), 2), round(float(r['bbox'][0] + r['bbox'][2]), 2), round(float(r['bbox'][1] + r['bbox'][3]), 2)]
        })
    counts = {}
    for d in detections:
        counts[d['class']] = counts.get(d['class'], 0) + 1
    visual = next((p for p in output_dir.iterdir() if p.suffix.lower() in {'.jpg','.jpeg','.png'}), None)
    return {'total': len(detections), 'counts': dict(sorted(counts.items())), 'detections': detections, 'visual_path': str(visual) if visual else None, 'log': log[-4000:]}