import io
import os
import uuid
import logging
from typing import Optional

from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, Header
from fastapi.responses import FileResponse, JSONResponse
from pydantic import BaseModel
import torch
import asyncio
import time

import numpy as np
from PIL import Image

# Lazy import performed in get_model() to avoid import-time failures on Space


LOGGER = logging.getLogger("hair_server")
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")

EXPECTED_BEARER = "logicgo@123"

# Optional Mongo persistence
from pymongo import MongoClient
MONGO_URI = os.environ.get("MONGO_URI", "")
mongo_client = MongoClient(MONGO_URI) if MONGO_URI else None
mongo_db = mongo_client.get_database("HairSwapDB") if mongo_client else None
uploads_col = mongo_db.get_collection("uploads") if mongo_db else None
results_col = mongo_db.get_collection("results") if mongo_db else None


def verify_bearer(authorization: Optional[str] = Header(None)):
    if not authorization:
        raise HTTPException(status_code=401, detail="Missing Authorization header")
    try:
        scheme, token = authorization.split(" ", 1)
    except ValueError:
        raise HTTPException(status_code=401, detail="Invalid Authorization header format")
    if scheme.lower() != "bearer":
        raise HTTPException(status_code=401, detail="Invalid auth scheme")
    if token != EXPECTED_BEARER:
        raise HTTPException(status_code=401, detail="Invalid token")
    return True


app = FastAPI(title="Hair Swap API", version="1.0.0")


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


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


@app.get("/health/gpu")
def health_gpu():
    try:
        import torch as _t
        return {
            "cuda_available": bool(_t.cuda.is_available()),
            "cuda_device_count": int(_t.cuda.device_count()) if _t.cuda.is_available() else 0,
            "cuda_device_name": _t.cuda.get_device_name(0) if _t.cuda.is_available() else None,
            "torch_version": getattr(_t, "__version__", None),
            "torch_cuda_version": getattr(_t.version, "cuda", None),
        }
    except Exception as e:
        return {"error": str(e)}


class HairSwapRequest(BaseModel):
    source_id: str
    reference_id: str
    converter_scale: float = 1.0
    scale: float = 1.0
    guidance_scale: float = 1.5
    controlnet_conditioning_scale: float = 1.0


# Initialize model lazily on first request
_model = None  # type: ignore[assignment]


def get_model():
    global _model
    if _model is None:
        try:
            LOGGER.info("Loading StableHair model ...")
            # Enforce GPU-only execution on Spaces with L4/other CUDA GPUs
            if not torch.cuda.is_available():
                raise RuntimeError("CUDA GPU not available. Install CUDA-enabled PyTorch and run on a GPU Space.")
            device = "cuda"
            dtype = torch.float16
            # Speed knobs for NVIDIA GPUs
            try:
                torch.backends.cuda.matmul.allow_tf32 = True  # type: ignore[attr-defined]
                torch.backends.cudnn.allow_tf32 = True  # type: ignore[attr-defined]
                torch.backends.cudnn.benchmark = True  # type: ignore[attr-defined]
            except Exception:
                pass
            LOGGER.info(f"Using device: {device}, dtype: {dtype}")

            # Ensure HF token env var is where downstream libs expect it
            if os.environ.get("HUGGINGFACEHUB_API_TOKEN") and not os.environ.get("HUGGINGFACE_HUB_TOKEN"):
                os.environ["HUGGINGFACE_HUB_TOKEN"] = os.environ["HUGGINGFACEHUB_API_TOKEN"]

            # Backward-compat shim: some diffusers versions import a helper only present in newer hub versions.
            try:
                import huggingface_hub as _hfh  # type: ignore

                if not hasattr(_hfh, "split_torch_state_dict_into_shards"):
                    def _split_torch_state_dict_into_shards(state_dict, max_shard_size="10GB"):
                        # Minimal shim: return a single shard mapping expected by callers
                        return {"pytorch_model.bin": state_dict}

                    _hfh.split_torch_state_dict_into_shards = _split_torch_state_dict_into_shards  # type: ignore[attr-defined]
            except Exception:
                pass
            
            # Import here to defer importing diffusers/transformers until needed
            from infer_full import StableHair  # noqa: WPS433
            _model = StableHair(config="./configs/hair_transfer.yaml", device=device, weight_dtype=dtype)
            LOGGER.info("Model loaded successfully")
        except Exception as e:
            LOGGER.error(f"Failed to load model: {str(e)}")
            raise Exception(f"Model loading failed: {str(e)}")
    return _model


# Use a writable location on Hugging Face Spaces
BASE_DATA_DIR = os.environ.get("DATA_DIR", "/data")
UPLOAD_DIR = os.path.join(BASE_DATA_DIR, "uploads")
RESULTS_DIR = os.path.join(BASE_DATA_DIR, "results")
LOGS_DIR = os.path.join(BASE_DATA_DIR, "logs")
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(RESULTS_DIR, exist_ok=True)
os.makedirs(LOGS_DIR, exist_ok=True)


@app.post("/upload")
async def upload_image(image: UploadFile = File(...), _=Depends(verify_bearer)):
    if not image.filename:
        raise HTTPException(status_code=400, detail="No file name provided")
    contents = await image.read()
    try:
        Image.open(io.BytesIO(contents)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid image file")

    image_id = str(uuid.uuid4())
    ext = os.path.splitext(image.filename)[1] or ".png"
    path = os.path.join(UPLOAD_DIR, image_id + ext)
    with open(path, "wb") as f:
        f.write(contents)
    # Save metadata to Mongo
    if uploads_col:
        try:
            uploads_col.insert_one({"_id": image_id, "filename": os.path.basename(path), "path": path})
        except Exception:
            pass
    return {"id": image_id, "filename": os.path.basename(path)}


@app.post("/get-hairswap")
def get_hairswap(req: HairSwapRequest, _=Depends(verify_bearer)):
    try:
        # Resolve file paths
        def find_file(image_id: str) -> str:
            for name in os.listdir(UPLOAD_DIR):
                if name.startswith(image_id):
                    return os.path.join(UPLOAD_DIR, name)
            raise HTTPException(status_code=404, detail=f"Image id not found: {image_id}")

        source_path = find_file(req.source_id)
        reference_path = find_file(req.reference_id)
        
        LOGGER.info(f"Found source: {source_path}, reference: {reference_path}")

        # Load model with error handling
        try:
            model = get_model()
            LOGGER.info("Model loaded successfully")
        except Exception as e:
            LOGGER.error(f"Model loading failed: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Model loading failed: {str(e)}")

        # Perform hair transfer with error handling
        try:
            LOGGER.info("Starting hair transfer...")
            # Log current schedulers for visibility
            try:
                sched_main = type(model.pipeline.scheduler).__name__ if hasattr(model, "pipeline") else None
                sched_bald = type(model.remove_hair_pipeline.scheduler).__name__ if hasattr(model, "remove_hair_pipeline") else None
                LOGGER.info(f"Schedulers -> main: {sched_main}, remove_hair: {sched_bald}")
            except Exception:
                pass
            id_np, out_np, bald_np, ref_np = model.Hair_Transfer(
                source_image=source_path,
                reference_image=reference_path,
                random_seed=-1,
                step=20,
                guidance_scale=req.guidance_scale,
                scale=req.scale,
                controlnet_conditioning_scale=req.controlnet_conditioning_scale,
                size=448,
            )
            LOGGER.info("Hair transfer completed successfully")
        except Exception as e:
            import traceback
            tb = traceback.format_exc()
            LOGGER.error(f"Hair transfer failed: {str(e)} | device={model.device if hasattr(model, 'device') else 'n/a'} cuda_available={torch.cuda.is_available()}\n{tb}")
            raise HTTPException(status_code=500, detail=f"Hair transfer failed: {str(e)}")

        # Save result
        try:
            result_id = str(uuid.uuid4())
            out_img = Image.fromarray((out_np * 255.).astype(np.uint8))
            filename = f"{result_id}.png"
            out_path = os.path.join(RESULTS_DIR, filename)
            out_img.save(out_path)
            LOGGER.info(f"Result saved: {out_path}")
            
            if results_col:
                try:
                    results_col.insert_one({
                        "_id": result_id,
                        "filename": filename,
                        "path": out_path,
                        "source_id": req.source_id,
                        "reference_id": req.reference_id,
                    })
                except Exception as e:
                    LOGGER.warning(f"MongoDB save failed: {str(e)}")

            return {"result": filename}
        except Exception as e:
            LOGGER.error(f"Result saving failed: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Result saving failed: {str(e)}")
            
    except HTTPException:
        raise
    except Exception as e:
        LOGGER.error(f"Unexpected error in get_hairswap: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")


@app.get("/download/{filename}")
def download(filename: str, _=Depends(verify_bearer)):
    path = os.path.join(RESULTS_DIR, filename)
    if not os.path.exists(path):
        raise HTTPException(status_code=404, detail="File not found")
    return FileResponse(path, media_type="image/png", filename=filename)


@app.get("/logs")
def logs(_=Depends(verify_bearer)):
    if uploads_col and results_col:
        uploads = list(uploads_col.find({}, {"_id": 1, "filename": 1}).limit(20))
        results = list(results_col.find({}, {"_id": 1, "filename": 1, "source_id": 1, "reference_id": 1}).limit(20))
        return JSONResponse({"uploads": uploads, "results": results})
    return JSONResponse({"logs": ["service running"], "db": "not_configured"})


# -------------------- Async job API --------------------
MAX_CONCURRENCY = int(os.environ.get("MAX_CONCURRENCY", "1"))
_sem = asyncio.Semaphore(MAX_CONCURRENCY)
_jobs = {}


@app.post("/get-hairswap-async")
async def get_hairswap_async(req: HairSwapRequest, _=Depends(verify_bearer)):
    job_id = str(uuid.uuid4())
    _jobs[job_id] = {"status": "queued", "result": None, "error": None, "started_at": None, "ended_at": None}

    async def _run_job():
        _jobs[job_id]["status"] = "running"
        _jobs[job_id]["started_at"] = time.time()
        try:
            async with _sem:
                # reuse the same core flow as sync endpoint
                def find_file(image_id: str) -> str:
                    for name in os.listdir(UPLOAD_DIR):
                        if name.startswith(image_id):
                            return os.path.join(UPLOAD_DIR, name)
                    raise HTTPException(status_code=404, detail=f"Image id not found: {image_id}")

                source_path = find_file(req.source_id)
                reference_path = find_file(req.reference_id)
                LOGGER.info(f"[job {job_id}] Found source: {source_path}, reference: {reference_path}")

                model = get_model()
                LOGGER.info(f"[job {job_id}] Model loaded successfully")

                LOGGER.info(f"[job {job_id}] Starting hair transfer...")
                try:
                    try:
                        sched_main = type(model.pipeline.scheduler).__name__ if hasattr(model, "pipeline") else None
                        sched_bald = type(model.remove_hair_pipeline.scheduler).__name__ if hasattr(model, "remove_hair_pipeline") else None
                        LOGGER.info(f"[job {job_id}] Schedulers -> main: {sched_main}, remove_hair: {sched_bald}")
                    except Exception:
                        pass

                    id_np, out_np, bald_np, ref_np = model.Hair_Transfer(
                        source_image=source_path,
                        reference_image=reference_path,
                        random_seed=-1,
                        step=20,
                        guidance_scale=req.guidance_scale,
                        scale=req.scale,
                        controlnet_conditioning_scale=req.controlnet_conditioning_scale,
                        size=448,
                    )
                    result_id = str(uuid.uuid4())
                    out_img = Image.fromarray((out_np * 255.).astype(np.uint8))
                    filename = f"{result_id}.png"
                    out_path = os.path.join(RESULTS_DIR, filename)
                    out_img.save(out_path)
                    _jobs[job_id]["result"] = {"filename": filename}
                    _jobs[job_id]["status"] = "completed"
                    LOGGER.info(f"[job {job_id}] Completed -> {out_path}")
                except Exception as e:
                    LOGGER.error(f"[job {job_id}] Hair transfer failed: {str(e)}")
                    _jobs[job_id]["error"] = str(e)
                    _jobs[job_id]["status"] = "failed"
        except Exception as e:
            _jobs[job_id]["error"] = str(e)
            _jobs[job_id]["status"] = "failed"
        finally:
            _jobs[job_id]["ended_at"] = time.time()

    asyncio.create_task(_run_job())
    return {"job_id": job_id, "status": "queued"}


@app.get("/job/{job_id}")
def job_status(job_id: str, _=Depends(verify_bearer)):
    data = _jobs.get(job_id)
    if not data:
        raise HTTPException(status_code=404, detail="job not found")
    return data