from __future__ import annotations

import json
from typing import Any

from pydantic import BaseModel, Field


class KieBase64UploadRequest(BaseModel):
    base64Data: str = Field(..., description="Base64 data URL payload.")
    uploadPath: str = Field("images/base64", description="Kie upload path bucket.")
    fileName: str = Field("reference.png", description="Uploaded file name.")


class KieUploadFileData(BaseModel):
    fileName: str | None = None
    filePath: str | None = None
    downloadUrl: str | None = None
    mimeType: str | None = None


class KieBase64UploadResponse(BaseModel):
    code: int
    msg: str
    data: KieUploadFileData


class KieNanoBanana2Input(BaseModel):
    prompt: str
    aspect_ratio: str = "auto"
    resolution: str = "2K"
    output_format: str = "png"
    google_search: bool = False
    image_input: list[str] = Field(default_factory=list)


class KieCreateTaskRequest(BaseModel):
    model: str = "nano-banana-2"
    input: KieNanoBanana2Input
    callBackUrl: str | None = None


class KieTaskCreateData(BaseModel):
    taskId: str


class KieCreateTaskResponse(BaseModel):
    code: int
    msg: str
    data: KieTaskCreateData


class KieTaskRecordData(BaseModel):
    taskId: str
    state: str
    resultJson: str | dict[str, Any] | None = None
    failCode: str | None = None
    failMsg: str | None = None

    def parsed_result(self) -> dict[str, Any]:
        if self.resultJson is None or self.resultJson == "":
            return {}
        if isinstance(self.resultJson, dict):
            return self.resultJson
        try:
            parsed = json.loads(self.resultJson)
        except json.JSONDecodeError:
            return {}
        return parsed if isinstance(parsed, dict) else {}

    def result_urls(self) -> list[str]:
        payload = self.parsed_result()
        urls = payload.get("resultUrls", [])
        if not isinstance(urls, list):
            return []
        return [url for url in urls if isinstance(url, str) and url]

    def failure_reason(self) -> str:
        fail_code = (self.failCode or "").strip()
        fail_msg = (self.failMsg or "").strip()
        if fail_code and fail_msg:
            return f"{fail_code}: {fail_msg}"
        return fail_code or fail_msg


class KieTaskRecordResponse(BaseModel):
    code: int
    msg: str
    data: KieTaskRecordData


KieBase64UploadResponse.model_rebuild()
KieCreateTaskRequest.model_rebuild()
KieCreateTaskResponse.model_rebuild()
KieTaskRecordResponse.model_rebuild()