import uvicorn
import jwt
import os
import httpx 
import asyncio
import argparse
from fastapi import FastAPI, HTTPException, Header
from pydantic import BaseModel

app = FastAPI(title="Aris Node: LLM Specialist")

# --- CONFIG ---
REGISTRY_URL = os.getenv("ARIS_REGISTRY", "http://localhost:8000/register")
MY_DID = "did:aris:llm-node-01"
MY_ENDPOINT = "http://localhost:9006"
ARIS_PUBLIC_KEY = "s3cret_k3y_for_signing_tokens" 
OLLAMA_URL = "http://localhost:11434/api/generate"

# --- AUTO-REGISTER ON STARTUP ---
@app.on_event("startup")
async def start_heartbeat():
    async def heartbeat():
        while True:
            async with httpx.AsyncClient() as client:
                try:
                    await client.post(REGISTRY_URL, json={
                        "did": MY_DID,
                        "endpoint": MY_ENDPOINT,
                        "capabilities": ["ai.generate"]
                    })
                    print("💓 Heartbeat: Registry updated.")
                except Exception:
                    print("⚠️ Heartbeat: Registry unreachable...")
            await asyncio.sleep(30) # Check in every 30 seconds
            
    asyncio.create_task(heartbeat())

# --- MODELS & ENDPOINTS ---
class PromptRequest(BaseModel):
    model: str = "tinyllama"
    prompt: str

@app.post("/generate")
async def generate_text(job: PromptRequest, x_aris_token: str = Header(...)):
    try:
        # Verify the session token issued by the Registry
        payload = jwt.decode(x_aris_token, ARIS_PUBLIC_KEY, algorithms=["HS256"], audience=MY_DID)
    except jwt.InvalidTokenError:
        raise HTTPException(401, "Invalid or Expired Token")

    print(f"🧠 [LLM Node] Processing for {payload['sub']}...")

    async with httpx.AsyncClient() as client:
        resp = await client.post(OLLAMA_URL, json={
            "model": job.model, "prompt": job.prompt, "stream": False
        }, timeout=60.0)
        return {"result": resp.json().get("response", ""), "status": "success"}

def main():
    parser = argparse.ArgumentParser(description="Aris Worker Node")
    parser.add_argument("--port", type=int, default=9006, help="Port to run the agent on")
    parser.add_argument("--model", type=str, default="tinyllama", help="Ollama model to use")
    args = parser.parse_args()

    print(f"🚀 Starting Aris Node on port {args.port} using {args.model}...")
    uvicorn.run(app, host="0.0.0.0", port=args.port)
    
    
if __name__ == "__main__":
    main()