from __future__ import annotations

from typing import Any, Optional

from sqlalchemy import and_, select
from sqlalchemy.ext.asyncio import AsyncSession

from ...models import (
    QualityAsset,
    QualityAssetPair,
    TrainedModelArtifact,
    TrainedModelDeployment,
    TrainingDataset,
    TrainingDatasetItem,
    TrainingEvaluation,
    TrainingJob,
    TrainingJobRun,
)
from ...utils.timezone import now_cn_naive


EVALUATION_TYPES = ("offline_judge", "structure_compliance", "shadow_replay")


class TrainingPipelineService:
    def __init__(self, db: AsyncSession):
        self.db = db

    @staticmethod
    def _dataset_snapshot(row: TrainingDataset) -> dict[str, Any]:
        return {
            "id": int(row.id),
            "tenant_id": int(row.tenant_id),
            "project_id": int(row.project_id) if row.project_id is not None else None,
            "dataset_type": row.dataset_type,
            "version": row.version,
            "status": row.status,
            "sample_count": int(row.sample_count or 0),
            "selection_rules": row.selection_rules or {},
            "channel_scope": row.channel_scope or {},
            "notes": row.notes,
        }

    def _job_snapshot(self, row: TrainingJob, *, artifact: Optional[TrainedModelArtifact] = None) -> dict[str, Any]:
        payload = {
            "id": int(row.id),
            "tenant_id": int(row.tenant_id),
            "project_id": int(row.project_id) if row.project_id is not None else None,
            "dataset_id": int(row.dataset_id),
            "job_type": row.job_type,
            "base_model_family": row.base_model_family,
            "target_slot": row.target_slot,
            "status": row.status,
            "provider_backend": row.provider_backend,
            "artifact_ref": row.artifact_ref,
            "metrics_json": row.metrics_json or {},
        }
        if artifact is not None:
            payload["artifact"] = self._artifact_snapshot(artifact)
        return payload

    @staticmethod
    def _artifact_snapshot(row: TrainedModelArtifact) -> dict[str, Any]:
        return {
            "id": int(row.id),
            "tenant_id": int(row.tenant_id),
            "training_job_id": int(row.training_job_id),
            "artifact_ref": row.artifact_ref,
            "base_model_family": row.base_model_family,
            "target_slot": row.target_slot,
            "status": row.status,
            "metrics_json": row.metrics_json or {},
        }

    @staticmethod
    def _deployment_snapshot(row: TrainedModelDeployment) -> dict[str, Any]:
        return {
            "id": int(row.id),
            "tenant_id": int(row.tenant_id),
            "artifact_id": int(row.artifact_id),
            "deployment_status": row.deployment_status,
            "bound_tier": row.bound_tier,
            "effective_scope": row.effective_scope or {},
            "decision": row.decision,
            "notes": row.notes,
            "decided_by": int(row.decided_by) if row.decided_by is not None else None,
            "decided_at": row.decided_at,
        }

    @staticmethod
    def _evaluation_snapshot(row: TrainingEvaluation) -> dict[str, Any]:
        return {
            "id": int(row.id),
            "tenant_id": int(row.tenant_id),
            "training_job_id": int(row.training_job_id),
            "evaluation_type": row.evaluation_type,
            "status": row.status,
            "score": float(row.score or 0.0) if row.score is not None else None,
            "evaluation_summary": row.evaluation_summary or {},
        }

    async def create_dataset(
        self,
        *,
        tenant_id: int,
        actor_id: Optional[int],
        dataset_type: str,
        version: str,
        asset_ids: list[int],
        pair_ids: list[int],
        selection_rules: Optional[dict[str, Any]],
        notes: Optional[str],
    ) -> dict[str, Any]:
        dataset_type = str(dataset_type or "").strip().lower()
        if dataset_type not in {"sft", "preference", "eval"}:
            raise ValueError("unsupported dataset_type")
        asset_ids = [int(x) for x in asset_ids or []]
        pair_ids = [int(x) for x in pair_ids or []]

        assets = []
        if asset_ids:
            assets = (
                await self.db.execute(
                    select(QualityAsset).where(
                        QualityAsset.tenant_id == int(tenant_id),
                        QualityAsset.id.in_(asset_ids),
                    )
                )
            ).scalars().all()
        pairs = []
        if pair_ids:
            pairs = (
                await self.db.execute(
                    select(QualityAssetPair).where(
                        QualityAssetPair.tenant_id == int(tenant_id),
                        QualityAssetPair.id.in_(pair_ids),
                    )
                )
            ).scalars().all()

        if dataset_type == "sft":
            if not assets or any(asset.asset_type != "golden" or asset.status != "approved" for asset in assets):
                raise ValueError("SFT dataset only accepts approved golden assets")
        if dataset_type == "preference":
            if not pairs or any(pair.status != "approved" for pair in pairs):
                raise ValueError("Preference dataset only accepts approved asset pairs")
        if dataset_type == "eval" and not (assets or pairs):
            raise ValueError("Eval dataset requires assets or pairs")

        project_id = None
        if assets:
            project_id = assets[0].project_id
        elif pairs:
            project_id = pairs[0].project_id
        dataset = TrainingDataset(
            tenant_id=int(tenant_id),
            project_id=int(project_id) if project_id is not None else None,
            dataset_type=dataset_type,
            version=str(version or "").strip() or f"{dataset_type}-v1",
            status="ready",
            sample_count=len(assets) + len(pairs),
            selection_rules=selection_rules or {},
            channel_scope={},
            notes=str(notes or "").strip() or None,
            created_by=actor_id,
        )
        self.db.add(dataset)
        await self.db.flush()

        for asset in assets:
            self.db.add(
                TrainingDatasetItem(
                    tenant_id=int(tenant_id),
                    dataset_id=int(dataset.id),
                    quality_asset_id=int(asset.id),
                    redacted_snapshot={
                        "question_text": asset.question_text,
                        "answer_text": asset.answer_text,
                        "corrected_answer": asset.corrected_answer,
                        "prompt_key": asset.prompt_key,
                    },
                )
            )
        for pair in pairs:
            self.db.add(
                TrainingDatasetItem(
                    tenant_id=int(tenant_id),
                    dataset_id=int(dataset.id),
                    quality_asset_pair_id=int(pair.id),
                    redacted_snapshot={
                        "pair_type": pair.pair_type,
                        "reason_text": pair.reason_text,
                    },
                )
            )
        await self.db.commit()
        await self.db.refresh(dataset)
        return self._dataset_snapshot(dataset)

    async def create_job(
        self,
        *,
        tenant_id: int,
        actor_id: Optional[int],
        dataset_id: int,
        job_type: str,
        base_model_family: str,
        target_slot: str,
        provider_backend: str,
    ) -> dict[str, Any]:
        normalized_job_type = str(job_type or "").strip().lower()
        if normalized_job_type not in {"sft", "preference"}:
            raise ValueError("unsupported training job_type")
        normalized_target_slot = str(target_slot or "").strip().lower()
        if normalized_target_slot not in {"low", "high"}:
            raise ValueError("unsupported target_slot")
        dataset = (
            await self.db.execute(
                select(TrainingDataset).where(
                    TrainingDataset.tenant_id == int(tenant_id),
                    TrainingDataset.id == int(dataset_id),
                )
            )
        ).scalar_one_or_none()
        if dataset is None:
            raise ValueError("training dataset not found")

        job = TrainingJob(
            tenant_id=int(tenant_id),
            project_id=int(dataset.project_id) if dataset.project_id is not None else None,
            dataset_id=int(dataset.id),
            job_type=normalized_job_type,
            base_model_family=str(base_model_family or "").strip(),
            target_slot=normalized_target_slot,
            status="queued",
            provider_backend=str(provider_backend or "").strip(),
            artifact_ref=None,
            metrics_json={},
            created_by=actor_id,
        )
        self.db.add(job)
        await self.db.flush()

        run = TrainingJobRun(
            tenant_id=int(tenant_id),
            training_job_id=int(job.id),
            status="queued",
            provider_run_ref=None,
            metrics_json={},
        )
        self.db.add(run)

        artifact = TrainedModelArtifact(
            tenant_id=int(tenant_id),
            training_job_id=int(job.id),
            artifact_ref=f"artifact:{job.job_type}:{job.id}",
            base_model_family=job.base_model_family,
            target_slot=job.target_slot,
            status="pending_eval",
            metrics_json={},
        )
        self.db.add(artifact)
        await self.db.flush()
        job.artifact_ref = artifact.artifact_ref

        for evaluation_type in EVALUATION_TYPES:
            self.db.add(
                TrainingEvaluation(
                    tenant_id=int(tenant_id),
                    training_job_id=int(job.id),
                    evaluation_type=evaluation_type,
                    status="pending",
                    score=None,
                    evaluation_summary={},
                )
            )

        await self.db.commit()
        await self.db.refresh(job)
        await self.db.refresh(artifact)
        return self._job_snapshot(job, artifact=artifact)

    async def list_datasets(
        self,
        *,
        tenant_id: int,
        dataset_type: Optional[str] = None,
        limit: int = 20,
        offset: int = 0,
    ) -> dict[str, Any]:
        stmt = select(TrainingDataset).where(TrainingDataset.tenant_id == int(tenant_id))
        if dataset_type:
            stmt = stmt.where(TrainingDataset.dataset_type == str(dataset_type).strip().lower())
        rows = (
            await self.db.execute(
                stmt.order_by(TrainingDataset.id.desc()).offset(int(offset)).limit(int(limit))
            )
        ).scalars().all()
        return {"items": [self._dataset_snapshot(row) for row in rows], "total": len(rows), "limit": int(limit), "offset": int(offset)}

    async def list_jobs(
        self,
        *,
        tenant_id: int,
        status: Optional[str] = None,
        limit: int = 20,
        offset: int = 0,
    ) -> dict[str, Any]:
        stmt = select(TrainingJob).where(TrainingJob.tenant_id == int(tenant_id))
        if status:
            stmt = stmt.where(TrainingJob.status == str(status).strip().lower())
        rows = (
            await self.db.execute(
                stmt.order_by(TrainingJob.id.desc()).offset(int(offset)).limit(int(limit))
            )
        ).scalars().all()
        return {"items": [self._job_snapshot(row) for row in rows], "total": len(rows), "limit": int(limit), "offset": int(offset)}

    async def list_evaluations(self, *, tenant_id: int, job_id: int) -> dict[str, Any]:
        rows = (
            await self.db.execute(
                select(TrainingEvaluation)
                .where(
                    TrainingEvaluation.tenant_id == int(tenant_id),
                    TrainingEvaluation.training_job_id == int(job_id),
                )
                .order_by(TrainingEvaluation.id.asc())
            )
        ).scalars().all()
        return {"items": [self._evaluation_snapshot(row) for row in rows], "total": len(rows)}

    async def promote_artifact(
        self,
        *,
        tenant_id: int,
        artifact_id: int,
        actor_id: Optional[int],
    ) -> dict[str, Any]:
        artifact = (
            await self.db.execute(
                select(TrainedModelArtifact).where(
                    TrainedModelArtifact.tenant_id == int(tenant_id),
                    TrainedModelArtifact.id == int(artifact_id),
                )
            )
        ).scalar_one_or_none()
        if artifact is None:
            raise ValueError("trained model artifact not found")
        artifact.status = "candidate"
        row = TrainedModelDeployment(
            tenant_id=int(tenant_id),
            artifact_id=int(artifact.id),
            deployment_status="candidate",
            bound_tier=str(artifact.target_slot or "").strip().lower() or "low",
            effective_scope={},
            decision=None,
            notes=None,
            decided_by=actor_id,
            decided_at=now_cn_naive().isoformat(),
        )
        self.db.add(row)
        await self.db.commit()
        await self.db.refresh(row)
        return self._deployment_snapshot(row)

    async def decide_deployment(
        self,
        *,
        tenant_id: int,
        deployment_id: int,
        actor_id: Optional[int],
        decision: str,
        notes: Optional[str],
    ) -> dict[str, Any]:
        row = (
            await self.db.execute(
                select(TrainedModelDeployment).where(
                    and_(
                        TrainedModelDeployment.tenant_id == int(tenant_id),
                        TrainedModelDeployment.id == int(deployment_id),
                    )
                )
            )
        ).scalar_one_or_none()
        if row is None:
            raise ValueError("trained model deployment not found")
        normalized = str(decision or "").strip().lower()
        if normalized not in {"accepted", "rejected", "rolled_back"}:
            raise ValueError("unsupported decision")
        row.decision = normalized
        row.notes = str(notes or "").strip() or None
        row.decided_by = actor_id
        row.decided_at = now_cn_naive().isoformat()
        row.deployment_status = "active" if normalized == "accepted" else "rollback"
        await self.db.commit()
        await self.db.refresh(row)
        return self._deployment_snapshot(row)

    async def list_deployments(
        self,
        *,
        tenant_id: int,
        deployment_status: Optional[str] = None,
        limit: int = 20,
        offset: int = 0,
    ) -> dict[str, Any]:
        stmt = select(TrainedModelDeployment).where(TrainedModelDeployment.tenant_id == int(tenant_id))
        if deployment_status:
            stmt = stmt.where(
                TrainedModelDeployment.deployment_status == str(deployment_status).strip().lower()
            )
        rows = (
            await self.db.execute(
                stmt.order_by(TrainedModelDeployment.id.desc()).offset(int(offset)).limit(int(limit))
            )
        ).scalars().all()
        return {
            "items": [self._deployment_snapshot(row) for row in rows],
            "total": len(rows),
            "limit": int(limit),
            "offset": int(offset),
        }