from fastapi import FastAPI
from pydantic import BaseModel
from identity import load_identity
from reasoning import reason_check
from reflection import reflect_notes
from rag import query_rag
from ingest_handler import handle_ingest
from refine import refine_answer


# ---------------------------------------------------
# Create the app BEFORE using it
# ---------------------------------------------------
app = FastAPI()

# ---------------------------------------------------
# Models
# ---------------------------------------------------
class ReasonRequest(BaseModel):
    prompt: str
    session_id: str | None = None

class IngestRequest(BaseModel):
    user: str
    assistant: str | None = None
    session_id: str | None = None

# ---------------------------------------------------
# Load identity
# ---------------------------------------------------
IDENTITY = load_identity()

# ---------------------------------------------------
# Routes MUST come after app = FastAPI()
# ---------------------------------------------------

@app.get("/health")
def health():
    return {
        "status": "ok",
        "identity_loaded": IDENTITY is not None
    }

@app.post("/ingest")
async def ingest(data: IngestRequest):
    await handle_ingest(data)
    return {"status": "ok"}

@app.post("/reason")
async def reason(data: ReasonRequest):
    user_prompt = data.prompt
    
    intake_summary = "recent summary"  

    identity_block = IDENTITY
    rag_block = query_rag(user_prompt)

    reflection_data = await reflect_notes(intake_summary, identity_block)
    notes = reflection_data.get("notes", [])

    draft = await reason_check(
        user_prompt,
        identity_block,
        rag_block,
        notes
    )
    # --- REFINE STEP ----------------------------------------------------
    refine_result = refine_answer(
    draft_output=draft,
    reflection_notes=notes,
    identity_block=identity_block,
    rag_block=rag_block,
)   

    final_output = refine_result["final_output"]

    return {
    "draft_output": draft,
    "reflection_notes": notes,
    "refined_output": final_output,
    "refine_meta": {
        "used_primary_backend": refine_result.get("used_primary_backend"),
        "fallback_used": refine_result.get("fallback_used")
    },
    "identity_used": identity_block is not None,
    "rag_used": rag_block is not None
}
