import uuid

import mlflow
import pandas as pd
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import settings
from app.models import Prediction, Variant

_model = None


def load_model():
    global _model
    if _model is None:
        mlflow.set_tracking_uri(settings.mlflow_tracking_uri)
        _model = mlflow.pyfunc.load_model(f"models:/{settings.model_name}/{settings.model_stage}")
    return _model


def featurise(variants: list[Variant]) -> pd.DataFrame:
    # Mirror ml/rarelens_ml/features.py exactly; shared package later.
    return pd.DataFrame(
        {
            "impact": [v.impact for v in variants],
            "consequence": [v.consequence for v in variants],
            "gnomad_af": [v.gnomad_af if v.gnomad_af is not None else 0.0 for v in variants],
            "cadd_phred": [v.annotations.get("CADD_PHRED") for v in variants],
            "am_pathogenicity": [v.annotations.get("am_pathogenicity") for v in variants],
        }
    )


async def score_job(job_id: uuid.UUID, session: AsyncSession) -> int:
    variants = (await session.scalars(select(Variant).where(Variant.job_id == job_id))).all()
    if not variants:
        return 0
    model = load_model()
    scores = model.predict(featurise(variants))
    for v, s in zip(variants, scores):
        session.add(Prediction(variant_id=v.id, model_name=settings.model_name,
                               model_version=settings.model_stage, score=float(s)))
    await session.commit()
    return len(variants)