import time
import json
import jwt  # pip install pyjwt
from fastapi import FastAPI, HTTPException, Header
from pydantic import BaseModel
from typing import List, Dict, Optional

# --- CONFIG ---
ARIS_PRIVATE_KEY = "s3cret_k3y_for_signing_tokens"
MVP_PRICE_INR = 20.00
LEDGER_FILE = "registry/accounts.json"

app = FastAPI(title="Aris Registry (Paid)", version="0.2")

# --- DATA STORES ---
service_registry: Dict[str, List[dict]] = {}
ledger: Dict[str, dict] = {} 

# --- HELPER FUNCTIONS ---
def load_accounts():
    try:
        with open(LEDGER_FILE, "r") as f:
            return json.load(f)
    except FileNotFoundError:
        return {}

def save_accounts(data):
    with open(LEDGER_FILE, "w") as f:
        json.dump(data, f, indent=4)
        
def validate_api_key(api_key: str):
    # Load your ledger
    with open("registry/accounts.json", "r") as f:
        accounts = json.load(f)
    
    # Check 1: Does key exist?
    if api_key not in accounts:
        raise HTTPException(status_code=401, detail="Invalid API Key")
    
    # Check 2: Is the account enabled?
    user = accounts[api_key]
    if not user.get("active", False):
        raise HTTPException(status_code=403, detail="Account suspended")
        
    # Check 3: Does it have money?
    if user.get("balance", 0) <= 0:
        raise HTTPException(status_code=402, detail="Insufficient funds")
        
    return user

# --- MODELS ---
class AgentRegistration(BaseModel):
    did: str
    endpoint: str
    capabilities: List[str]

class SessionRequest(BaseModel):
    payer_did: str
    target_did: str
    capability: str

# --- ENDPOINTS ---

@app.post("/register")
async def register_agent(agent: AgentRegistration):
    for cap in agent.capabilities:
        if cap not in service_registry:
            service_registry[cap] = []
        # Update existing agent or add new
        service_registry[cap] = [a for a in service_registry[cap] if a['did'] != agent.did]
        service_registry[cap].append(agent.dict())
    
    print(f"✅ [Registry] Registered {agent.did}")
    return {"status": "registered"}

@app.get("/discover")
async def discover(capability: str):
    agents = service_registry.get(capability, [])
    if not agents:
        raise HTTPException(404, "No agents found")
    return {"agents": agents}

@app.post("/handshake")
async def handshake(req: SessionRequest, x_api_key: Optional[str] = Header(None)):
    """
    💰 PAYMENT LOGIC:
    1. Check Key
    2. Check Balance
    3. Deduct Fee
    4. Issue Token
    """
    # 1. AUTH
    if not x_api_key:
        raise HTTPException(401, "Missing API Key")

    accounts = load_accounts()
    if x_api_key not in accounts:
        raise HTTPException(403, "Invalid API Key")

    user_account = accounts[x_api_key]

    if not user_account["active"]:
        raise HTTPException(403, "Account Suspended")

    # 2. CHECK BALANCE
    if user_account["balance"] < MVP_PRICE_INR:
        raise HTTPException(402, "Insufficient Balance. Please Top Up.")

    # 3. CHARGE
    user_account["balance"] -= MVP_PRICE_INR
    save_accounts(accounts) # Save to disk
    
    print(f"💰 [Billing] Charged {x_api_key[:10]}... | New Balance: ₹{user_account['balance']}")

    # 4. ISSUE TOKEN
    txn_id = f"txn_{int(time.time())}"
    payload = {
        "iss": "aris-registry",
        "sub": req.payer_did,
        "aud": req.target_did,
        "scope": req.capability,
        "txn_id": txn_id,
        "exp": time.time() + 300 # 5 mins
    }
    
    token = jwt.encode(payload, ARIS_PRIVATE_KEY, algorithm="HS256")

    return {
        "session_token": token,
        "txn_id": txn_id,
        "remaining_balance": user_account["balance"]
    }
    
class TopUpRequest(BaseModel):
    api_key: str
    amount: float

@app.post("/admin/topup")
async def top_up(req: TopUpRequest):
    accounts = load_accounts()
    if req.api_key not in accounts:
        raise HTTPException(404, "Key not found")
    
    accounts[req.api_key]["balance"] += req.amount
    save_accounts(accounts)
    
    return {
        "status": "success", 
        "new_balance": accounts[req.api_key]["balance"]
    }
    
if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)
    

def start():
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

if __name__ == "__main__":
    start()