# app.py - FastAPI Web UI for HOLAS Defender
import os
import asyncio
import requests
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel
from typing import List, Optional

# ─────────────────────────────────────────────────────
# ⚙️ Configuration
# ─────────────────────────────────────────────────────
APP_NAME = "🐨 Holas Defender Ultimate V13"
MODEL = os.getenv("MODEL_NAME", "llama3.2")
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://localhost:11434/api/generate")
HF_API_URL = os.getenv("HF_API_URL", "https://api-inference.huggingface.co/models/meta-llama/Llama-3.2-3B-Instruct")
HF_TOKEN = os.getenv("HF_TOKEN", "")
USE_OLLAMA = os.getenv("USE_OLLAMA", "false").lower() == "true"

SYSTEM_PROMPT = """You are HOLAS Defender, an advanced cybersecurity assistant.
🔹 Help identify threats, phishing, vulnerabilities
🔹 Provide clear, actionable security advice
🔹 Explain concepts simply; stay ethical & professional
🔹 Never assist with malicious activities
Respond concisely in markdown. Start with 🐨"""

app = FastAPI(title=APP_NAME)

# Mount static files & templates (optional)
# app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")

# ─────────────────────────────────────────────────────
# 🧠 Inference Functions
# ─────────────────────────────────────────────────────
def query_ollama(prompt: str, history: List[dict]) -> str:
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for msg in history:
        messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": prompt})
    
    payload = {
        "model": MODEL,
        "prompt": prompt,
        "system": SYSTEM_PROMPT,
        "stream": False,
        "options": {"temperature": 0.3, "num_predict": 512}
    }
    
    try:
        resp = requests.post(OLLAMA_URL, json=payload, timeout=60)
        resp.raise_for_status()
        return resp.json().get("response", "⚠️ No response")
    except Exception as e:
        return f"⚠️ Ollama error: {str(e)}"

def query_huggingface(prompt: str, history: List[dict]) -> str:
    headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
    full_prompt = f"<|system|>\n{SYSTEM_PROMPT}<|user|>\n{prompt}<|assistant|>\n"
    
    payload = {
        "inputs": full_prompt,
        "parameters": {"max_new_tokens": 512, "temperature": 0.3, "return_full_text": False}
    }
    
    try:
        resp = requests.post(HF_API_URL, headers=headers, json=payload, timeout=90)
        if resp.status_code == 503:
            return "🔄 Model loading... Try again in ~30 seconds."
        resp.raise_for_status()
        result = resp.json()
        if isinstance(result, list):
            return result[0].get("generated_text", "⚠️ Empty response")
        return str(result)
    except Exception as e:
        return f"⚠️ HF API error: {str(e)}"

def generate_response(prompt: str, history: List[dict]) -> str:
    if USE_OLLAMA:
        return query_ollama(prompt, history)
    else:
        return query_huggingface(prompt, history)

# ─────────────────────────────────────────────────────
# 🌐 API Routes
# ─────────────────────────────────────────────────────
class ChatMessage(BaseModel):
    role: str  # "user" or "assistant"
    content: str

class ChatRequest(BaseModel):
    message: str
    history: Optional[List[ChatMessage]] = []

@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
    """Serve the main chat UI"""
    return templates.TemplateResponse("index.html", {
        "request": request,
        "title": APP_NAME,
        "emoji": "🐨"
    })

@app.post("/api/chat")
async def chat_endpoint(req: ChatRequest):
    """Chat API endpoint"""
    if not req.message.strip():
        raise HTTPException(status_code=400, detail="Message cannot be empty")
    
    # Convert history to dict format
    history_dicts = [{"role": m.role, "content": m.content} for m in (req.history or [])]
    
    # Generate response (run in thread to avoid blocking)
    loop = asyncio.get_event_loop()
    response = await loop.run_in_executor(
        None, 
        lambda: generate_response(req.message, history_dicts)
    )
    
    return {"response": response, "model": MODEL}

@app.get("/health")
async def health_check():
    """Health check for HF Spaces"""
    return {"status": "healthy", "app": APP_NAME}

@app.get("/api/models")
async def list_models():
    """List available models (for debugging)"""
    return {
        "current": MODEL,
        "ollama_enabled": USE_OLLAMA,
        "hf_token_set": bool(HF_TOKEN)
    }

# ─────────────────────────────────────────────────────
# 🚀 Launch
# ─────────────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn
    port = int(os.getenv("PORT", 7860))
    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=port,
        log_level="info",
        reload=False
    )