from typing import List, Optional
import uuid
import os
import shutil
from fastapi import FastAPI, File, UploadFile, Body, Form, BackgroundTasks
from ultralytics import YOLO
from dotenv import load_dotenv
import io
from PIL import Image

from src.infrastructure.vision_service import VertexVisionService
from src.infrastructure.repository import MongoPalmOilRepository
from src.application.analyze_bunch import AnalyzeBunchUseCase, AnalyzeBatchUseCase, SearchSimilarUseCase

# Load environment variables
load_dotenv()

# Set Google Cloud credentials globally
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "gemini-embedding-service-key.json"

app = FastAPI(title="Palm Oil Ripeness Service (DDD)")

# Initialize YOLO model
yolo_model = YOLO('best.pt')

# Global state for the confidence threshold
current_conf = 0.25

# Initialize DDD Components
vision_service = VertexVisionService(
    project_id=os.getenv("PROJECT_ID", "your-project-id"),
    location=os.getenv("LOCATION", "us-central1")
)
repo = MongoPalmOilRepository(
    uri=os.getenv("MONGO_URI"),
    db_name=os.getenv("DB_NAME", "palm_oil_db"),
    collection_name=os.getenv("COLLECTION_NAME", "ffb_records")
)
analyze_use_case = AnalyzeBunchUseCase(vision_service, repo)
analyze_batch_use_case = AnalyzeBatchUseCase(vision_service, repo)
search_use_case = SearchSimilarUseCase(vision_service, repo)

@app.get("/get_confidence")
# ... (rest of the code remains same until analyze)
async def get_confidence():
    """Returns the current confidence threshold used by the model."""
    return {
        "status": "success",
        "current_confidence": current_conf,
        "model_version": "best.pt"
    }

@app.post("/set_confidence")
async def set_confidence(threshold: float = Body(..., embed=True)):
    """Updates the confidence threshold globally."""
    global current_conf
    if 0.0 <= threshold <= 1.0:
        current_conf = threshold
        return {"status": "success", "new_confidence": current_conf}
    else:
        return {"status": "error", "message": "Threshold must be between 0.0 and 1.0"}


@app.post("/analyze")
async def analyze_only(file: UploadFile = File(...)):
    """Local YOLO detection only. Guaranteed to work without Billing."""
    image_bytes = await file.read()
    img = Image.open(io.BytesIO(image_bytes))
    results = yolo_model(img, conf=current_conf)
    
    detections = []
    for r in results:
        for box in r.boxes:
            detections.append({
                "class": yolo_model.names[int(box.cls)],
                "confidence": round(float(box.conf), 2),
                "box": box.xyxy.tolist()[0]
            })
            
    return {
        "status": "success", 
        "current_threshold": current_conf,
        "detections": detections
    }

@app.post("/vectorize_and_store")
async def vectorize_and_store(file: UploadFile = File(...), detection_data: str = Form(...)):
    """Cloud-dependent. Requires active billing."""
    import json
    try:
        primary_detection = json.loads(detection_data)
    except Exception:
        return {"status": "error", "message": "Invalid detection_data format"}

    unique_id = uuid.uuid4().hex
    temp_path = f"temp_vec_{unique_id}_{file.filename}"
    
    # Reset file pointer since it might have been read (though here it's a new request)
    # Actually, in a new request, we read it for the first time.
    with open(temp_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    try:
        record_id = analyze_use_case.execute(temp_path, primary_detection)
        return {
            "status": "success",
            "record_id": record_id,
            "message": "FFB vectorized and archived successfully"
        }
    except RuntimeError as e:
        return {"status": "error", "message": str(e)}
    except Exception as e:
        return {"status": "error", "message": f"An unexpected error occurred: {str(e)}"}
    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)

@app.post("/process_batch")
async def process_batch(files: List[UploadFile] = File(...)):
    """Handles multiple images: Detect -> Vectorize -> Store. Graceful handling of cloud errors."""
    batch_results = []
    temp_files = []

    try:
        for file in files:
            # 1. Save Temp
            unique_id = uuid.uuid4().hex
            path = f"temp_batch_{unique_id}_{file.filename}"
            with open(path, "wb") as f_out:
                shutil.copyfileobj(file.file, f_out)
            temp_files.append(path)

            # 2. YOLO Detect
            img = Image.open(path)
            yolo_res = yolo_model(img, conf=current_conf)
            
            # 3. Take the primary detection per image
            if yolo_res and yolo_res[0].boxes:
                box = yolo_res[0].boxes[0]
                batch_results.append({
                    "path": path,
                    "yolo": {
                        "class": yolo_model.names[int(box.cls)],
                        "confidence": float(box.conf),
                        "box": box.xyxy.tolist()[0]
                    }
                })

        if not batch_results:
            return {"status": "no_detection", "message": "No bunches detected in batch"}

        # 4. Process Batch Use Case with error handling for cloud services
        try:
            record_ids = analyze_batch_use_case.execute(batch_results)
            return {
                "status": "success", 
                "processed_count": len(record_ids), 
                "record_ids": record_ids,
                "message": f"Successfully processed {len(record_ids)} bunches"
            }
        except RuntimeError as e:
            return {
                "status": "partial_success",
                "message": f"Detections completed, but cloud archival failed: {str(e)}",
                "detections_count": len(batch_results)
            }

    except Exception as e:
        return {"status": "error", "message": f"Batch processing failed: {str(e)}"}

    finally:
        # 5. Clean up all temp files
        for path in temp_files:
            if os.path.exists(path):
                os.remove(path)
@app.post("/search_hybrid")
async def search_hybrid(
    file: Optional[UploadFile] = File(None), 
    text_query: Optional[str] = Form(None), 
    limit: int = Form(3)
):
    """Hybrid Search: Supports Visual Similarity and Natural Language Search."""
    temp_path = None
    try:
        try:
            if file:
                unique_id = uuid.uuid4().hex
                temp_path = f"temp_search_{unique_id}_{file.filename}"
                with open(temp_path, "wb") as buffer:
                    shutil.copyfileobj(file.file, buffer)
                
                results = search_use_case.execute(image_path=temp_path, limit=limit)
            elif text_query:
                results = search_use_case.execute(text_query=text_query, limit=limit)
            else:
                return {"status": "error", "message": "No search input provided"}

            return {"status": "success", "results": results}
        except RuntimeError as e:
            return {"status": "error", "message": f"Search unavailable: {str(e)}"}
        
    finally:
        if temp_path and os.path.exists(temp_path):
            os.remove(temp_path)

@app.get("/get_image/{record_id}")
async def get_image(record_id: str):
    """Retrieve the Base64 image data for a specific record."""
    record = repo.get_by_id(record_id)
    if not record:
        return {"status": "error", "message": "Record not found"}
    return {
        "status": "success",
        "image_data": record.get("image_data")
    }

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