"""Training task management service.

Background training runs in a daemon thread. A threading.Event per task_id
provides cancellation support. Progress is persisted to DB at most every 2 %.
"""
from __future__ import annotations

import threading
from pathlib import Path
from typing import Any

from app.database import db_session
from app.ml.lstm_trainer import train_lstm
from app.models import DataFile, DataPackage, DataPackageFile
from app.models.train_management import SavedModel, TrainTask
from app.services.data_management_service import DataManagementService

# ── module-level cancel registry ─────────────────────────────────────────────
_cancel_events: dict[int, threading.Event] = {}
_registry_lock = threading.Lock()


class TrainService:
    def __init__(self) -> None:
        self._dm = DataManagementService()
        self._base_dir = Path(__file__).resolve().parents[2]
        self._models_dir = self._base_dir / 'saved_models'
        self._models_dir.mkdir(parents=True, exist_ok=True)

    # ── packages (for dropdown) ──────────────────────────────────────────────

    def list_packages(self, category_id: str = '', name: str = '') -> list[dict[str, Any]]:
        with db_session() as session:
            query = session.query(DataPackage)
            if category_id not in ('', 'all', None):
                try:
                    db_id = int(category_id)
                    query = query.filter(DataPackage.category_id == db_id)
                except (ValueError, TypeError):
                    pass
            if name:
                query = query.filter(DataPackage.name.like(f'%{name}%'))
            rows = query.order_by(DataPackage.created_at.desc()).all()
            return [{'id': p.id, 'name': p.name, 'data_count': p.data_count, 'category_id': p.category_id} for p in rows]

    # ── tasks ────────────────────────────────────────────────────────────────

    def list_tasks(self) -> list[dict[str, Any]]:
        with db_session() as session:
            tasks = session.query(TrainTask).order_by(TrainTask.created_at.desc()).all()
            return [self._task_to_dict(t) for t in tasks]

    def get_task(self, task_id: int) -> dict[str, Any]:
        with db_session() as session:
            task = session.query(TrainTask).filter(TrainTask.id == task_id).first()
            if not task:
                raise ValueError('训练任务不存在')
            return self._task_to_dict(task)

    def create_task(self, model_name: str, train_package_id: int, val_package_id: int, params: dict) -> dict[str, Any]:
        with db_session() as session:
            train_pkg = session.query(DataPackage).filter(DataPackage.id == train_package_id).first()
            if not train_pkg:
                raise ValueError('训练集数据包不存在')
            val_pkg = session.query(DataPackage).filter(DataPackage.id == val_package_id).first()
            if not val_pkg:
                raise ValueError('验证集数据包不存在')

            task = TrainTask(
                model_name=model_name,
                package_id=train_package_id,
                package_name=train_pkg.name,
                val_package_id=val_package_id,
                val_package_name=val_pkg.name,
                params=params,
                status='pending',
                progress=0,
            )
            session.add(task)
            session.commit()
            session.refresh(task)
            task_dict = self._task_to_dict(task)

        self._launch_thread(task_dict['id'], train_package_id, val_package_id, params)
        return task_dict

    def cancel_task(self, task_id: int) -> None:
        with _registry_lock:
            event = _cancel_events.get(task_id)

        if event:
            event.set()
        else:
            # Task may still be in 'pending' state (thread not started yet)
            with db_session() as session:
                task = session.query(TrainTask).filter(TrainTask.id == task_id).first()
                if task and task.status in ('pending', 'running'):
                    task.status = 'cancelled'
                    session.commit()

    def restart_task(self, task_id: int) -> dict[str, Any]:
        with db_session() as session:
            task = session.query(TrainTask).filter(TrainTask.id == task_id).first()
            if not task:
                raise ValueError('任务不存在')
            if task.status in ('pending', 'running'):
                raise ValueError('任务仍在进行中，请先取消')

            new_task = TrainTask(
                model_name=task.model_name,
                package_id=task.package_id,
                package_name=task.package_name,
                val_package_id=task.val_package_id,
                val_package_name=task.val_package_name,
                params=task.params,
                status='pending',
                progress=0,
            )
            session.add(new_task)
            session.commit()
            session.refresh(new_task)
            task_dict = self._task_to_dict(new_task)

        self._launch_thread(task_dict['id'], task_dict['package_id'], task_dict['val_package_id'], task_dict['params'])
        return task_dict

    def delete_task(self, task_id: int) -> None:
        with db_session() as session:
            task = session.query(TrainTask).filter(TrainTask.id == task_id).first()
            if not task:
                raise ValueError('任务不存在')
            if task.status in ('pending', 'running'):
                raise ValueError('请先取消正在进行的任务')

            # Remove model file if exists (and model not saved to saved_models table)
            model_path = self._models_dir / f'task_{task_id}.pt'
            if model_path.exists() and not task.is_saved:
                model_path.unlink(missing_ok=True)

            session.delete(task)
            session.commit()

    def save_model(self, task_id: int, model_name: str | None = None, description: str | None = None) -> dict[str, Any]:
        with db_session() as session:
            task = session.query(TrainTask).filter(TrainTask.id == task_id).first()
            if not task:
                raise ValueError('任务不存在')
            if task.status != 'completed':
                raise ValueError('只能保存已完成的训练任务')
            if task.is_saved:
                raise ValueError('模型已保存')

            model_path = self._models_dir / f'task_{task_id}.pt'
            if not model_path.exists():
                raise ValueError('模型文件不存在，可能已被清除，请重新训练')

            saved = SavedModel(
                task_id=task_id,
                model_name=(model_name.strip() if model_name and model_name.strip() else task.model_name),
                description=description,
                package_id=task.package_id,
                package_name=task.package_name,
                params=task.params,
                file_path=model_path.relative_to(self._base_dir).as_posix(),
                train_loss=task.train_loss,
                val_loss=task.val_loss,
                test_loss=task.test_loss,
            )
            session.add(saved)
            task.is_saved = 1
            session.commit()
            session.refresh(saved)
            return self._saved_to_dict(saved)
    def update_saved_model(self, model_id: int, model_name: str | None = None, description: str | None = None) -> dict[str, Any]:
        with db_session() as session:
            model = session.query(SavedModel).filter(SavedModel.id == model_id).first()
            if not model:
                raise ValueError('模型不存在')
            if model_name is not None and model_name.strip():
                model.model_name = model_name.strip()
            if description is not None:
                model.description = description
            session.commit()
            session.refresh(model)
            return self._saved_to_dict(model)
    # ── saved models ─────────────────────────────────────────────────────────

    def list_saved_models(self) -> list[dict[str, Any]]:
        with db_session() as session:
            models = session.query(SavedModel).order_by(SavedModel.created_at.desc()).all()
            return [self._saved_to_dict(m) for m in models]

    def delete_saved_model(self, model_id: int) -> None:
        with db_session() as session:
            model = session.query(SavedModel).filter(SavedModel.id == model_id).first()
            if not model:
                raise ValueError('模型不存在')

            # Reset is_saved on the originating task if it still exists
            task = session.query(TrainTask).filter(TrainTask.id == model.task_id).first()
            if task:
                task.is_saved = 0

            # Delete the .pt file only if nothing else references it
            model_path = self._resolve_model_path(model.file_path)
            if model_path.exists():
                model_path.unlink(missing_ok=True)

            session.delete(model)
            session.commit()

    # ── private helpers ───────────────────────────────────────────────────────

    def _resolve_model_path(self, file_path: str) -> Path:
        """Return absolute Path whether file_path is relative or absolute."""
        p = Path(file_path)
        if p.is_absolute():
            return p
        return self._base_dir / p

    def _launch_thread(self, task_id: int, train_package_id: int, val_package_id: int | None, params: dict) -> None:
        cancel_event = threading.Event()
        with _registry_lock:
            _cancel_events[task_id] = cancel_event

        thread = threading.Thread(
            target=self._training_worker,
            args=(task_id, train_package_id, val_package_id, params, cancel_event),
            daemon=True,
            name=f'train-task-{task_id}',
        )
        thread.start()

    def _training_worker(
        self,
        task_id: int,
        train_package_id: int,
        val_package_id: int | None,
        params: dict,
        cancel_event: threading.Event,
    ) -> None:
        try:
            self._update_task(task_id, status='running', progress=0)

            train_records = self._load_package_records(train_package_id)
            if not train_records:
                raise ValueError('训练集数据包没有有效数据，请检查关联文件')

            val_records: list[dict[str, Any]] = []
            if val_package_id:
                val_records = self._load_package_records(val_package_id)

            save_path = self._models_dir / f'task_{task_id}.pt'
            epoch_log_buffer: list[dict[str, Any]] = []

            def on_progress(pct: int, epoch: int, train_loss: float, val_loss: float | None) -> None:
                if cancel_event.is_set():
                    return
                epoch_log_buffer.append({
                    'epoch': epoch,
                    'train_loss': round(float(train_loss), 6),
                    'val_loss': round(float(val_loss), 6) if val_loss is not None else None,
                })
                self._update_task(
                    task_id,
                    progress=pct,
                    train_loss=round(float(train_loss), 6),
                    val_loss=round(float(val_loss), 6) if val_loss is not None else None,
                    epoch_logs=list(epoch_log_buffer),
                )

            result = train_lstm(
                train_records=train_records,
                val_records=val_records,
                params=params,
                save_path=save_path,
                on_progress=on_progress,
                cancel_event=cancel_event,
            )

            self._update_task(
                task_id,
                status='completed',
                progress=100,
                train_loss=result['train_loss'],
                val_loss=result.get('val_loss'),
                test_loss=result.get('test_loss'),
            )

        except InterruptedError:
            self._update_task(task_id, status='cancelled')

        except Exception as exc:
            self._update_task(task_id, status='failed', error_msg=str(exc)[:2000])

        finally:
            with _registry_lock:
                _cancel_events.pop(task_id, None)

    def _update_task(self, task_id: int, **kwargs: Any) -> None:
        with db_session() as session:
            task = session.query(TrainTask).filter(TrainTask.id == task_id).first()
            if not task:
                return
            for key, value in kwargs.items():
                setattr(task, key, value)
            session.commit()

    def _load_package_records(self, package_id: int) -> list[dict[str, Any]]:
        with db_session() as session:
            pf_rows = (
                session.query(DataPackageFile)
                .filter(DataPackageFile.package_id == package_id)
                .order_by(DataPackageFile.sort_order.asc())
                .all()
            )
            file_ids = [pf.file_id for pf in pf_rows]
            if not file_ids:
                return []
            files = session.query(DataFile).filter(DataFile.id.in_(file_ids)).all()
            file_map = {f.id: f for f in files}

        all_records: list[dict[str, Any]] = []
        for fid in file_ids:
            if fid not in file_map:
                continue
            fmeta = file_map[fid]
            path = self._dm._resolve_local_file_path(fmeta.file_path, fmeta.stored_name)
            records, _ = self._dm._read_records(path, limit=None)
            all_records.extend(records)

        return all_records

    @staticmethod
    def _task_to_dict(task: TrainTask) -> dict[str, Any]:
        return {
            'id': task.id,
            'model_name': task.model_name,
            'package_id': task.package_id,
            'package_name': task.package_name,
            'val_package_id': task.val_package_id,
            'val_package_name': task.val_package_name,
            'epoch_logs': task.epoch_logs or [],
            'params': task.params,
            'status': task.status,
            'progress': task.progress,
            'train_loss': task.train_loss,
            'val_loss': task.val_loss,
            'test_loss': task.test_loss,
            'error_msg': task.error_msg,
            'is_saved': bool(task.is_saved),
            'created_at': task.created_at.strftime('%Y-%m-%d %H:%M:%S') if task.created_at else '',
        }

    @staticmethod
    def _saved_to_dict(model: SavedModel) -> dict[str, Any]:
        return {
            'id': model.id,
            'task_id': model.task_id,
            'model_name': model.model_name,
            'description': model.description or '',
            'package_id': model.package_id,
            'package_name': model.package_name,
            'params': model.params,
            'train_loss': model.train_loss,
            'val_loss': model.val_loss,
            'test_loss': model.test_loss,
            'created_at': model.created_at.strftime('%Y-%m-%d %H:%M:%S') if model.created_at else '',
        }
