marzban_tg_bot/server.py

from fastapi import FastAPI, HTTPException, Request
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import uvicorn
from datetime import datetime
import logging
import json

from database import db
from config import PLANS, CONFIG
from marzban import marzban

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("server")

app = FastAPI()

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

@app.on_event("startup")
async def startup():
    await db.connect()
    logger.info("Database connected")

@app.get("/api/plans")
async def get_plans():
    plans_list = []
    for pid, p in PLANS.items():
        plans_list.append({
            "id": pid,
            **p
        })
    return plans_list

from aiogram.types import LabeledPrice
from pydantic import BaseModel

class BuyPlanRequest(BaseModel):
    user_id: int
    plan_id: str
    promo_code: str = None

class PromoCheckRequest(BaseModel):
    code: str

@app.post("/api/check-promo")
async def check_promo_code(req: PromoCheckRequest):
    promo = await db.get_promo_code(req.code)
    if not promo:
        return JSONResponse(status_code=404, content={"error": "Invalid or expired promo code"})
    
    return {
        "code": promo["code"],
        "discount": promo["discount"],
        "bonus_days": promo["bonus_days"],
        "is_unlimited": promo["is_unlimited"],
        "description": f"Discount {promo['discount']}%" + (f" + {promo['bonus_days']} Days" if promo['bonus_days'] else "")
    }

@app.post("/api/create-invoice")
async def create_invoice(req: BuyPlanRequest, request: Request):
    bot = getattr(request.app.state, "bot", None)
    if not bot:
        return JSONResponse(status_code=500, content={"error": "Bot instance not initialized"})

    plan = PLANS.get(req.plan_id)
    if not plan:
        return JSONResponse(status_code=404, content={"error": "Plan not found"})

    price = plan['price']

    # Validating Promo
    discount = 0
    if req.promo_code:
        promo = await db.get_promo_code(req.promo_code)
        if promo:
             discount = promo['discount']
    
    final_price = int(price * (100 - discount) / 100)
    if final_price < 1: final_price = 1

    try:
        invoice_link = await bot.create_invoice_link(
            title=f"Sub: {plan['name']}",
            description=f"{plan['data_limit']}GB / {plan['days']} days",
            payload=f"{req.plan_id}:{req.promo_code or ''}",
            provider_token="", # Empty for Stars
            currency="XTR",
            prices=[LabeledPrice(label=plan['name'], amount=final_price)]
        )
        return {"invoice_link": invoice_link}
    except Exception as e:
        logger.error(f"Error generating invoice: {e}")
        return JSONResponse(status_code=500, content={"error": str(e)})

@app.get("/api/user/{user_id}")
async def get_user_stats(user_id: int):
    user = await db.get_user(user_id)
    if not user:
        return JSONResponse(status_code=404, content={"error": "User not found"})

    sub_until = user['subscription_until']
    days_left = 0
    status = "Inactive"
    expire_str = "No active subscription"

    if sub_until:
        if isinstance(sub_until, str):
            try:
                sub_until = datetime.fromisoformat(sub_until)
            except:
                pass
        
        if isinstance(sub_until, datetime):
            expire_str = sub_until.strftime("%d.%m.%Y")
            if sub_until > datetime.now():
                delta = sub_until - datetime.now()
                days_left = delta.days
                status = "Active"
            else:
                status = "Expired"

    # Fetch detailed stats from Marzban
    sub_url = ""
    used_traffic = 0
    
    if user['marzban_username']:
        try:
            m_user = await marzban.get_user(user['marzban_username'])
            # Check for error in response
            if isinstance(m_user, dict) and not m_user.get('detail'):
                used_traffic = m_user.get('used_traffic', 0)
                sub_url = m_user.get('subscription_url', "")
                
                # Fix relative URL
                if sub_url and sub_url.startswith('/'):
                    base = CONFIG.get('BASE_URL') or CONFIG['MARZBAN_URL']
                    sub_url = f"{base.rstrip('/')}{sub_url}"
        except Exception as e:
            logger.error(f"Marzban fetch error: {e}")

    return {
        "status": status,
        "days_left": days_left,
        "expire_date": expire_str,
        "data_limit_gb": user['data_limit'] or 0,
        "used_traffic_gb": round(used_traffic / (1024**3), 2),
        "plan": "Premium",
        "subscription_url": sub_url,
        "username": user['username'],
        "marzban_username": user['marzban_username']
    }

# Serve Static Files (must be last)
app.mount("/", StaticFiles(directory="web_app/static", html=True), name="static")

if __name__ == "__main__":
    from config import CONFIG
    uvicorn.run(app, host="0.0.0.0", port=CONFIG["WEB_APP_PORT"])