import logging
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# ------------------------------
# Logging setup
# ------------------------------
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)
logger.info("🚀 Starting Privacy Audit AI Backend...")

# ------------------------------
# Load model and tokenizer
# ------------------------------
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")

# ------------------------------
# FastAPI app
# ------------------------------
app = FastAPI(title="Privacy Audit AI", version="1.0.0")

# ------------------------------
# Request model
# ------------------------------
class AnalyzeInput(BaseModel):
    os_apps: list[str] = []
    browser_extensions: list[str] = []
    account_apps: list[str] = []

# ------------------------------
# Utility functions
# ------------------------------
def generate_plain_text(input_text: str):
    """Generate plain-language explanation using the model."""
    logger.info(f"Generating explanation for input: {input_text[:60]}...")
    inputs = tokenizer(f"Explain privacy risks in plain language:\n{input_text}", return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=150)
    explanation = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return explanation

def explain_risks(os_apps, browser_exts, account_apps):
    summary_text = (
        f"OS Apps: {', '.join(os_apps[:10])}, "
        f"Browser Extensions: {', '.join(browser_exts[:5])}, "
        f"Connected Apps: {', '.join(account_apps)}"
    )
    return generate_plain_text(summary_text)

# ------------------------------
# Endpoints
# ------------------------------
@app.get("/ping")
def ping():
    return {"status": "ok", "message": "Backend is alive!"}

@app.post("/analyze")
def analyze(data: AnalyzeInput):
    explanation = explain_risks(data.os_apps, data.browser_extensions, data.account_apps)
    return {"plain_language": explanation}

@app.get("/audit")
def audit_mvp():
    findings = "App X has camera access, App Y has location access, Chrome has 5 extensions"
    explanation = generate_plain_text(findings)
    return {
        "findings": findings,
        "plain_language": explanation,
        "risk_level": "Medium"
    }

@app.post("/audit")
def audit_mvp_post(data: AnalyzeInput):
    findings_text = (
        f"Detected {len(data.os_apps)} installed apps, "
        f"{len(data.browser_extensions)} browser extensions, and "
        f"{len(data.account_apps)} connected account apps."
    )
    explanation = generate_plain_text(findings_text)
    return {
        "findings": findings_text,
        "plain_language": explanation,
        "risk_level": "Medium"
    }


uvicorn app:app --host 0.0.0.0 --port 7860