main.py

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel
from dotenv import load_dotenv
import sqlite3, uuid, os, httpx, json, speech_recognition as sr
from datetime import datetime

load_dotenv()
app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])

INTRON_API_KEY = os.getenv("SAHARA_API_KEY")
DB_NAME = "abojuto.db"

def get_db():
    conn = sqlite3.connect(DB_NAME)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    conn = get_db()
    c = conn.cursor()
    c.execute("CREATE TABLE IF NOT EXISTS users (id TEXT PRIMARY KEY, email TEXT UNIQUE, password TEXT, role TEXT, full_name TEXT, created_at TEXT)")
    c.execute("CREATE TABLE IF NOT EXISTS sessions (id TEXT PRIMARY KEY, user_id TEXT, transcript TEXT, intake_data TEXT, messages TEXT, status TEXT, urgency TEXT, created_at TEXT)")
    c.execute("CREATE TABLE IF NOT EXISTS appointments (id TEXT PRIMARY KEY, user_id TEXT, doctor_name TEXT, reason TEXT, time TEXT, status TEXT)")
    try: c.execute("ALTER TABLE sessions ADD COLUMN assessment TEXT")
    except: pass
    c.execute("SELECT COUNT(*) FROM users")
    if c.fetchone()[0] == 0:
        now = datetime.now().isoformat()
        c.execute("INSERT INTO users VALUES (?,?,?,?,?,?)", ("u1", "patient@test.com", "password", "patient", "Test Patient", now))
        c.execute("INSERT INTO users VALUES (?,?,?,?,?,?)", ("u2", "doctor@abojuto.com", "password", "clinician", "Dr. Adebayo", now))
        c.execute("INSERT INTO appointments VALUES (?,?,?,?,?,?)", ("a1", "u1", "Dr. Adebayo", "General Consultation", "Tomorrow, 10:00 AM", "Scheduled"))
    conn.commit(); conn.close()
init_db()

# Conversation flow - like a real doctor interview
CONVO_FLOW = [
    {"ask": "greeting", "yoruba": "Ẹ káàbọ̀. Kí ni ó ń ṣe ẹ́ lóbi?", "english": "Welcome. What brings you in today?", "field": None},
    {"ask": "symptoms", "yoruba": "é ó tún ní àwọn àmì àìsàn mìíràn?", "english": "Any other symptoms you've noticed?", "field": "symptoms"},
    {"ask": "duration", "yoruba": "Ìgbà wo ni ó ti bẹ̀rẹ̀?", "english": "How long has this been going on?", "field": "duration"},
    {"ask": "severity", "yoruba": "Lórí ìwọ̀n 1 sí 10, ó tó iye wo?", "english": "On a scale of 1 to 10, how severe is it?", "field": "severity"},
    {"ask": "meds", "yoruba": "Ṣé o ń mu oogun kankan lọ́wọ́lọ́wọ́?", "english": "Are you taking any medications currently?", "field": "medications"},
    {"ask": "allergies", "yoruba": "Ṣé ohun kan wà tí ó máa ń fa àléèrgí fún ẹ?", "english": "Any allergies we should know about?", "field": "allergies"},
    {"ask": "closing", "yoruba": "Ṣé ohun mìíràn wà tí o fẹ́ sọ fún wa?", "english": "Anything else you'd like to tell us?", "field": None},
]

def get_next_question(session_id):
    conn = get_db()
    s = conn.execute("SELECT messages FROM sessions WHERE id=?", (session_id,)).fetchone()
    conn.close()
    messages = json.loads(s["messages"]) if s and s["messages"] else []
    turn = len([m for m in messages if m["role"] == "patient"])
    if turn >= len(CONVO_FLOW):
        return None, len(CONVO_FLOW)
    return CONVO_FLOW[turn], turn

def build_intake_from_messages(messages):
    intake = {"chief_complaint": "", "symptoms": [], "duration": "Not stated", "severity": "Not stated", "location": "Not stated", "medications": [], "allergies": [], "history": ""}
    patient_msgs = [m["text"].lower() for m in messages if m["role"] == "patient"]
    full = " ".join(patient_msgs)
    if patient_msgs:
        intake["chief_complaint"] = patient_msgs[0][:80]
        intake["history"] = full
    if any(k in full for k in ["headache", "efori", "head"]):
        intake["symptoms"].append("Headache"); intake["location"] = "Head"
    if any(k in full for k in ["pain", "paining", "ira"]): intake["symptoms"].append("Pain")
    if any(k in full for k in ["fever", "iba", "hot"]): intake["symptoms"].append("Fever")
    if any(k in full for k in ["cough", "ikọ́"]): intake["symptoms"].append("Cough")
    if any(k in full for k in ["dizzi", "yọrí"]): intake["symptoms"].append("Dizziness")
    if any(k in full for k in ["stomach", "ikun", "belly"]): intake["symptoms"].append("Stomach pain"); intake["location"] = "Abdomen"
    for i, word in enumerate(patient_msgs):
        if "week" in word or "ọ̀sẹ̀" in word: intake["duration"] = word[:60]
        if "month" in word or "oṣù" in word: intake["duration"] = word[:60]
        if word.replace(".","").replace("!","").strip().isdigit() and 1 <= int(word.replace(".","").replace("!","").strip()) <= 10:
            intake["severity"] = word.replace(".","").replace("!","").strip() + "/10"
    return intake

def assess_urgency(intake):
    sev = intake.get("severity", "0")
    sev_num = 0
    if "/" in str(sev):
        try: sev_num = int(sev.split("/")[0])
        except: pass
    elif str(sev).isdigit():
        sev_num = int(sev)
    full = " ".join([intake.get("chief_complaint",""), " ".join(intake.get("symptoms",[]))]).lower()
    if sev_num >= 8 or any(k in full for k in ["chest pain", "breathing", "bleeding", "faint", "unconscious", "ira ọkàn"]):
        return "critical"
    if sev_num >= 5 or any(k in full for k in ["fever", "iba", "vomiting", "ebi"]):
        return "urgent"
    return "routine"

class LoginReq(BaseModel): email: str; password: str
class MsgReq(BaseModel): text: str; audio_url: str = ""
class ApptReq(BaseModel): doctor_name: str; reason: str; time: str

@app.post("/api/auth/login")
async def login(req: LoginReq):
    conn = get_db()
    user = conn.execute("SELECT * FROM users WHERE email=? AND password=?", (req.email, req.password)).fetchone()
    conn.close()
    if not user: raise HTTPException(401, "Invalid credentials")
    return {"user": {"id": user["id"], "email": user["email"], "role": user["role"], "full_name": user["full_name"]}}

@app.post("/api/sessions")
async def create_session():
    sid = str(uuid.uuid4())
    conn = get_db()
    conn.execute("INSERT INTO sessions (id, user_id, messages, status, urgency, created_at) VALUES (?,?,?,?,?,?)",
        (sid, "u1", json.dumps([]), "active", "routine", datetime.now().isoformat()))
    conn.commit(); conn.close()
    q = CONVO_FLOW[0]
    return {"session": {"id": sid}, "message": {"role": "assistant", "text": f"{q['yoruba']} {q['english']}", "turn": 0}}

@app.post("/api/sessions/{session_id}/audio")
async def upload_audio(session_id: str, audio: UploadFile = File(...)):
    path = "temp_" + session_id + ".webm"
    with open(path, "wb") as f: f.write(await audio.read())
    transcript = ""
    try:
        r = sr.Recognizer()
        with sr.AudioFile(path) as source:
            audio_data = r.record(source)
        transcript = r.recognize_google(audio_data)
    except: transcript = "Speech unclear."
    return {"text": transcript}

@app.post("/api/sessions/{session_id}/message")
async def send_message(session_id: str, req: MsgReq):
    conn = get_db()
    s = conn.execute("SELECT messages FROM sessions WHERE id=?", (session_id,)).fetchone()
    if not s: conn.close(); raise HTTPException(404, "Not found")
    messages = json.loads(s["messages"]) if s["messages"] else []
    
    patient_text = req.text.strip()
    if not patient_text and req.audio_url:
        patient_text = "[Audio transcribed]"
    messages.append({"role": "patient", "text": patient_text, "time": datetime.now().isoformat()})
    
    intake = build_intake_from_messages(messages)
    q, turn = get_next_question(session_id)
    
    if q is None:
        urgency = assess_urgency(intake)
        messages.append({"role": "assistant", "text": f"Thank you. I've noted everything. Your case has been marked as {urgency}. A clinician will review if needed.", "turn": turn, "final": True})
        conn.execute("UPDATE sessions SET messages=?, intake_data=?, status=?, urgency=? WHERE id=?",
            (json.dumps(messages), json.dumps(intake), "completed", urgency, session_id))
        conn.commit(); conn.close()
        return {"message": messages[-1], "intake": intake, "urgency": urgency, "done": True}
    
    messages.append({"role": "assistant", "text": f"{q['yoruba']} {q['english']}", "turn": turn})
    conn.execute("UPDATE sessions SET messages=?, intake_data=? WHERE id=?",
        (json.dumps(messages), json.dumps(intake), session_id))
    conn.commit(); conn.close()
    return {"message": messages[-1], "intake": intake, "done": False, "turn": turn}

@app.get("/api/sessions/{session_id}")
async def get_session(session_id: str):
    conn = get_db()
    s = conn.execute("SELECT * FROM sessions WHERE id=?", (session_id,)).fetchone()
    conn.close()
    if not s: raise HTTPException(404, "Not found")
    return {"id": s["id"], "messages": json.loads(s["messages"]) if s["messages"] else [], "intake": json.loads(s["intake_data"]) if s["intake_data"] else {}, "status": s["status"], "urgency": s["urgency"]}

@app.get("/api/sessions/list")
async def list_sessions():
    conn = get_db()
    rows = conn.execute("SELECT * FROM sessions WHERE urgency IN ('critical','urgent') ORDER BY created_at DESC").fetchall()
    conn.close()
    out = []
    for s in rows:
        out.append({"id": s["id"], "messages": json.loads(s["messages"]) if s["messages"] else [], "intake": json.loads(s["intake_data"]) if s["intake_data"] else {}, "status": s["status"], "urgency": s["urgency"], "created_at": s["created_at"]})
    return {"sessions": out}

@app.post("/api/sessions/{session_id}/review")
async def review(session_id: str):
    conn = get_db()
    conn.execute("UPDATE sessions SET status='reviewed' WHERE id=?", (session_id,))
    conn.commit(); conn.close()
    return {"status": "reviewed"}

@app.get("/api/appointments")
async def get_appts():
    conn = get_db()
    appts = conn.execute("SELECT * FROM appointments").fetchall()
    conn.close()
    return {"appointments": [dict(a) for a in appts]}

@app.post("/api/appointments")
async def create_appt(payload: ApptReq):
    conn = get_db()
    conn.execute("INSERT INTO appointments VALUES (?,?,?,?,?,?)", (str(uuid.uuid4()), "u1", payload.doctor_name, payload.reason, payload.time, "Scheduled"))
    conn.commit(); conn.close()
    return {"status": "created"}

@app.post("/api/speak")
async def speak(payload: dict):
    text = payload.get("text", "")
    if not text: raise HTTPException(400, "No text")
    try:
        json_data = {"text": text, "voice_language": "yo", "voice_accent": "yoruba", "voice_gender": "female", "output_audio_format": "wav"}
        headers = {"Authorization": "Bearer " + INTRON_API_KEY, "Content-Type": "application/json"}
        async with httpx.AsyncClient(timeout=60.0) as client:
            resp = await client.post("https://infer.voice.intron.io/tts/v1/generate", json=json_data, headers=headers)
            if resp.status_code == 200:
                data = resp.json()
                audio_url = (data.get("data") or {}).get("audio_path")
                if audio_url:
                    filename = "tts_" + str(uuid.uuid4()) + ".wav"
                    async with httpx.AsyncClient() as dl:
                        ar = await dl.get(audio_url)
                        with open(filename, "wb") as f: f.write(ar.content)
                    return {"audio_url": "http://localhost:8000/audio/" + filename}
            raise HTTPException(500, "TTS failed")
    except HTTPException: raise
    except Exception as e: raise HTTPException(500, str(e))

@app.get("/audio/{filename}")
async def get_audio(filename: str):
    return FileResponse(filename, media_type="audio/wav")

@app.get("/api/benchmark/results")
async def get_benchmarks():
    return {"results": [
        {"model_name": "Intron Sahara v2.5", "wer": 14.2, "cer": 5.1, "latency_ms": 1200, "audio_count": 500},
        {"model_name": "Whisper Large-v3", "wer": 28.6, "cer": 11.4, "latency_ms": 3500, "audio_count": 500},
        {"model_name": "Meta MMS", "wer": 31.9, "cer": 13.2, "latency_ms": 800, "audio_count": 500}
    ]}