from fastapi import FastAPI, APIRouter, HTTPException, Depends, status, UploadFile, File
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from dotenv import load_dotenv
from starlette.middleware.cors import CORSMiddleware
from motor.motor_asyncio import AsyncIOMotorClient
import os
import logging
from pathlib import Path
from pydantic import BaseModel, Field, ConfigDict, EmailStr
from typing import List, Optional, Dict, Any
import uuid
from datetime import datetime, timezone, timedelta
import bcrypt
import jwt
from enum import Enum
import io
from fastapi.responses import StreamingResponse

ROOT_DIR = Path(__file__).parent
#load_dotenv(ROOT_DIR / '.env')
load_dotenv(ROOT_DIR / '.env', override=True)
load_dotenv(Path(__file__).resolve().parent.parent / '.env', override=True)  
# app root fallback

# MongoDB connection
#mongo_url = os.environ['MONGO_URL']
mongo_url = os.environ.get("MONGO_URL")
if not mongo_url:
    raise RuntimeError("MONGO_URL not set. Check /backend/.env or app env vars.")
# client = AsyncIOMotorClient(mongo_url)
client = AsyncIOMotorClient(
    mongo_url,
    serverSelectionTimeoutMS=5000,
    connectTimeoutMS=5000,
    socketTimeoutMS=5000,
)
#db = client[os.environ['DB_NAME']]
db_name = os.environ.get("DB_NAME")
if not db_name:
    raise RuntimeError("DB_NAME not set.")
db = client[db_name]

# JWT Configuration
JWT_SECRET = os.environ.get('JWT_SECRET', 'kmc-secret-key-2024-production')
JWT_ALGORITHM = "HS256"
JWT_EXPIRATION_HOURS = 24

# Create the main app
app = FastAPI(title="KMC Performance Monitoring System")
api_router = APIRouter(prefix="/api")
security = HTTPBearer()

# Enums
class UserRole(str, Enum):
    SUPERVISOR = "supervisor"
    ADMIN_OFFICER = "admin_officer"
    SCM = "scm"
    KIBARANI_OFFICER = "kibarani_officer"
    MANAGEMENT = "management"
    SYSTEM_ADMIN = "system_admin"

class SalesStatus(str, Enum):
    DRAFT = "draft"
    SUBMITTED = "submitted"
    RETURNED = "returned"
    POSTED = "posted"

class StockStatus(str, Enum):
    DRAFT = "draft"
    POSTED = "posted"

class PaymentMethod(str, Enum):
    COOP = "coop"
    KCB = "kcb"
    ECITIZEN = "ecitizen"
    SWIPE = "swipe"

# Models
class UserCreate(BaseModel):
    email: EmailStr
    password: str
    full_name: str
    role: UserRole
    shop_id: Optional[str] = None

class UserLogin(BaseModel):
    email: EmailStr
    password: str

class UserResponse(BaseModel):
    model_config = ConfigDict(extra="ignore")
    id: str
    email: str
    full_name: str
    role: UserRole
    shop_id: Optional[str] = None
    is_active: bool = True
    created_at: str

class TokenResponse(BaseModel):
    access_token: str
    token_type: str = "bearer"
    user: UserResponse

class ShopCreate(BaseModel):
    name: str
    location: str
    opening_date: str
    is_active: bool = True

class ShopResponse(BaseModel):
    model_config = ConfigDict(extra="ignore")
    id: str
    name: str
    location: str
    opening_date: str
    is_active: bool
    created_at: str

class PaymentBreakdown(BaseModel):
    coop: float = 0
    kcb: float = 0
    ecitizen: float = 0
    swipe: float = 0

class SalesLineItem(BaseModel):
    category: str
    product_name: str
    quantity_kg: float
    unit_price: float
    total_amount: float

class SalesEntryCreate(BaseModel):
    shop_id: str
    date: str
    kmc_products_total: float
    kenchic_total: float = 0
    payment_breakdown: PaymentBreakdown
    line_items: List[SalesLineItem] = []
    notes: Optional[str] = None

class SalesEntryResponse(BaseModel):
    model_config = ConfigDict(extra="ignore")
    id: str
    shop_id: str
    shop_name: Optional[str] = None
    date: str
    kmc_products_total: float
    kenchic_total: float
    total_sales: float
    payment_breakdown: PaymentBreakdown
    line_items: List[SalesLineItem]
    status: SalesStatus
    submitted_by: Optional[str] = None
    submitted_at: Optional[str] = None
    posted_by: Optional[str] = None
    posted_at: Optional[str] = None
    return_reason: Optional[str] = None
    notes: Optional[str] = None
    created_at: str
    updated_at: str

class StockLineItem(BaseModel):
    category: str
    subcategory: Optional[str] = None
    product_name: str
    quantity_kg: float
    unit: str = "kg"
    status: str = "available"  # available, low, out_of_stock

class StockEntryCreate(BaseModel):
    location: str  # "hq" or "kibarani"
    date: str
    stock_items: List[StockLineItem]
    notes: Optional[str] = None

class StockEntryResponse(BaseModel):
    model_config = ConfigDict(extra="ignore")
    id: str
    location: str
    date: str
    stock_items: List[StockLineItem]
    status: StockStatus
    submitted_by: Optional[str] = None
    submitted_at: Optional[str] = None
    notes: Optional[str] = None
    created_at: str
    updated_at: str

class LivestockEntry(BaseModel):
    date: str
    cattle_received: int = 0
    shoats_received: int = 0
    cattle_slaughtered: int = 0
    shoats_slaughtered: int = 0
    notes: Optional[str] = None

class LivestockResponse(BaseModel):
    model_config = ConfigDict(extra="ignore")
    id: str
    date: str
    cattle_received: int
    shoats_received: int
    cattle_slaughtered: int
    shoats_slaughtered: int
    notes: Optional[str]
    created_at: str
    updated_at: str

class ProductionOps(BaseModel):
    date: str
    operations: List[str]  # ["deboning", "slaughter", "value_addition", etc.]
    slaughter_start_time: Optional[str] = None
    batches_completed: int = 0
    dispatch_details: Optional[str] = None
    notes: Optional[str] = None

class AuditLogEntry(BaseModel):
    action: str
    entity_type: str
    entity_id: str
    user_id: str
    user_email: str
    changes: Dict[str, Any]
    timestamp: str

class DashboardStats(BaseModel):
    livestock_received_cattle: int
    livestock_received_shoats: int
    livestock_slaughtered_cattle: int
    livestock_slaughtered_shoats: int
    total_sales_today: float
    total_sales_week: float
    cold_room_stock_kg: float
    production_status: List[str]
    shops_count: int
    pending_submissions: int

# Helper functions
def hash_password(password: str) -> str:
    return bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')

def verify_password(password: str, hashed: str) -> bool:
    return bcrypt.checkpw(password.encode('utf-8'), hashed.encode('utf-8'))

def create_token(user_id: str, email: str, role: str) -> str:
    payload = {
        "user_id": user_id,
        "email": email,
        "role": role,
        "exp": datetime.now(timezone.utc) + timedelta(hours=JWT_EXPIRATION_HOURS)
    }
    return jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM)

async def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(security)):
    try:
        payload = jwt.decode(credentials.credentials, JWT_SECRET, algorithms=[JWT_ALGORITHM])
        user = await db.users.find_one({"id": payload["user_id"]}, {"_id": 0})
        if not user:
            raise HTTPException(status_code=401, detail="User not found")
        return user
    except jwt.ExpiredSignatureError:
        raise HTTPException(status_code=401, detail="Token expired")
    except jwt.InvalidTokenError:
        raise HTTPException(status_code=401, detail="Invalid token")

def require_roles(*roles: UserRole):
    async def role_checker(user: dict = Depends(get_current_user)):
        if user["role"] not in [r.value for r in roles]:
            raise HTTPException(status_code=403, detail="Insufficient permissions")
        return user
    return role_checker

def sanitize_for_json(obj):
    """Recursively convert ObjectId and other non-serializable types to strings"""
    if isinstance(obj, dict):
        return {k: sanitize_for_json(v) for k, v in obj.items() if k != '_id'}
    elif isinstance(obj, list):
        return [sanitize_for_json(item) for item in obj]
    elif hasattr(obj, '__str__') and type(obj).__name__ == 'ObjectId':
        return str(obj)
    return obj

async def log_audit(action: str, entity_type: str, entity_id: str, user: dict, changes: dict):
    # Sanitize changes to remove ObjectId and _id fields
    sanitized_changes = sanitize_for_json(changes)
    audit_entry = {
        "id": str(uuid.uuid4()),
        "action": action,
        "entity_type": entity_type,
        "entity_id": entity_id,
        "user_id": user["id"],
        "user_email": user["email"],
        "changes": sanitized_changes,
        "timestamp": datetime.now(timezone.utc).isoformat()
    }
    await db.audit_logs.insert_one(audit_entry)

# Auth Routes
@api_router.post("/auth/register", response_model=UserResponse)
async def register(user_data: UserCreate):
    existing = await db.users.find_one({"email": user_data.email})
    if existing:
        raise HTTPException(status_code=400, detail="Email already registered")
    
    user = {
        "id": str(uuid.uuid4()),
        "email": user_data.email,
        "password_hash": hash_password(user_data.password),
        "full_name": user_data.full_name,
        "role": user_data.role.value,
        "shop_id": user_data.shop_id,
        "is_active": True,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "updated_at": datetime.now(timezone.utc).isoformat()
    }
    await db.users.insert_one(user)
    return UserResponse(
        id=user["id"],
        email=user["email"],
        full_name=user["full_name"],
        role=user["role"],
        shop_id=user["shop_id"],
        is_active=user["is_active"],
        created_at=user["created_at"]
    )

@api_router.post("/auth/login", response_model=TokenResponse)
async def login(credentials: UserLogin):
    user = await db.users.find_one({"email": credentials.email}, {"_id": 0})
    if not user or not verify_password(credentials.password, user["password_hash"]):
        raise HTTPException(status_code=401, detail="Invalid credentials")
    
    if not user.get("is_active", True):
        raise HTTPException(status_code=401, detail="Account disabled")
    
    token = create_token(user["id"], user["email"], user["role"])
    return TokenResponse(
        access_token=token,
        user=UserResponse(
            id=user["id"],
            email=user["email"],
            full_name=user["full_name"],
            role=user["role"],
            shop_id=user.get("shop_id"),
            is_active=user.get("is_active", True),
            created_at=user["created_at"]
        )
    )

@api_router.get("/auth/me", response_model=UserResponse)
async def get_me(user: dict = Depends(get_current_user)):
    return UserResponse(
        id=user["id"],
        email=user["email"],
        full_name=user["full_name"],
        role=user["role"],
        shop_id=user.get("shop_id"),
        is_active=user.get("is_active", True),
        created_at=user["created_at"]
    )

# User Management Routes
@api_router.get("/users", response_model=List[UserResponse])
async def get_users(user: dict = Depends(require_roles(UserRole.SYSTEM_ADMIN))):
    users = await db.users.find({}, {"_id": 0, "password_hash": 0}).to_list(1000)
    return [UserResponse(**u) for u in users]

@api_router.put("/users/{user_id}/toggle-active")
async def toggle_user_active(user_id: str, user: dict = Depends(require_roles(UserRole.SYSTEM_ADMIN))):
    target_user = await db.users.find_one({"id": user_id})
    if not target_user:
        raise HTTPException(status_code=404, detail="User not found")
    
    new_status = not target_user.get("is_active", True)
    await db.users.update_one({"id": user_id}, {"$set": {"is_active": new_status}})
    await log_audit("toggle_active", "user", user_id, user, {"is_active": new_status})
    return {"success": True, "is_active": new_status}

# Shop Routes
@api_router.get("/shops", response_model=List[ShopResponse])
async def get_shops(user: dict = Depends(get_current_user)):
    shops = await db.shops.find({}, {"_id": 0}).to_list(100)
    return [ShopResponse(**s) for s in shops]

@api_router.post("/shops", response_model=ShopResponse)
async def create_shop(shop_data: ShopCreate, user: dict = Depends(require_roles(UserRole.SYSTEM_ADMIN))):
    shop = {
        "id": str(uuid.uuid4()),
        "name": shop_data.name,
        "location": shop_data.location,
        "opening_date": shop_data.opening_date,
        "is_active": shop_data.is_active,
        "created_at": datetime.now(timezone.utc).isoformat()
    }
    await db.shops.insert_one(shop)
    await log_audit("create", "shop", shop["id"], user, shop)
    return ShopResponse(**shop)

# Sales Routes
@api_router.get("/sales", response_model=List[SalesEntryResponse])
async def get_sales(
    shop_id: Optional[str] = None,
    status: Optional[SalesStatus] = None,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    user: dict = Depends(get_current_user)
):
    query = {}
    
    # Filter by shop for supervisors
    if user["role"] == UserRole.SUPERVISOR.value and user.get("shop_id"):
        query["shop_id"] = user["shop_id"]
    elif shop_id:
        query["shop_id"] = shop_id
    
    if status:
        query["status"] = status.value
    if start_date:
        query["date"] = {"$gte": start_date}
    if end_date:
        query.setdefault("date", {})["$lte"] = end_date
    
    sales = await db.sales.find(query, {"_id": 0}).sort("date", -1).to_list(1000)
    
    # Add shop names
    shops = {s["id"]: s["name"] for s in await db.shops.find({}, {"_id": 0}).to_list(100)}
    for sale in sales:
        sale["shop_name"] = shops.get(sale["shop_id"], "Unknown")
    
    return [SalesEntryResponse(**s) for s in sales]

@api_router.post("/sales", response_model=SalesEntryResponse)
async def create_sales_entry(
    sales_data: SalesEntryCreate,
    user: dict = Depends(require_roles(UserRole.SUPERVISOR, UserRole.ADMIN_OFFICER, UserRole.SYSTEM_ADMIN))
):
    # Validate shop exists
    shop = await db.shops.find_one({"id": sales_data.shop_id})
    if not shop:
        raise HTTPException(status_code=404, detail="Shop not found")
    
    # Calculate total
    total_sales = sales_data.kmc_products_total + sales_data.kenchic_total
    
    # Validate payment breakdown matches total
    payment_total = (
        sales_data.payment_breakdown.coop +
        sales_data.payment_breakdown.kcb +
        sales_data.payment_breakdown.ecitizen +
        sales_data.payment_breakdown.swipe
    )
    
    if abs(payment_total - total_sales) > 0.01:
        raise HTTPException(status_code=400, detail=f"Payment breakdown ({payment_total}) must match total sales ({total_sales})")
    
    now = datetime.now(timezone.utc).isoformat()
    sales_entry = {
        "id": str(uuid.uuid4()),
        "shop_id": sales_data.shop_id,
        "date": sales_data.date,
        "kmc_products_total": sales_data.kmc_products_total,
        "kenchic_total": sales_data.kenchic_total,
        "total_sales": total_sales,
        "payment_breakdown": sales_data.payment_breakdown.model_dump(),
        "line_items": [item.model_dump() for item in sales_data.line_items],
        "status": SalesStatus.DRAFT.value,
        "submitted_by": None,
        "submitted_at": None,
        "posted_by": None,
        "posted_at": None,
        "return_reason": None,
        "notes": sales_data.notes,
        "created_at": now,
        "updated_at": now
    }
    
    await db.sales.insert_one(sales_entry)
    sales_entry["shop_name"] = shop["name"]
    await log_audit("create", "sales", sales_entry["id"], user, {"status": "draft"})
    
    return SalesEntryResponse(**sales_entry)

@api_router.put("/sales/{sales_id}", response_model=SalesEntryResponse)
async def update_sales_entry(
    sales_id: str,
    sales_data: SalesEntryCreate,
    user: dict = Depends(require_roles(UserRole.SUPERVISOR, UserRole.ADMIN_OFFICER))
):
    existing = await db.sales.find_one({"id": sales_id})
    if not existing:
        raise HTTPException(status_code=404, detail="Sales entry not found")
    
    if existing["status"] == SalesStatus.POSTED.value:
        raise HTTPException(status_code=400, detail="Cannot edit posted sales entry")
    
    total_sales = sales_data.kmc_products_total + sales_data.kenchic_total
    
    update_data = {
        "kmc_products_total": sales_data.kmc_products_total,
        "kenchic_total": sales_data.kenchic_total,
        "total_sales": total_sales,
        "payment_breakdown": sales_data.payment_breakdown.model_dump(),
        "line_items": [item.model_dump() for item in sales_data.line_items],
        "notes": sales_data.notes,
        "updated_at": datetime.now(timezone.utc).isoformat()
    }
    
    await db.sales.update_one({"id": sales_id}, {"$set": update_data})
    await log_audit("update", "sales", sales_id, user, update_data)
    
    updated = await db.sales.find_one({"id": sales_id}, {"_id": 0})
    shop = await db.shops.find_one({"id": updated["shop_id"]})
    updated["shop_name"] = shop["name"] if shop else "Unknown"
    
    return SalesEntryResponse(**updated)

@api_router.post("/sales/{sales_id}/submit")
async def submit_sales(sales_id: str, user: dict = Depends(require_roles(UserRole.SUPERVISOR, UserRole.ADMIN_OFFICER, UserRole.SYSTEM_ADMIN))):
    existing = await db.sales.find_one({"id": sales_id})
    if not existing:
        raise HTTPException(status_code=404, detail="Sales entry not found")
    
    if existing["status"] not in [SalesStatus.DRAFT.value, SalesStatus.RETURNED.value]:
        raise HTTPException(status_code=400, detail="Can only submit draft or returned entries")
    
    now = datetime.now(timezone.utc).isoformat()
    await db.sales.update_one({"id": sales_id}, {
        "$set": {
            "status": SalesStatus.SUBMITTED.value,
            "submitted_by": user["id"],
            "submitted_at": now,
            "updated_at": now
        }
    })
    await log_audit("submit", "sales", sales_id, user, {"status": "submitted"})
    return {"success": True, "status": "submitted"}

@api_router.post("/sales/{sales_id}/return")
async def return_sales(
    sales_id: str,
    reason: str,
    user: dict = Depends(require_roles(UserRole.ADMIN_OFFICER, UserRole.SYSTEM_ADMIN))
):
    existing = await db.sales.find_one({"id": sales_id})
    if not existing:
        raise HTTPException(status_code=404, detail="Sales entry not found")
    
    if existing["status"] != SalesStatus.SUBMITTED.value:
        raise HTTPException(status_code=400, detail="Can only return submitted entries")
    
    now = datetime.now(timezone.utc).isoformat()
    await db.sales.update_one({"id": sales_id}, {
        "$set": {
            "status": SalesStatus.RETURNED.value,
            "return_reason": reason,
            "updated_at": now
        }
    })
    await log_audit("return", "sales", sales_id, user, {"status": "returned", "reason": reason})
    return {"success": True, "status": "returned"}

@api_router.post("/sales/{sales_id}/post")
async def post_sales(sales_id: str, user: dict = Depends(require_roles(UserRole.ADMIN_OFFICER, UserRole.SYSTEM_ADMIN))):
    existing = await db.sales.find_one({"id": sales_id})
    if not existing:
        raise HTTPException(status_code=404, detail="Sales entry not found")
    
    if existing["status"] != SalesStatus.SUBMITTED.value:
        raise HTTPException(status_code=400, detail="Can only post submitted entries")
    
    now = datetime.now(timezone.utc).isoformat()
    await db.sales.update_one({"id": sales_id}, {
        "$set": {
            "status": SalesStatus.POSTED.value,
            "posted_by": user["id"],
            "posted_at": now,
            "updated_at": now
        }
    })
    await log_audit("post", "sales", sales_id, user, {"status": "posted"})
    return {"success": True, "status": "posted"}

# Stock Routes
@api_router.get("/stock", response_model=List[StockEntryResponse])
async def get_stock(
    location: Optional[str] = None,
    date: Optional[str] = None,
    user: dict = Depends(get_current_user)
):
    query = {}
    if location:
        query["location"] = location
    if date:
        query["date"] = date
    
    stock = await db.stock.find(query, {"_id": 0}).sort("date", -1).to_list(1000)
    return [StockEntryResponse(**s) for s in stock]

@api_router.post("/stock", response_model=StockEntryResponse)
async def create_stock_entry(
    stock_data: StockEntryCreate,
    user: dict = Depends(require_roles(UserRole.SCM, UserRole.KIBARANI_OFFICER, UserRole.SYSTEM_ADMIN))
):
    now = datetime.now(timezone.utc).isoformat()
    stock_entry = {
        "id": str(uuid.uuid4()),
        "location": stock_data.location,
        "date": stock_data.date,
        "stock_items": [item.model_dump() for item in stock_data.stock_items],
        "status": StockStatus.POSTED.value,
        "submitted_by": user["id"],
        "submitted_at": now,
        "notes": stock_data.notes,
        "created_at": now,
        "updated_at": now
    }
    
    await db.stock.insert_one(stock_entry)
    await log_audit("create", "stock", stock_entry["id"], user, {"location": stock_data.location})
    
    return StockEntryResponse(**stock_entry)

@api_router.get("/stock/latest")
async def get_latest_stock(user: dict = Depends(get_current_user)):
    """Get the latest stock entries for both HQ and Kibarani"""
    hq_stock = await db.stock.find_one({"location": "hq"}, {"_id": 0}, sort=[("date", -1)])
    kibarani_stock = await db.stock.find_one({"location": "kibarani"}, {"_id": 0}, sort=[("date", -1)])
    
    return {
        "hq": hq_stock,
        "kibarani": kibarani_stock
    }

# Livestock Routes
@api_router.get("/livestock", response_model=List[LivestockResponse])
async def get_livestock(
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    user: dict = Depends(get_current_user)
):
    query = {}
    if start_date:
        query["date"] = {"$gte": start_date}
    if end_date:
        query.setdefault("date", {})["$lte"] = end_date
    
    livestock = await db.livestock.find(query, {"_id": 0}).sort("date", -1).to_list(1000)
    return [LivestockResponse(**l) for l in livestock]

@api_router.post("/livestock", response_model=LivestockResponse)
async def create_livestock_entry(
    entry: LivestockEntry,
    user: dict = Depends(require_roles(UserRole.SCM, UserRole.SYSTEM_ADMIN))
):
    now = datetime.now(timezone.utc).isoformat()
    livestock_entry = {
        "id": str(uuid.uuid4()),
        "date": entry.date,
        "cattle_received": entry.cattle_received,
        "shoats_received": entry.shoats_received,
        "cattle_slaughtered": entry.cattle_slaughtered,
        "shoats_slaughtered": entry.shoats_slaughtered,
        "notes": entry.notes,
        "created_at": now,
        "updated_at": now
    }
    
    await db.livestock.insert_one(livestock_entry)
    await log_audit("create", "livestock", livestock_entry["id"], user, livestock_entry)
    
    return LivestockResponse(**livestock_entry)

# Dashboard Routes
@api_router.get("/dashboard/stats", response_model=DashboardStats)
async def get_dashboard_stats(user: dict = Depends(get_current_user)):
    today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    week_ago = (datetime.now(timezone.utc) - timedelta(days=7)).strftime("%Y-%m-%d")
    
    # Get today's livestock
    today_livestock = await db.livestock.find_one({"date": today}, {"_id": 0})
    
    # Get today's sales
    today_sales = await db.sales.find({
        "date": today,
        "status": SalesStatus.POSTED.value
    }, {"_id": 0}).to_list(1000)
    total_sales_today = sum(s.get("total_sales", 0) for s in today_sales)
    
    # Get week's sales
    week_sales = await db.sales.find({
        "date": {"$gte": week_ago, "$lte": today},
        "status": SalesStatus.POSTED.value
    }, {"_id": 0}).to_list(1000)
    total_sales_week = sum(s.get("total_sales", 0) for s in week_sales)
    
    # Get latest stock for cold room
    latest_stock = await db.stock.find_one({"location": "hq"}, {"_id": 0}, sort=[("date", -1)])
    cold_room_stock = 0
    if latest_stock:
        cold_room_stock = sum(item.get("quantity_kg", 0) for item in latest_stock.get("stock_items", []))
    
    # Get production ops
    today_ops = await db.production_ops.find_one({"date": today}, {"_id": 0})
    production_status = today_ops.get("operations", []) if today_ops else []
    
    # Get shop count
    shops_count = await db.shops.count_documents({"is_active": True})
    
    # Get pending submissions
    pending_submissions = await db.sales.count_documents({"status": SalesStatus.SUBMITTED.value})
    
    return DashboardStats(
        livestock_received_cattle=today_livestock.get("cattle_received", 0) if today_livestock else 0,
        livestock_received_shoats=today_livestock.get("shoats_received", 0) if today_livestock else 0,
        livestock_slaughtered_cattle=today_livestock.get("cattle_slaughtered", 0) if today_livestock else 0,
        livestock_slaughtered_shoats=today_livestock.get("shoats_slaughtered", 0) if today_livestock else 0,
        total_sales_today=total_sales_today,
        total_sales_week=total_sales_week,
        cold_room_stock_kg=cold_room_stock,
        production_status=production_status,
        shops_count=shops_count,
        pending_submissions=pending_submissions
    )

@api_router.get("/dashboard/sales-by-shop")
async def get_sales_by_shop(days: int = 7, user: dict = Depends(get_current_user)):
    start_date = (datetime.now(timezone.utc) - timedelta(days=days)).strftime("%Y-%m-%d")
    
    sales = await db.sales.find({
        "date": {"$gte": start_date},
        "status": SalesStatus.POSTED.value
    }, {"_id": 0}).to_list(1000)
    
    shops = {s["id"]: s["name"] for s in await db.shops.find({}, {"_id": 0}).to_list(100)}
    
    # Aggregate by shop
    shop_sales = {}
    for sale in sales:
        shop_id = sale["shop_id"]
        shop_name = shops.get(shop_id, "Unknown")
        if shop_name not in shop_sales:
            shop_sales[shop_name] = {
                "total": 0,
                "coop": 0,
                "kcb": 0,
                "ecitizen": 0,
                "swipe": 0
            }
        shop_sales[shop_name]["total"] += sale.get("total_sales", 0)
        pb = sale.get("payment_breakdown", {})
        shop_sales[shop_name]["coop"] += pb.get("coop", 0)
        shop_sales[shop_name]["kcb"] += pb.get("kcb", 0)
        shop_sales[shop_name]["ecitizen"] += pb.get("ecitizen", 0)
        shop_sales[shop_name]["swipe"] += pb.get("swipe", 0)
    
    return [{"shop": k, **v} for k, v in shop_sales.items()]

@api_router.get("/dashboard/payment-breakdown")
async def get_payment_breakdown(days: int = 7, user: dict = Depends(get_current_user)):
    start_date = (datetime.now(timezone.utc) - timedelta(days=days)).strftime("%Y-%m-%d")
    
    sales = await db.sales.find({
        "date": {"$gte": start_date},
        "status": SalesStatus.POSTED.value
    }, {"_id": 0}).to_list(1000)
    
    totals = {"coop": 0, "kcb": 0, "ecitizen": 0, "swipe": 0}
    for sale in sales:
        pb = sale.get("payment_breakdown", {})
        totals["coop"] += pb.get("coop", 0)
        totals["kcb"] += pb.get("kcb", 0)
        totals["ecitizen"] += pb.get("ecitizen", 0)
        totals["swipe"] += pb.get("swipe", 0)
    
    return [
        {"name": "Co-op", "value": totals["coop"]},
        {"name": "KCB", "value": totals["kcb"]},
        {"name": "E-Citizen", "value": totals["ecitizen"]},
        {"name": "Swipe/PDQ", "value": totals["swipe"]}
    ]

@api_router.get("/dashboard/stock-position")
async def get_stock_position(user: dict = Depends(get_current_user)):
    latest_stock = await db.stock.find_one({"location": "hq"}, {"_id": 0}, sort=[("date", -1)])
    
    if not latest_stock:
        return []
    
    # Aggregate by category
    categories = {}
    for item in latest_stock.get("stock_items", []):
        cat = item.get("category", "Other")
        if cat not in categories:
            categories[cat] = 0
        categories[cat] += item.get("quantity_kg", 0)
    
    return [{"category": k, "quantity": v} for k, v in categories.items()]

@api_router.get("/dashboard/sales-trend")
async def get_sales_trend(days: int = 7, user: dict = Depends(get_current_user)):
    dates = [(datetime.now(timezone.utc) - timedelta(days=i)).strftime("%Y-%m-%d") for i in range(days-1, -1, -1)]
    
    trend = []
    for date in dates:
        sales = await db.sales.find({
            "date": date,
            "status": SalesStatus.POSTED.value
        }, {"_id": 0}).to_list(1000)
        
        livestock = await db.livestock.find_one({"date": date}, {"_id": 0})
        
        trend.append({
            "date": date,
            "sales": sum(s.get("total_sales", 0) for s in sales),
            "cattle_slaughtered": livestock.get("cattle_slaughtered", 0) if livestock else 0,
            "shoats_slaughtered": livestock.get("shoats_slaughtered", 0) if livestock else 0
        })
    
    return trend

@api_router.get("/dashboard/out-of-stock")
async def get_out_of_stock(user: dict = Depends(get_current_user)):
    """Get out-of-stock heatmap data"""
    latest_hq = await db.stock.find_one({"location": "hq"}, {"_id": 0}, sort=[("date", -1)])
    latest_kibarani = await db.stock.find_one({"location": "kibarani"}, {"_id": 0}, sort=[("date", -1)])
    
    categories = ["MOB Beef", "Loose Beef", "Lamb", "Shoats", "Offals", "KMC Products"]
    
    def get_status(stock_entry, category):
        if not stock_entry:
            return "unknown"
        items = stock_entry.get("stock_items", [])
        cat_items = [i for i in items if i.get("category") == category]
        if not cat_items:
            return "out_of_stock"
        total = sum(i.get("quantity_kg", 0) for i in cat_items)
        if total == 0:
            return "out_of_stock"
        elif total < 50:
            return "low"
        return "available"
    
    heatmap = []
    for cat in categories:
        heatmap.append({
            "category": cat,
            "hq": get_status(latest_hq, cat),
            "kibarani": get_status(latest_kibarani, cat)
        })
    
    return heatmap

# Audit Log Routes
@api_router.get("/audit-logs")
async def get_audit_logs(
    entity_type: Optional[str] = None,
    limit: int = 100,
    user: dict = Depends(require_roles(UserRole.SYSTEM_ADMIN, UserRole.MANAGEMENT))
):
    query = {}
    if entity_type:
        query["entity_type"] = entity_type
    
    logs = await db.audit_logs.find(query, {"_id": 0}).sort("timestamp", -1).limit(limit).to_list(limit)
    # Sanitize any remaining ObjectIds in changes
    return [sanitize_for_json(log) for log in logs]

# Seed Data Route (for initial setup)
@api_router.post("/seed-data")
async def seed_data():
    """Seed initial data for testing"""
    # Check if already seeded
    existing_shops = await db.shops.count_documents({})
    if existing_shops > 0:
        return {"message": "Data already seeded"}
    
    # Create shops
    shops = [
        {"id": str(uuid.uuid4()), "name": "Athi River", "location": "Athi River", "opening_date": "2024-01-01", "is_active": True, "created_at": datetime.now(timezone.utc).isoformat()},
        {"id": str(uuid.uuid4()), "name": "Landhies", "location": "Nairobi", "opening_date": "2024-01-01", "is_active": True, "created_at": datetime.now(timezone.utc).isoformat()},
        {"id": str(uuid.uuid4()), "name": "Makupa", "location": "Mombasa", "opening_date": "2024-01-01", "is_active": True, "created_at": datetime.now(timezone.utc).isoformat()},
        {"id": str(uuid.uuid4()), "name": "Kitengela", "location": "Kitengela", "opening_date": "2025-08-01", "is_active": True, "created_at": datetime.now(timezone.utc).isoformat()},
        {"id": str(uuid.uuid4()), "name": "Thika", "location": "Thika", "opening_date": "2025-09-01", "is_active": True, "created_at": datetime.now(timezone.utc).isoformat()},
        {"id": str(uuid.uuid4()), "name": "Utawala", "location": "Utawala", "opening_date": "2025-10-01", "is_active": True, "created_at": datetime.now(timezone.utc).isoformat()},
    ]
    await db.shops.insert_many(shops)
    
    # Create admin user
    admin_user = {
        "id": str(uuid.uuid4()),
        "email": "admin@kmc.go.ke",
        "password_hash": hash_password("admin123"),
        "full_name": "System Administrator",
        "role": UserRole.SYSTEM_ADMIN.value,
        "shop_id": None,
        "is_active": True,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "updated_at": datetime.now(timezone.utc).isoformat()
    }
    await db.users.insert_one(admin_user)
    
    # Create sample users
    sample_users = [
        {"email": "supervisor@kmc.go.ke", "full_name": "John Supervisor", "role": UserRole.SUPERVISOR.value, "shop_id": shops[0]["id"]},
        {"email": "admin.officer@kmc.go.ke", "full_name": "Jane Admin", "role": UserRole.ADMIN_OFFICER.value, "shop_id": None},
        {"email": "scm@kmc.go.ke", "full_name": "Peter SCM", "role": UserRole.SCM.value, "shop_id": None},
        {"email": "kibarani@kmc.go.ke", "full_name": "Mary Kibarani", "role": UserRole.KIBARANI_OFFICER.value, "shop_id": None},
        {"email": "management@kmc.go.ke", "full_name": "Director Management", "role": UserRole.MANAGEMENT.value, "shop_id": None},
    ]
    
    for u in sample_users:
        user_doc = {
            "id": str(uuid.uuid4()),
            "email": u["email"],
            "password_hash": hash_password("password123"),
            "full_name": u["full_name"],
            "role": u["role"],
            "shop_id": u["shop_id"],
            "is_active": True,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "updated_at": datetime.now(timezone.utc).isoformat()
        }
        await db.users.insert_one(user_doc)
    
    # Create sample sales data
    import random
    for shop in shops:
        for days_ago in range(14):
            date = (datetime.now(timezone.utc) - timedelta(days=days_ago)).strftime("%Y-%m-%d")
            kmc_total = random.uniform(150000, 800000)
            kenchic_total = random.uniform(0, 20000)
            total = kmc_total + kenchic_total
            
            # Random payment breakdown
            coop = total * random.uniform(0.4, 0.6)
            kcb = total * random.uniform(0.1, 0.2)
            ecitizen = total * random.uniform(0.1, 0.2)
            swipe = total - coop - kcb - ecitizen
            
            sales_doc = {
                "id": str(uuid.uuid4()),
                "shop_id": shop["id"],
                "date": date,
                "kmc_products_total": round(kmc_total, 2),
                "kenchic_total": round(kenchic_total, 2),
                "total_sales": round(total, 2),
                "payment_breakdown": {
                    "coop": round(coop, 2),
                    "kcb": round(kcb, 2),
                    "ecitizen": round(ecitizen, 2),
                    "swipe": round(swipe, 2)
                },
                "line_items": [],
                "status": SalesStatus.POSTED.value,
                "submitted_by": admin_user["id"],
                "submitted_at": datetime.now(timezone.utc).isoformat(),
                "posted_by": admin_user["id"],
                "posted_at": datetime.now(timezone.utc).isoformat(),
                "return_reason": None,
                "notes": None,
                "created_at": datetime.now(timezone.utc).isoformat(),
                "updated_at": datetime.now(timezone.utc).isoformat()
            }
            await db.sales.insert_one(sales_doc)
    
    # Create sample stock data
    stock_categories = [
        {"category": "MOB Beef", "product_name": "MOB Beef", "quantity_kg": random.uniform(5000, 20000)},
        {"category": "Loose Beef", "product_name": "Loose Beef", "quantity_kg": random.uniform(100, 500)},
        {"category": "Lamb", "product_name": "Whole Lamb", "quantity_kg": random.uniform(200, 1000)},
        {"category": "Shoats", "product_name": "Whole Goat", "quantity_kg": random.uniform(500, 2000)},
        {"category": "Offals", "subcategory": "Ox Matumbo Mix", "product_name": "Ox Matumbo Mix", "quantity_kg": random.uniform(100, 500)},
        {"category": "Offals", "subcategory": "Ox Liver", "product_name": "Ox Liver", "quantity_kg": random.uniform(50, 200)},
        {"category": "KMC Products", "subcategory": "VAP", "product_name": "Sausages", "quantity_kg": random.uniform(100, 500)},
        {"category": "KMC Products", "subcategory": "VAP", "product_name": "Smokies", "quantity_kg": random.uniform(100, 400)},
        {"category": "KMC Products", "subcategory": "VAP", "product_name": "Burgers", "quantity_kg": random.uniform(50, 200)},
        {"category": "KMC Products", "subcategory": "VAP", "product_name": "Meatballs", "quantity_kg": random.uniform(50, 150)},
        {"category": "KMC Products", "subcategory": "VAP", "product_name": "Corned Beef", "quantity_kg": random.uniform(100, 300)},
    ]
    
    for location in ["hq", "kibarani"]:
        for days_ago in range(7):
            date = (datetime.now(timezone.utc) - timedelta(days=days_ago)).strftime("%Y-%m-%d")
            stock_items = []
            for cat in stock_categories:
                item = cat.copy()
                item["quantity_kg"] = round(random.uniform(50, item["quantity_kg"]), 2)
                item["unit"] = "kg"
                item["status"] = "available" if item["quantity_kg"] > 100 else ("low" if item["quantity_kg"] > 0 else "out_of_stock")
                stock_items.append(item)
            
            stock_doc = {
                "id": str(uuid.uuid4()),
                "location": location,
                "date": date,
                "stock_items": stock_items,
                "status": StockStatus.POSTED.value,
                "submitted_by": admin_user["id"],
                "submitted_at": datetime.now(timezone.utc).isoformat(),
                "notes": None,
                "created_at": datetime.now(timezone.utc).isoformat(),
                "updated_at": datetime.now(timezone.utc).isoformat()
            }
            await db.stock.insert_one(stock_doc)
    
    # Create sample livestock data
    for days_ago in range(14):
        date = (datetime.now(timezone.utc) - timedelta(days=days_ago)).strftime("%Y-%m-%d")
        livestock_doc = {
            "id": str(uuid.uuid4()),
            "date": date,
            "cattle_received": random.randint(50, 200),
            "shoats_received": random.randint(100, 300),
            "cattle_slaughtered": random.randint(40, 150),
            "shoats_slaughtered": random.randint(80, 250),
            "notes": None,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "updated_at": datetime.now(timezone.utc).isoformat()
        }
        await db.livestock.insert_one(livestock_doc)
    
    # Create sample production ops
    ops_list = ["Slaughter", "Deboning", "Value Addition", "Milling", "Rendering", "Dispatch"]
    for days_ago in range(7):
        date = (datetime.now(timezone.utc) - timedelta(days=days_ago)).strftime("%Y-%m-%d")
        ops_doc = {
            "id": str(uuid.uuid4()),
            "date": date,
            "operations": random.sample(ops_list, random.randint(3, 6)),
            "slaughter_start_time": "05:30",
            "batches_completed": random.randint(2, 8),
            "dispatch_details": f"Dispatched to {random.randint(3, 6)} locations",
            "notes": None,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "updated_at": datetime.now(timezone.utc).isoformat()
        }
        await db.production_ops.insert_one(ops_doc)
    
    return {"message": "Sample data seeded successfully", "shops": len(shops), "users": len(sample_users) + 1}

# Health Check
@api_router.get("/health")
async def health_check():
    return {"status": "healthy", "timestamp": datetime.now(timezone.utc).isoformat()}

# Include the router
app.include_router(api_router)

# CORS Middleware
app.add_middleware(
    CORSMiddleware,
    allow_credentials=True,
    allow_origins=os.environ.get('CORS_ORIGINS', '*').split(','),
    allow_methods=["*"],
    allow_headers=["*"],
)

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

@app.on_event("shutdown")
async def shutdown_db_client():
    client.close()