app.py

"""
PensionSight — FastAPI Backend
Run: uvicorn main:app --reload --port 8000
"""

import os
import json
import re
from typing import Optional, List
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import google.generativeai as genai
import uvicorn
from dotenv import load_dotenv

load_dotenv()

from pension_engine import (
    SubscriberProfile,
    PensionCalculationEngine,
    NPS_SCHEMES,
    SCENARIO_RETURNS,
)

# ── App Setup ─────────────────────────────────────────────────────────────────
app = FastAPI(title="PensionSight API", version="1.0.0")

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

engine = PensionCalculationEngine()

API_KEY = os.getenv("GEMINI_API_KEY")
if API_KEY:
    genai.configure(api_key=API_KEY)

# ── Pydantic Models ───────────────────────────────────────────────────────────
class ProfileInput(BaseModel):
    current_age: int
    retirement_age: int
    monthly_contribution: float
    existing_corpus: float = 0
    sector: str = "non_government"
    scheme: str = "LC50"
    annuity_percent: float = 0.40
    annuity_rate: float = 0.06
    annual_contribution_increase: float = 0
    employer_contribution: float = 0
    is_gig_worker: bool = False
    monthly_incomes: List[float] = []
    deferral_age: Optional[int] = None
    desired_monthly_pension: Optional[float] = None

class ReversePlanInput(BaseModel):
    desired_monthly_pension: float
    current_age: int
    retirement_age: int
    scheme: str = "LC50"
    scenario: str = "realistic"
    annuity_percent: float = 0.40
    existing_corpus: float = 0
    annual_step_up: float = 0

class ChatMessage(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[ChatMessage]
    session_id: Optional[str] = None

# ── Helpers ───────────────────────────────────────────────────────────────────
def fmt_inr(amount: float) -> str:
    if amount >= 1e7:
        return f"₹{amount/1e7:.2f} Cr"
    elif amount >= 1e5:
        return f"₹{amount/1e5:.2f} Lakh"
    return f"₹{amount:,.0f}"

def safe_float(v, d=0.0): return float(v) if v is not None else d
def safe_int(v, d=0):     return int(v)   if v is not None else d

def profile_from_input(p: ProfileInput) -> SubscriberProfile:
    return SubscriberProfile(
        current_age=p.current_age,
        retirement_age=p.retirement_age,
        monthly_contribution=p.monthly_contribution,
        existing_corpus=p.existing_corpus,
        sector=p.sector,
        scheme=p.scheme,
        annuity_percent=p.annuity_percent,
        annuity_rate=p.annuity_rate,
        annual_contribution_increase=p.annual_contribution_increase,
        employer_contribution=p.employer_contribution,
        is_gig_worker=p.is_gig_worker,
        monthly_incomes=p.monthly_incomes,
        deferral_age=p.deferral_age,
        desired_monthly_pension=p.desired_monthly_pension,
    )

def scenario_to_dict(r, profile: SubscriberProfile) -> dict:
    ap  = int(profile.annuity_percent * 100)
    lp  = 100 - ap
    return {
        "scenario":           r.scenario,
        "label":              r.label,
        "blended_return":     r.blended_return,
        "projected_corpus":   r.projected_corpus,
        "real_corpus":        r.real_corpus,
        "lumpsum":            r.lumpsum_withdrawal,
        "lumpsum_pct":        lp,
        "annuity_corpus":     r.annuity_corpus,
        "annuity_pct":        ap,
        "monthly_pension":    r.monthly_pension,
        "real_monthly_pension": r.real_monthly_pension,
        "total_contributions": r.total_contributions,
        "wealth_gained":      r.wealth_gained,
        "investment_years":   r.investment_years,
    }

# ── Routes ────────────────────────────────────────────────────────────────────
@app.get("/")
def root():
    return {"status": "PensionSight API running", "version": "1.0.0"}

@app.post("/api/project")
def project(profile_input: ProfileInput):
    try:
        profile = profile_from_input(profile_input)
        results = engine.project_all_scenarios(profile)
        nudges  = engine.generate_nudges(profile)
        return {
            "success": True,
            "profile": {
                "current_age": profile.current_age,
                "retirement_age": profile.retirement_age,
                "investment_years": profile.investment_years,
                "monthly_sip": profile.total_monthly_contribution,
                "scheme": NPS_SCHEMES[profile.scheme]["name"],
                "scheme_key": profile.scheme,
                "annual_step_up": profile.annual_contribution_increase,
                "annuity_percent": profile.annuity_percent,
            },
            "scenarios": {
                key: scenario_to_dict(r, profile)
                for key, r in results.items()
            },
            "nudges": [
                {
                    "type":     n.nudge_type,
                    "message":  n.message,
                    "impact":   n.impact_rupees,
                    "impact_fmt": fmt_inr(n.impact_rupees),
                    "current":  n.current_value,
                    "improved": n.improved_value,
                }
                for n in nudges
            ],
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/api/reverse-plan")
def reverse_plan(inp: ReversePlanInput):
    try:
        result = engine.reverse_plan(
            desired_monthly_pension=inp.desired_monthly_pension,
            current_age=inp.current_age,
            retirement_age=inp.retirement_age,
            scheme=inp.scheme,
            scenario=inp.scenario,
            annuity_percent=inp.annuity_percent,
            existing_corpus=inp.existing_corpus,
            annual_step_up=inp.annual_step_up,
        )
        return {
            "success": True,
            "desired_monthly_pension": result.desired_monthly_pension,
            "desired_fmt":       fmt_inr(result.desired_monthly_pension),
            "required_corpus":   result.required_corpus,
            "required_corpus_fmt": fmt_inr(result.required_corpus),
            "required_monthly_sip": result.required_monthly_sip,
            "required_sip_fmt":  fmt_inr(result.required_monthly_sip),
            "scenario":          result.scenario,
            "investment_years":  result.investment_years,
            "current_age":       result.current_age,
            "retirement_age":    result.retirement_age,
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/api/schemes")
def get_schemes():
    return {
        key: {
            "name": val["name"],
            "max_equity": int(val["max_equity"] * 100),
        }
        for key, val in NPS_SCHEMES.items()
    }


# ── Gemini Chat ───────────────────────────────────────────────────────────────
SYSTEM_INSTRUCTION = """
You are PensionSight, India's smartest NPS retirement planning assistant for PFRDA Innovate4NPS 2026.

YOUR MOST IMPORTANT RULE:
You MUST NEVER calculate, estimate, or guess any financial numbers (corpus, pension, SIP amounts).
When you have all required data, emit a tool_call JSON block. The engine calculates — you explain.

TOOL CALL FORMAT:
```tool_call
{
  "tool": "project",
  "profile": {
    "current_age": 20,
    "retirement_age": 60,
    "monthly_contribution": 20000,
    "existing_corpus": 0,
    "sector": "non_government",
    "scheme": "LC75",
    "annuity_percent": 0.40,
    "annuity_rate": 0.06,
    "annual_contribution_increase": 5,
    "employer_contribution": 0,
    "is_gig_worker": true,
    "monthly_incomes": [80000, 80000, 80000],
    "deferral_age": null,
    "desired_monthly_pension": null
  }
}
```

For reverse planning:
```tool_call
{
  "tool": "reverse_plan",
  "desired_monthly_pension": 40000,
  "current_age": 28,
  "retirement_age": 60,
  "scheme": "LC50",
  "annuity_percent": 0.40,
  "existing_corpus": 0,
  "annual_step_up": 5
}
```

DATA TO COLLECT (1-2 questions at a time):
1. current_age (18–74)
2. retirement_age (60–75)
3. monthly_contribution (≥500)
4. sector (government/non_government; gig workers = non_government)
5. existing_corpus (0 if new)
6. scheme (LC25/LC50/LC75; recommend LC75 if age<35)
7. annuity_percent (0.40–1.0; min 40%)
8. annual_contribution_increase (0–20%; suggest 5%)
9. employer_contribution (0 for gig/self-employed)
10. is_gig_worker (true/false); if true → monthly_incomes list

NPS 2025 RULES:
- Min retirement: 60, max deferral: 75
- Min annuity: 40% (max lumpsum 60%, tax-free)
- Government: employer 14% of basic (Budget 2025)
- Schemes: LC25 (low), LC50 (moderate), LC75 (high equity), AGGRESSIVE, ACTIVE

TONE: Warm, concise, no jargon. Max 2 questions per reply. Use ₹ symbol.
After projection results, offer: (1) reverse planning, (2) nudges, (3) annuity split change.
"""

GEMINI_MODEL = None
if API_KEY:
    GEMINI_MODEL = genai.GenerativeModel(
        model_name="gemini-2.5-flash-lite",
        system_instruction=SYSTEM_INSTRUCTION
    )

def extract_tool_call(text: str):
    for pattern in [
        r"```tool_call\s*([\s\S]*?)```",
        r"```json\s*([\s\S]*?)```",
        r"```\s*(\{[\s\S]*?\})\s*```",
    ]:
        m = re.search(pattern, text)
        if m:
            try:
                d = json.loads(m.group(1).strip())
                if "tool" in d:
                    return d
            except Exception:
                pass
    return None

def strip_code_blocks(text: str) -> str:
    return re.sub(r"```[\w]*\s*[\s\S]*?```", "", text).strip()

def run_tool(tj: dict) -> dict:
    tool = tj.get("tool")
    if tool == "project":
        p = tj["profile"]
        inp = ProfileInput(
            current_age=safe_int(p.get("current_age"), 30),
            retirement_age=safe_int(p.get("retirement_age"), 60),
            monthly_contribution=safe_float(p.get("monthly_contribution"), 1000),
            existing_corpus=safe_float(p.get("existing_corpus")),
            sector=p.get("sector", "non_government"),
            scheme=p.get("scheme", "LC50"),
            annuity_percent=safe_float(p.get("annuity_percent"), 0.40),
            annuity_rate=safe_float(p.get("annuity_rate"), 0.06),
            annual_contribution_increase=safe_float(p.get("annual_contribution_increase")),
            employer_contribution=safe_float(p.get("employer_contribution")),
            is_gig_worker=bool(p.get("is_gig_worker", False)),
            monthly_incomes=list(p.get("monthly_incomes") or []),
            deferral_age=safe_int(p["deferral_age"]) if p.get("deferral_age") else None,
        )
        return {"type": "project", "data": project(inp)}
    elif tool == "reverse_plan":
        inp = ReversePlanInput(
            desired_monthly_pension=safe_float(tj["desired_monthly_pension"]),
            current_age=safe_int(tj["current_age"]),
            retirement_age=safe_int(tj["retirement_age"]),
            scheme=tj.get("scheme", "LC50"),
            annuity_percent=safe_float(tj.get("annuity_percent"), 0.40),
            existing_corpus=safe_float(tj.get("existing_corpus")),
            annual_step_up=safe_float(tj.get("annual_step_up")),
        )
        return {"type": "reverse_plan", "data": reverse_plan(inp)}
    return {"type": "error", "data": {"message": f"Unknown tool: {tool}"}}


@app.post("/api/chat")
def chat(req: ChatRequest):
    if not GEMINI_MODEL:
        raise HTTPException(status_code=500, detail="GEMINI_API_KEY not configured")

    try:
        # Build Gemini history (exclude last user message)
        history = []
        messages = req.messages
        for msg in messages[:-1]:
            history.append({"role": msg.role if msg.role != "assistant" else "model",
                           "parts": [msg.content]})

        chat_session = GEMINI_MODEL.start_chat(history=history)
        last_msg     = messages[-1].content
        response     = chat_session.send_message(last_msg)
        bot_text     = response.text

        tool_json    = extract_tool_call(bot_text)
        display_text = strip_code_blocks(bot_text)
        tool_result  = None
        follow_up    = None

        if tool_json:
            tool_result = run_tool(tool_json)
            # Get Gemini to interpret the result
            result_summary = json.dumps(tool_result["data"], indent=2)[:3000]
            fu_resp  = chat_session.send_message(
                f"[TOOL RESULT]\n{result_summary}\n\n"
                "Summarise the realistic scenario for the user in 3-4 sentences. "
                "Then offer: (1) reverse planning, (2) nudges, (3) change annuity split."
            )
            follow_up = fu_resp.text

        return {
            "success":      True,
            "text":         display_text,
            "tool_result":  tool_result,
            "follow_up":    follow_up,
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    
if __name__ == "__main__":
    uvicorn.run("app:app", host="0.0.0.0", port=7860)