app.py

import os
import io
import uuid
import base64
import tempfile

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from PIL import Image
from src.rag_pipeline import KidneyStoneRAGPipeline
from src.llm_client import GeminiClient
from src.convert_to_pdf import md_text_to_pdf_bytes
from src.correction_service import CorrectionService
from src.correction_store import CorrectionStore
from src.models import CorrectionRequest, CorrectionResponse, VerifyRequest
from pydantic import BaseModel

app = FastAPI()

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

# Startup'ta bir kez yükle
pipeline = KidneyStoneRAGPipeline(
    detector_model_path="yolo26-seg_best.pt",
    kb_persist_directory="./chroma_db"
)
llm = GeminiClient()

correction_store = CorrectionStore()
correction_svc = CorrectionService(kb=pipeline.kb, store=correction_store)
correction_svc.replay_corrections_on_startup()

class RegenerateRequest(BaseModel):
    session_id: str
    features: dict

class TranslateRequest(BaseModel):
    report: str

class PDFRequest(BaseModel):
    report: str
    annotated_image: str | None = None

@app.get("/")
@app.get("/health")
def health():
    return {"status": "ok"}

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    # Geçici dosyaya kaydet (pipeline path istiyor)
    img_bytes = await file.read()

    # Gelen format ne olursa olsun (JPEG, DICOM-export PNG vb.) orijinal baytları koru
    suffix = os.path.splitext(file.filename or "")[1] or ".png"
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(img_bytes)
        tmp_path = tmp.name
    
    # Pipeline'ı çalıştır
    result = pipeline.process(tmp_path, output_dir="/tmp")
    
    # Annotated görüntüyü base64'e çevir
    with open(result["annotated_image_path"], "rb") as f:
        annotated_b64 = base64.b64encode(f.read()).decode()
    
    return {
        "session_id": str(uuid.uuid4()),
        "has_stone": result["features"].get("stone_detected", False),
        "detections": result["detection"],
        "features": result["features"],
        "annotated_image": annotated_b64,
        "report": result["report"],
        "retrieved_context": result["retrieved_context"],
    }

@app.post("/download-pdf")
async def download_pdf(body: PDFRequest):
    from fastapi import HTTPException
    try:
        pdf_bytes = md_text_to_pdf_bytes(body.report, body.annotated_image)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"PDF oluşturulamadı: {str(e)}")
    return StreamingResponse(
        io.BytesIO(pdf_bytes),
        media_type="application/pdf",
        headers={"Content-Disposition": "attachment; filename=kidney_stone_report.pdf"},
    )

@app.post("/translate")
async def translate(body: TranslateRequest):
    system_prompt = "You are a medical translator. Translate the following radiology report to Turkish. Keep medical terminology accurate and professional. Preserve the exact structure, headings, and formatting of the original report."

    translated = llm.generate(system_prompt=system_prompt, user_prompt=body.report)

    return {"report_tr": translated}


@app.post("/correct", response_model=CorrectionResponse)
async def correct_report(body: CorrectionRequest):
    """
    Doktor bir rapordaki hatayı bildirir.
    Sistem kök nedeni analiz eder ve ChromaDB'yi günceller.
    """
    try:
        return correction_svc.process_correction(body)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/verify-correction")
async def verify_correction(body: VerifyRequest):
    """
    Düzeltmenin etkili olup olmadığını kontrol eder.
    Aynı özelliklerle retrieval yeniden çalıştırılır.
    """
    try:
        return correction_svc.verify_correction(body.correction_id, body.features)
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/corrections")
async def list_corrections(limit: int = 50):
    """Tüm düzeltmelerin audit trail listesi."""
    return correction_store.list_corrections(limit)


@app.get("/corrections/{correction_id}")
async def get_correction(correction_id: str):
    """Belirli bir düzeltmenin detaylarını getir."""
    record = correction_store.get_correction(correction_id)
    if not record:
        raise HTTPException(status_code=404, detail="Correction not found")
    return record

@app.post("/regenerate-report")
async def regenerate_report(body: RegenerateRequest):
    from src.query_builder import build_all_queries
    from src.report_generator import build_report_prompt

    features = body.features
    queries = build_all_queries(features) if features.get("stone_detected") else []

    largest_mm = features.get("largest_stone_mm")
    size_filter = pipeline._stone_size_range(largest_mm) if largest_mm else None

    retrieved = pipeline.retrieve_context(queries, size_filter=size_filter) if queries else []

    system_prompt, user_prompt = build_report_prompt(features, retrieved)
    report = llm.generate(system_prompt=system_prompt, user_prompt=user_prompt, temperature=0.3)

    return {
        "session_id": body.session_id,
        "report": report,
        "retrieved_context": retrieved,
    }


@app.get("/admin/export-correction-store")
async def export_correction_store():
    """HF Spaces'taki güncel correction_store.db'yi indir."""
    import os
    from fastapi.responses import Response
    db_path = correction_store.db_path
    if not os.path.exists(db_path):
        raise HTTPException(status_code=404, detail="correction_store.db bulunamadı")
    with open(db_path, "rb") as f:
        db_bytes = f.read()
    return Response(
        content=db_bytes,
        media_type="application/octet-stream",
        headers={"Content-Disposition": "attachment; filename=correction_store.db"},
    )


#Test için daha sonra kaldırılıcak
@app.get("/admin/chunk-edits")
async def admin_chunk_edits():
    """correction_store içindeki tüm chunk değişikliklerini listeler."""
    import sqlite3
    conn = sqlite3.connect(correction_store.db_path)
    conn.row_factory = sqlite3.Row
    rows = conn.execute(
        "SELECT edit_id, correction_id, chroma_id, edit_type, "
        "original_content, new_content, timestamp "
        "FROM chunk_edits ORDER BY timestamp DESC"
    ).fetchall()
    conn.close()
    return {"chunk_edits": [dict(r) for r in rows], "total": len(rows)}