backend/app/api/ai_builder.py

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from typing import List, Dict, Any
import json

from app.db.session import get_db
from app.services.user_service import user_service
from app.models.database import UserAIModel, User as DBUser
from app.core.auth import get_user_or_demo

# Import from our new model_builder package
from app.model_builder import (
    ModelTrainer, 
    ModelPublisher,
    UniverseItemSchema,
    ModelSaveSchema,
    ModelResponseSchema,
    TrainRequestSchema
)

router = APIRouter()

# --- Universe Management ---

@router.get("/universe", response_model=List[str])
async def get_universe(db: AsyncSession = Depends(get_db), user: DBUser = Depends(get_user_or_demo)):
    items = await user_service.get_universe(db, user.id)
    return [item.symbol for item in items]

@router.post("/universe")
async def add_to_universe(item: UniverseItemSchema, db: AsyncSession = Depends(get_db), user: DBUser = Depends(get_user_or_demo)):
    await user_service.add_to_universe(db, user.id, item.symbol, item.name)
    return {"status": "success"}

@router.post("/universe/sync-all")
async def sync_all(db: AsyncSession = Depends(get_db), user: DBUser = Depends(get_user_or_demo)):
    items = await user_service.get_universe(db, user.id)
    for item in items:
        await user_service.sync_historical_data(db, item.symbol)
    return {"status": "success"}

@router.get("/features-status")
async def get_features_status(db: AsyncSession = Depends(get_db), user: DBUser = Depends(get_user_or_demo)):
    return await user_service.get_features_status(db, user.id)

# --- AI Model Management ---

@router.get("/models", response_model=List[ModelResponseSchema])
async def get_models(db: AsyncSession = Depends(get_db), user: DBUser = Depends(get_user_or_demo)):
    result = await db.execute(
        select(UserAIModel)
        .where(UserAIModel.user_id == user.id)
        .order_by(UserAIModel.created_at.desc())
    )
    models = result.scalars().all()
    
    return [
        {
            "id": m.id,
            "name": m.name,
            "model_type": m.model_type,
            "target_symbol": m.target_symbol,
            "created_at": m.created_at.isoformat(),
            "metrics": json.loads(m.metrics) if m.metrics else {}
        }
        for m in models
    ]

@router.post("/models")
async def save_model(data: ModelSaveSchema, db: AsyncSession = Depends(get_db), user: DBUser = Depends(get_user_or_demo)):
    model = UserAIModel(
        user_id=user.id,
        name=data.name,
        model_type=data.model_type,
        target_symbol=data.target_symbol,
        parameters=json.dumps(data.parameters),
        metrics=json.dumps(data.metrics)
    )
    db.add(model)
    await db.commit()
    return {"status": "success", "id": model.id}

@router.delete("/models/{model_id}")
async def delete_model(model_id: int, db: AsyncSession = Depends(get_db), user: DBUser = Depends(get_user_or_demo)):
    publisher = ModelPublisher(db, user.id)
    return await publisher.delete_model(model_id)

# --- Training & Publishing ---

@router.post("/train")
async def train_model(data: TrainRequestSchema, db: AsyncSession = Depends(get_db), user: DBUser = Depends(get_user_or_demo)):
    trainer = ModelTrainer(db, user.id)
    result = await trainer.train(
        target_symbol=data.target_symbol,
        model_type=data.model_type,
        features=data.features,
        test_size=data.test_size
    )
    return result

@router.post("/models/{model_id}/upload-to-hf")
async def upload_model_to_hf(model_id: int, db: AsyncSession = Depends(get_db), user: DBUser = Depends(get_user_or_demo)):
    publisher = ModelPublisher(db, user.id)
    return await publisher.upload_to_hf(model_id)