Blame for main.py

"""
DataFlow Server — Main FastAPI Application
"""
import uuid
import json
from datetime import datetime
from typing import List

from fastapi import FastAPI, Depends, HTTPException, status, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select

from database import get_db, init_db, pwd_context
from models import User, Project, Task, Expense, Document, Notification, AIForecast
import schemas
import ai_service

app = FastAPI(title="DataFlow API", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.on_event("startup")
async def startup_event():
    await init_db()

# ─── Auth Routes ───
@app.post("/auth/login", response_model=schemas.Token)
async def login(credentials: schemas.UserLogin, db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(User).where(User.email == credentials.email))
    user = result.scalar_one_or_none()
    
    if not user or not pwd_context.verify(credentials.password, user.hashed_password):
        raise HTTPException(status_code=401, detail="Incorrect email or password")
        
    return {
        "access_token": "fake-jwt-token-for-demo",
        "token_type": "bearer",
        "user_id": user.id,
        "full_name": user.full_name,
        "role": user.role
    }

@app.post("/auth/register", response_model=schemas.Token)
async def register(user_data: schemas.UserRegister, db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(User).where(User.email == user_data.email))
    if result.scalar_one_or_none():
        raise HTTPException(status_code=400, detail="Email already registered")
        
    new_user = User(
        id=str(uuid.uuid4()),
        full_name=user_data.full_name,
        email=user_data.email,
        hashed_password=pwd_context.hash(user_data.password),
        position=user_data.position,
        role=user_data.role
    )
    db.add(new_user)
    await db.commit()
    
    return {
        "access_token": "fake-jwt-token-for-demo",
        "token_type": "bearer",
        "user_id": new_user.id,
        "full_name": new_user.full_name,
        "role": new_user.role
    }

# ─── Project Routes ───
@app.get("/projects", response_model=List[schemas.ProjectOut])
async def get_projects(db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(Project))
    return result.scalars().all()

@app.post("/projects", response_model=schemas.ProjectOut)
async def create_project(project: schemas.ProjectCreate, db: AsyncSession = Depends(get_db)):
    new_project = Project(
        id=str(uuid.uuid4()),
        **project.model_dump()
    )
    db.add(new_project)
    await db.commit()
    await db.refresh(new_project)
    return new_project

@app.patch("/projects/{project_id}", response_model=schemas.ProjectOut)
async def update_project(project_id: str, updates: schemas.ProjectUpdate, db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(Project).where(Project.id == project_id))
    project = result.scalar_one_or_none()
    if not project:
        raise HTTPException(status_code=404, detail="Project not found")
        
    for key, value in updates.model_dump(exclude_unset=True).items():
        setattr(project, key, value)
        
    await db.commit()
    await db.refresh(project)
    return project

# ─── Task Routes ───
@app.get("/tasks", response_model=List[schemas.TaskOut])
async def get_tasks(db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(Task))
    return result.scalars().all()

@app.post("/tasks", response_model=schemas.TaskOut)
async def create_task(task: schemas.TaskCreate, db: AsyncSession = Depends(get_db)):
    new_task = Task(id=str(uuid.uuid4()), **task.model_dump())
    db.add(new_task)
    await db.commit()
    await db.refresh(new_task)
    return new_task

@app.patch("/tasks/{task_id}", response_model=schemas.TaskOut)
async def update_task(task_id: str, updates: schemas.TaskUpdate, db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(Task).where(Task.id == task_id))
    task = result.scalar_one_or_none()
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
        
    for key, value in updates.model_dump(exclude_unset=True).items():
        setattr(task, key, value)
        
    await db.commit()
    await db.refresh(task)
    return task

# ─── Expense Routes ───
@app.get("/expenses", response_model=List[schemas.ExpenseOut])
async def get_expenses(db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(Expense))
    return result.scalars().all()

@app.post("/expenses", response_model=schemas.ExpenseOut)
async def create_expense(expense: schemas.ExpenseCreate, db: AsyncSession = Depends(get_db)):
    new_expense = Expense(
        id=str(uuid.uuid4()),
        **expense.model_dump(exclude={"date"}),
        date=expense.date or datetime.utcnow()
    )
    db.add(new_expense)
    
    # Update project spent amount
    project_result = await db.execute(select(Project).where(Project.id == expense.project_id))
    project = project_result.scalar_one_or_none()
    if project:
        project.spent += expense.amount
        
    await db.commit()
    await db.refresh(new_expense)
    return new_expense

@app.patch("/expenses/{expense_id}/status")
async def update_expense_status(expense_id: str, status: str, db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(Expense).where(Expense.id == expense_id))
    expense = result.scalar_one_or_none()
    if not expense:
        raise HTTPException(status_code=404, detail="Expense not found")
        
    expense.status = status
    await db.commit()
    return {"status": "success"}

# ─── Document & OCR Routes ───
@app.get("/documents", response_model=List[schemas.DocumentOut])
async def get_documents(db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(Document))
    return result.scalars().all()

@app.post("/documents/upload", response_model=schemas.DocumentOut)
async def upload_document(
    project_id: str = Form(...),
    file: UploadFile = File(...),
    db: AsyncSession = Depends(get_db)
):
    content = await file.read()
    file_ext = file.filename.split('.')[-1].lower() if '.' in file.filename else 'unknown'
    
    doc = Document(
        id=str(uuid.uuid4()),
        name=file.filename,
        project_id=project_id,
        type=file_ext,
        size_bytes=len(content)
    )
    
    # Trigger OCR if image or pdf
    if file_ext in ['jpg', 'jpeg', 'png', 'pdf']:
        ocr_result = await ai_service.ocr_extract(content, file.filename)
        if ocr_result.get("success"):
            doc.extracted_text = ocr_result.get("extracted_text")
            doc.recognition_confidence = ocr_result.get("confidence")
            doc.detected_type = ocr_result.get("detected_type")
            doc.extracted_amount = ocr_result.get("extracted_amount")
            
            # Create notification for OCR result
            notif = Notification(
                id=str(uuid.uuid4()),
                title="Документ распознан",
                body=f"Документ {file.filename} успешно обработан AI.",
                type="document"
            )
            db.add(notif)
            
    db.add(doc)
    await db.commit()
    await db.refresh(doc)
    return doc

# ─── AI Forecast Routes ───
@app.get("/forecasts", response_model=List[schemas.ForecastOut])
async def get_forecasts(db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(AIForecast))
    return result.scalars().all()

@app.post("/forecasts/generate", response_model=schemas.ForecastOut)
async def generate_forecast(req: schemas.ForecastRequest, db: AsyncSession = Depends(get_db)):
    # 1. Get project
    proj_res = await db.execute(select(Project).where(Project.id == req.project_id))
    project = proj_res.scalar_one_or_none()
    if not project:
        raise HTTPException(status_code=404, detail="Project not found")
        
    # 2. Get tasks to calc completion
    task_res = await db.execute(select(Task).where(Task.project_id == project.id))
    tasks = task_res.scalars().all()
    completed_tasks = sum(1 for t in tasks if t.status == "done")
    
    # 3. Get expenses for categories
    exp_res = await db.execute(select(Expense).where(Expense.project_id == project.id))
    expenses = exp_res.scalars().all()
    cats = {}
    for e in expenses:
        cats[e.category] = cats.get(e.category, 0) + e.amount
        
    # 4. Calc days
    now = datetime.utcnow()
    total_days = (project.end_date - project.start_date).days
    days_elapsed = (now - project.start_date).days
    if total_days <= 0: total_days = 1
    if days_elapsed < 0: days_elapsed = 0
    
    # 5. Call AI
    ai_result = await ai_service.generate_forecast(
        project_name=project.name,
        budget=project.budget,
        spent=project.spent,
        days_elapsed=days_elapsed,
        total_days=total_days,
        expense_categories=cats,
        completed_tasks=completed_tasks,
        total_tasks=len(tasks)
    )
    
    if not ai_result.get("success"):
        raise HTTPException(status_code=500, detail="Failed to generate AI forecast")
        
    # 6. Save or update forecast
    fc_res = await db.execute(select(AIForecast).where(AIForecast.project_id == project.id))
    forecast = fc_res.scalar_one_or_none()
    
    if forecast:
        forecast.current_spent = project.spent
        forecast.planned_budget = project.budget
        forecast.risk_level = ai_result.get("risk_level", "green")
        forecast.overrun_probability = ai_result.get("overrun_probability", 0)
        forecast.optimistic_forecast = ai_result.get("optimistic_forecast", 0)
        forecast.realistic_forecast = ai_result.get("realistic_forecast", 0)
        forecast.pessimistic_forecast = ai_result.get("pessimistic_forecast", 0)
        forecast.recommendations = json.dumps(ai_result.get("recommendations", []))
        forecast.generated_at = datetime.utcnow()
    else:
        forecast = AIForecast(
            id=str(uuid.uuid4()),
            project_id=project.id,
            current_spent=project.spent,
            planned_budget=project.budget,
            risk_level=ai_result.get("risk_level", "green"),
            overrun_probability=ai_result.get("overrun_probability", 0),
            optimistic_forecast=ai_result.get("optimistic_forecast", 0),
            realistic_forecast=ai_result.get("realistic_forecast", 0),
            pessimistic_forecast=ai_result.get("pessimistic_forecast", 0),
            recommendations=json.dumps(ai_result.get("recommendations", []))
        )
        db.add(forecast)
        
    await db.commit()
    await db.refresh(forecast)
    return forecast

# ─── Notifications Routes ───
@app.get("/notifications", response_model=List[schemas.NotificationOut])
async def get_notifications(db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(Notification).order_by(Notification.created_at.desc()))
    return result.scalars().all()

@app.patch("/notifications/{notif_id}/read")
async def mark_notification_read(notif_id: str, db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(Notification).where(Notification.id == notif_id))
    notif = result.scalar_one_or_none()
    if notif:
        notif.is_read = True
        await db.commit()
    return {"status": "success"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)