"""Model training API routes."""

import logging
import traceback
from datetime import datetime, timezone

from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from mixlift.deps import check_free_tier_run_limit, get_current_user, get_db
from mixlift.projects.models import Dataset, ModelRun, Project, User

logger = logging.getLogger(__name__)

router = APIRouter(tags=["modeling"])


@router.post("/projects/{project_id}/train", status_code=status.HTTP_202_ACCEPTED)
async def start_training(
    project_id: str,
    background_tasks: BackgroundTasks,
    user: User = Depends(check_free_tier_run_limit),
    session: AsyncSession = Depends(get_db),
):
    # Verify project
    result = await session.execute(
        select(Project).where(Project.id == project_id, Project.user_id == user.id)
    )
    project = result.scalar_one_or_none()
    if not project:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Project not found")

    # Get active dataset
    ds_result = await session.execute(
        select(Dataset)
        .where(Dataset.project_id == project_id, Dataset.is_active.is_(True))
        .order_by(Dataset.created_at.desc())
        .limit(1)
    )
    dataset = ds_result.scalar_one_or_none()
    if not dataset:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail="No active dataset. Upload data first."
        )

    # Check for already running model
    running = await session.execute(
        select(ModelRun).where(
            ModelRun.project_id == project_id, ModelRun.status.in_(["pending", "running"])
        )
    )
    if running.scalar_one_or_none():
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT, detail="A model is already training"
        )

    # Create model run record
    model_run = ModelRun(
        project_id=project.id,
        dataset_id=dataset.id,
        status="pending",
    )
    session.add(model_run)
    await session.flush()

    run_id = str(model_run.id)
    dataset_path = dataset.storage_path

    # Launch background training
    background_tasks.add_task(
        _train_model_background,
        run_id=run_id,
        dataset_path=dataset_path,
        target_metric=project.target_metric,
    )

    return {"run_id": run_id, "status": "pending", "message": "Training started"}


async def _train_model_background(
    run_id: str,
    dataset_path: str,
    target_metric: str,
):
    """Background task to fit the model."""
    import pandas as pd

    from mixlift.db import async_session
    from mixlift.modeling.engine import extract_saturation_curves, run_full_pipeline
    from mixlift.modeling.optimizer import optimize_budget
    from mixlift.modeling.preprocessing import prepare_model_data
    from mixlift.modeling.recommendations import generate_recommendations

    async with async_session() as session:
        result = await session.execute(select(ModelRun).where(ModelRun.id == run_id))
        model_run = result.scalar_one()
        model_run.status = "running"
        model_run.started_at = datetime.now(timezone.utc)
        await session.commit()

    try:
        # Load data
        df = pd.read_parquet(dataset_path)
        X, y, channel_columns = prepare_model_data(df)

        # Fit model
        model_results = run_full_pipeline(X, y, channel_columns)

        # Save model artifact
        from mixlift.config import settings

        artifact_path = str(settings.model_dir / f"{run_id}.nc")
        model_results.mmm.idata.to_netcdf(artifact_path)

        # Extract saturation curves
        sat_curves = extract_saturation_curves(model_results.mmm, X, channel_columns)

        # Optimize
        opt = optimize_budget(model_results.mmm, X, channel_columns)

        # Generate recommendations
        recommendations, summary = generate_recommendations(
            current_allocation=opt["current_allocation"],
            optimal_allocation=opt["optimal_allocation"],
            roas_estimates=model_results.roas_estimates,
            expected_lift_pct=opt["expected_lift_pct"],
        )

        # Save results
        async with async_session() as session:
            result = await session.execute(select(ModelRun).where(ModelRun.id == run_id))
            model_run = result.scalar_one()
            model_run.status = "completed"
            model_run.completed_at = datetime.now(timezone.utc)
            model_run.duration_secs = model_results.duration_secs
            model_run.channel_contributions = model_results.channel_contributions
            model_run.roas_estimates = model_results.roas_estimates
            model_run.roas_ci = {k: list(v) for k, v in model_results.roas_ci.items()}
            model_run.saturation_curves = sat_curves
            model_run.current_allocation = opt["current_allocation"]
            model_run.optimal_allocation = opt["optimal_allocation"]
            model_run.expected_lift_pct = opt["expected_lift_pct"]
            model_run.recommendations = [
                {
                    "rank": r.rank,
                    "action": r.action,
                    "channel": r.channel,
                    "change_amount": r.change_amount,
                    "current_spend": r.current_spend,
                    "optimal_spend": r.optimal_spend,
                    "roas": r.roas,
                    "reason": r.reason,
                }
                for r in recommendations
            ]
            model_run.summary_text = summary
            model_run.model_artifact_path = artifact_path
            await session.commit()

        logger.info(f"Training completed for run {run_id} in {model_results.duration_secs:.1f}s")

    except Exception as e:
        logger.error(f"Training failed for run {run_id}: {e}\n{traceback.format_exc()}")
        async with async_session() as session:
            result = await session.execute(select(ModelRun).where(ModelRun.id == run_id))
            model_run = result.scalar_one()
            model_run.status = "failed"
            model_run.error_message = str(e)
            model_run.completed_at = datetime.now(timezone.utc)
            await session.commit()


@router.get("/projects/{project_id}/runs/latest")
async def get_latest_run(
    project_id: str,
    user: User = Depends(get_current_user),
    session: AsyncSession = Depends(get_db),
):
    # Verify project ownership
    result = await session.execute(
        select(Project).where(Project.id == project_id, Project.user_id == user.id)
    )
    if not result.scalar_one_or_none():
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Project not found")

    run_result = await session.execute(
        select(ModelRun)
        .where(ModelRun.project_id == project_id)
        .order_by(ModelRun.created_at.desc())
        .limit(1)
    )
    run = run_result.scalar_one_or_none()
    if not run:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="No model runs found")

    return {
        "id": str(run.id),
        "status": run.status,
        "error_message": run.error_message,
        "started_at": str(run.started_at) if run.started_at else None,
        "completed_at": str(run.completed_at) if run.completed_at else None,
        "duration_secs": run.duration_secs,
        "channel_contributions": run.channel_contributions,
        "roas_estimates": run.roas_estimates,
        "roas_ci": run.roas_ci,
        "saturation_curves": run.saturation_curves,
        "current_allocation": run.current_allocation,
        "optimal_allocation": run.optimal_allocation,
        "expected_lift_pct": run.expected_lift_pct,
        "recommendations": run.recommendations,
        "summary_text": run.summary_text,
        "model_artifact_path": run.model_artifact_path,
    }
