from __future__ import annotations

import base64
import threading
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field

import kp


SERVICE_VERSION = "0.1.0"


@dataclass
class DeviceState:
    device_group: Optional[kp.DeviceGroup] = None
    port_id: Optional[int] = None
    model_desc: Optional[kp.ModelNefDescriptor] = None


STATE = DeviceState()
STATE_LOCK = threading.Lock()


app = FastAPI(title="Kneron LocalAPI", version=SERVICE_VERSION)


def _ok(data: Any) -> Dict[str, Any]:
    return {"ok": True, "data": data, "error": None}


def _err(code: str, message: str) -> Dict[str, Any]:
    return {"ok": False, "data": None, "error": {"code": code, "message": message}}


def _require_device() -> kp.DeviceGroup:
    if STATE.device_group is None:
        raise HTTPException(status_code=400, detail=_err("NO_DEVICE", "No connected device"))
    return STATE.device_group


def _image_format_from_str(value: str) -> kp.ImageFormat:
    value = value.upper()
    mapping = {
        "RGB565": kp.ImageFormat.KP_IMAGE_FORMAT_RGB565,
        "RGBA8888": kp.ImageFormat.KP_IMAGE_FORMAT_RGBA8888,
        "RAW8": kp.ImageFormat.KP_IMAGE_FORMAT_RAW8,
        "YUYV": kp.ImageFormat.KP_IMAGE_FORMAT_YUYV,
        "YUV420": kp.ImageFormat.KP_IMAGE_FORMAT_YUV420,
    }
    if value not in mapping:
        raise HTTPException(
            status_code=400,
            detail=_err("INVALID_IMAGE_FORMAT", f"Unsupported image_format: {value}"),
        )
    return mapping[value]


def _channels_ordering_from_str(value: str) -> kp.ChannelOrdering:
    value = value.upper()
    mapping = {
        "HCW": kp.ChannelOrdering.KP_CHANNEL_ORDERING_HCW,
        "CHW": kp.ChannelOrdering.KP_CHANNEL_ORDERING_CHW,
        "HWC": kp.ChannelOrdering.KP_CHANNEL_ORDERING_HWC,
        "DEFAULT": kp.ChannelOrdering.KP_CHANNEL_ORDERING_DEFAULT,
    }
    if value not in mapping:
        raise HTTPException(
            status_code=400,
            detail=_err("INVALID_CHANNEL_ORDERING", f"Unsupported channels_ordering: {value}"),
        )
    return mapping[value]


def _product_name_from_id(product_id: int) -> str:
    try:
        return kp.ProductId(product_id).name.replace("KP_DEVICE_", "")
    except ValueError:
        return "UNKNOWN"


@app.exception_handler(HTTPException)
def http_exception_handler(_: Request, exc: HTTPException) -> JSONResponse:
    if isinstance(exc.detail, dict) and "ok" in exc.detail:
        return JSONResponse(status_code=exc.status_code, content=exc.detail)
    return JSONResponse(
        status_code=exc.status_code,
        content=_err("HTTP_ERROR", str(exc.detail)),
    )


@app.exception_handler(Exception)
def unhandled_exception_handler(_: Request, exc: Exception) -> JSONResponse:
    return JSONResponse(
        status_code=500,
        content=_err("INTERNAL_ERROR", str(exc)),
    )


class ConnectRequest(BaseModel):
    port_id: Optional[int] = Field(default=None)
    scan_index: Optional[int] = Field(default=None)
    timeout_ms: Optional[int] = Field(default=5000)


class FirmwareLoadRequest(BaseModel):
    scpu_path: str
    ncpu_path: str


class ModelLoadRequest(BaseModel):
    nef_path: str


class InferenceRunRequest(BaseModel):
    model_id: int
    image_format: str
    width: int
    height: int
    image_base64: str
    channels_ordering: str = "DEFAULT"
    output_dtype: str = "float32"


@app.get("/health")
def health() -> Dict[str, Any]:
    return _ok({"status": "up"})


@app.get("/version")
def version() -> Dict[str, Any]:
    return _ok(
        {
            "service_version": SERVICE_VERSION,
            "kneronplus_version": kp.core.get_version(),
        }
    )


@app.get("/devices")
def devices() -> Dict[str, Any]:
    device_list = kp.core.scan_devices()
    devices_out = []
    for idx, device in enumerate(device_list.device_descriptor_list):
        devices_out.append(
            {
                "scan_index": idx,
                "usb_port_id": device.usb_port_id,
                "vendor_id": device.vendor_id,
                "product_id": f"0x{device.product_id:X}",
                "product_name": _product_name_from_id(device.product_id),
                "link_speed": device.link_speed.name,
                "usb_port_path": device.usb_port_path,
                "kn_number": device.kn_number,
                "is_connectable": device.is_connectable,
                "firmware": device.firmware,
            }
        )
    return _ok({"devices": devices_out})


@app.post("/devices/connect")
def connect(req: ConnectRequest) -> Dict[str, Any]:
    with STATE_LOCK:
        if STATE.device_group is not None:
            try:
                kp.core.disconnect_devices(STATE.device_group)
            except kp.ApiKPException:
                pass
            STATE.device_group = None
            STATE.port_id = None
            STATE.model_desc = None

        port_id = req.port_id
        if port_id is None:
            device_list = kp.core.scan_devices()
            if device_list.device_descriptor_number == 0:
                raise HTTPException(status_code=404, detail=_err("NO_DEVICE", "No device found"))
            if req.scan_index is None:
                scan_index = 0
            else:
                scan_index = req.scan_index
            if scan_index < 0 or scan_index >= device_list.device_descriptor_number:
                raise HTTPException(
                    status_code=400,
                    detail=_err("INVALID_SCAN_INDEX", f"Invalid scan_index: {scan_index}"),
                )
            port_id = device_list.device_descriptor_list[scan_index].usb_port_id

        try:
            device_group = kp.core.connect_devices([int(port_id)])
            if req.timeout_ms is not None:
                kp.core.set_timeout(device_group, int(req.timeout_ms))
        except kp.ApiKPException as exc:
            raise HTTPException(
                status_code=500,
                detail=_err(str(exc.api_return_code), str(exc)),
            )

        STATE.device_group = device_group
        STATE.port_id = int(port_id)
        return _ok({"connected": True, "port_id": STATE.port_id})


@app.post("/devices/disconnect")
def disconnect() -> Dict[str, Any]:
    with STATE_LOCK:
        if STATE.device_group is not None:
            try:
                kp.core.disconnect_devices(STATE.device_group)
            except kp.ApiKPException as exc:
                raise HTTPException(
                    status_code=500,
                    detail=_err(str(exc.api_return_code), str(exc)),
                )
            STATE.device_group = None
            STATE.port_id = None
            STATE.model_desc = None
    return _ok({"connected": False})


@app.post("/firmware/load")
def firmware_load(req: FirmwareLoadRequest) -> Dict[str, Any]:
    device_group = _require_device()
    try:
        kp.core.load_firmware_from_file(device_group, req.scpu_path, req.ncpu_path)
    except kp.ApiKPException as exc:
        raise HTTPException(
            status_code=500,
            detail=_err(str(exc.api_return_code), str(exc)),
        )
    return _ok({"loaded": True})


@app.post("/models/load")
def models_load(req: ModelLoadRequest) -> Dict[str, Any]:
    device_group = _require_device()
    try:
        model_desc = kp.core.load_model_from_file(device_group, req.nef_path)
    except kp.ApiKPException as exc:
        raise HTTPException(
            status_code=500,
            detail=_err(str(exc.api_return_code), str(exc)),
        )

    with STATE_LOCK:
        STATE.model_desc = model_desc

    models = []
    for model in model_desc.models:
        models.append(
            {
                "id": model.id,
                "input_nodes": len(model.input_nodes),
                "output_nodes": len(model.output_nodes),
                "max_raw_out_size": model.max_raw_out_size,
            }
        )

    return _ok({"models": models})


def _reset_device_and_clear_state(device_group: kp.DeviceGroup) -> None:
    kp.core.reset_device(device_group, kp.ResetMode.KP_RESET_REBOOT)
    kp.core.disconnect_devices(device_group)


@app.post("/models/clear")
def models_clear() -> Dict[str, Any]:
    device_group = _require_device()
    try:
        _reset_device_and_clear_state(device_group)
    except kp.ApiKPException as exc:
        raise HTTPException(
            status_code=500,
            detail=_err(str(exc.api_return_code), str(exc)),
        )

    with STATE_LOCK:
        STATE.device_group = None
        STATE.port_id = None
        STATE.model_desc = None

    return _ok({"cleared": True})


@app.post("/models/reset")
def models_reset() -> Dict[str, Any]:
    return models_clear()


@app.post("/inference/run")
def inference_run(req: InferenceRunRequest) -> Dict[str, Any]:
    device_group = _require_device()
    image_format = _image_format_from_str(req.image_format)
    channels_ordering = _channels_ordering_from_str(req.channels_ordering)
    if req.output_dtype.lower() != "float32":
        raise HTTPException(
            status_code=400,
            detail=_err("INVALID_OUTPUT_DTYPE", "Only float32 output is supported in PoC"),
        )

    try:
        if STATE.port_id is not None:
            kp.core.get_model_info(device_group, STATE.port_id)
    except kp.ApiKPException as exc:
        if exc.api_return_code == kp.ApiReturnCode.KP_ERROR_MODEL_NOT_LOADED_35:
            raise HTTPException(
                status_code=500,
                detail=_err(
                    "KP_ERROR_MODEL_NOT_LOADED_35",
                    str(kp.ApiReturnCode.KP_ERROR_MODEL_NOT_LOADED_35),
                ),
            )
        raise HTTPException(
            status_code=500,
            detail=_err(str(exc.api_return_code), str(exc)),
        )

    try:
        image_bytes = base64.b64decode(req.image_base64)
    except (ValueError, TypeError):
        raise HTTPException(
            status_code=400,
            detail=_err("INVALID_BASE64", "image_base64 is not valid base64 data"),
        )

    input_image = kp.GenericInputNodeImage(
        image=image_bytes,
        width=req.width,
        height=req.height,
        image_format=image_format,
    )

    input_desc = kp.GenericImageInferenceDescriptor(
        model_id=req.model_id,
        input_node_image_list=[input_image],
    )

    try:
        kp.inference.generic_image_inference_send(device_group, input_desc)
        result = kp.inference.generic_image_inference_receive(device_group)
    except kp.ApiKPException as exc:
        raise HTTPException(
            status_code=500,
            detail=_err(str(exc.api_return_code), str(exc)),
        )

    outputs = []
    for node_idx in range(result.header.num_output_node):
        try:
            node_output = kp.inference.generic_inference_retrieve_float_node(
                node_idx, result, channels_ordering
            )
        except kp.ApiKPException as exc:
            raise HTTPException(
                status_code=500,
                detail=_err(str(exc.api_return_code), str(exc)),
            )

        data_bytes = node_output.ndarray.astype("float32").tobytes()
        outputs.append(
            {
                "node_idx": node_idx,
                "name": node_output.name,
                "dtype": "float32",
                "shape": node_output.shape,
                "data_base64": base64.b64encode(data_bytes).decode("ascii"),
                "channels_ordering": channels_ordering.name,
            }
        )

    return _ok({"outputs": outputs})


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="127.0.0.1", port=4398)
