from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
from typing import Literal

app = FastAPI(title="Jarvis AI API")

# ---- Load models ----
chat_gen   = pipeline("text-generation", model="microsoft/DialoGPT-small")
qa_gen     = pipeline("text2text-generation", model="google/flan-t5-small")  # smarter Q&A
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
sentiment  = pipeline("sentiment-analysis")
en2hi      = pipeline("translation", model="Helsinki-NLP/opus-mt-en-hi")
hi2en      = pipeline("translation", model="Helsinki-NLP/opus-mt-hi-en")

# ---- Schemas ----
class ChatIn(BaseModel):
    message: str
    max_new_tokens: int | None = 80

class SummarizeIn(BaseModel):
    text: str
    max_tokens: int | None = 130
    min_tokens: int | None = 30

class TranslateIn(BaseModel):
    text: str
    direction: Literal["en-hi", "hi-en"] = "en-hi"

class SentimentIn(BaseModel):
    text: str

@app.get("/")
def home():
    return {
        "message": "Jarvis AI API alive 🚀",
        "endpoints": ["/chat","/summarize","/translate","/sentiment"]
    }

# ---- CHAT endpoint with intent handling ----
@app.post("/chat")
def chat(inp: ChatIn):
    prompt = (inp.message or "").strip()
    if not prompt:
        return {"reply": "Please say something."}

    # Detect if it's a QUESTION
    is_question = prompt.endswith("?") or any(
        w in prompt.lower() for w in ["who","what","when","where","why","how"]
    )

    if is_question:
        # Use FLAN-T5 for Q&A / reasoning
        res = qa_gen(prompt, max_length=inp.max_new_tokens or 80)[0]["generated_text"]
        return {"reply": res.strip()}

    else:
        # Use DialoGPT for small talk
        out = chat_gen(
            prompt,
            max_length=min(len(prompt) + (inp.max_new_tokens or 60), 160),
            do_sample=True,
            temperature=0.8,
            top_p=0.92,
            top_k=50,
            num_return_sequences=1,
            pad_token_id=50256,
        )
        text = out[0]["generated_text"]
        continuation = text[len(prompt):].strip()
        short = continuation.split("\n")[0]
        if "." in short:
            short = short.split(".")[0].strip() + "."
        return {"reply": short or continuation or "I'm here!"}

# ---- SUMMARIZE ----
@app.post("/summarize")
def summarize(inp: SummarizeIn):
    txt = (inp.text or "").strip()
    if not txt:
        return {"summary": "No text provided."}
    result = summarizer(
        txt,
        max_length=inp.max_tokens or 130,
        min_length=inp.min_tokens or 30,
        do_sample=False,
    )
    return {"summary": result[0]["summary_text"]}

# ---- TRANSLATE ----
@app.post("/translate")
def translate(inp: TranslateIn):
    txt = (inp.text or "").strip()
    if not txt:
        return {"translation": "No text provided."}
    if inp.direction == "en-hi":
        res = en2hi(txt)
    else:
        res = hi2en(txt)
    return {"translation": res[0]["translation_text"]}

# ---- SENTIMENT ----
@app.post("/sentiment")
def sentiment_route(inp: SentimentIn):
    txt = (inp.text or "").strip()
    if not txt:
        return {"label": "NEUTRAL", "score": 0.0}
    res = sentiment(txt)[0]
    return {"label": res["label"], "score": res["score"]}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("myapi:app", host="0.0.0.0", port=8000, reload=True)