"""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._models_dir = Path(__file__).resolve().parents[2] / 'saved_models'
        self._models_dir.mkdir(parents=True, exist_ok=True)

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

    def list_packages(self) -> list[dict[str, Any]]:
        with db_session() as session:
            rows = (
                session.query(DataPackage)
                .order_by(DataPackage.created_at.desc())
                .all()
            )
            return [{'id': p.id, 'name': p.name, 'data_count': p.data_count} 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, package_id: int, params: dict) -> dict[str, Any]:
        with db_session() as session:
            pkg = session.query(DataPackage).filter(DataPackage.id == package_id).first()
            if not pkg:
                raise ValueError('数据包不存在')

            task = TrainTask(
                model_name=model_name,
                package_id=package_id,
                package_name=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'], 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,
                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['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) -> 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=task.model_name,
                package_id=task.package_id,
                package_name=task.package_name,
                params=task.params,
                file_path=str(model_path),
                train_loss=task.train_loss,
                val_loss=task.val_loss,
            )
            session.add(saved)
            task.is_saved = 1
            session.commit()
            session.refresh(saved)
            return self._saved_to_dict(saved)

    # ── 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 = Path(model.file_path)
            if model_path.exists():
                model_path.unlink(missing_ok=True)

            session.delete(model)
            session.commit()

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

    def _launch_thread(self, task_id: int, package_id: int, 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, package_id, params, cancel_event),
            daemon=True,
            name=f'train-task-{task_id}',
        )
        thread.start()

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

            records = self._load_package_records(package_id)
            if not records:
                raise ValueError('数据包没有有效数据，请检查关联文件')

            save_path = self._models_dir / f'task_{task_id}.pt'
            last_pct = [0]

            def on_progress(pct: int, train_loss: float, val_loss: float | None) -> None:
                if cancel_event.is_set():
                    return
                # Throttle: persist at most every 2 % to reduce DB writes
                if pct - last_pct[0] >= 2 or pct == 100:
                    last_pct[0] = pct
                    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,
                    )

            result = train_lstm(
                records=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'),
            )

        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,
            'params': task.params,
            'status': task.status,
            'progress': task.progress,
            'train_loss': task.train_loss,
            'val_loss': task.val_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,
            'package_id': model.package_id,
            'package_name': model.package_name,
            'params': model.params,
            'train_loss': model.train_loss,
            'val_loss': model.val_loss,
            'created_at': model.created_at.strftime('%Y-%m-%d %H:%M:%S') if model.created_at else '',
        }
