#!/usr/bin/env python3
"""
LoRA Studio — FastAPI Backend
Seedream 4.0 Text-to-Image with LoRA + Style Reference
4K images in 2.39:1 cinematic ratio via Volcengine API.
"""

import os
import json
import time
import uuid
import requests
import io
import base64
from pathlib import Path
from typing import Optional
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from dotenv import load_dotenv
from loguru import logger
from PIL import Image

load_dotenv(os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env"), override=True)
load_dotenv(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".env"), override=False)

from volcenginesdkcore.signv4 import SignerV4

app = FastAPI(title="LoRA Studio")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

# Serve generated images
OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "outputs")
os.makedirs(OUTPUT_DIR, exist_ok=True)
app.mount("/outputs", StaticFiles(directory=OUTPUT_DIR), name="outputs")

# ============================================================
# VOLCENGINE CONFIG
# ============================================================
HOST = "visual.volcengineapi.com"
ENDPOINT = "https://visual.volcengineapi.com"
REGION = "cn-north-1"
SERVICE = "cv"
API_VERSION = "2025-06-01"
TEMPLATE_ID = "1762940890894000"
MG_MV_ID = "592942110437597155"
BASE_MODEL_ID = "592942110404108259"
BASE_MODEL_VER_ID = "592942110437597155"

LORA_PRESETS = {
    "krstyle": {
        "name": "KRSTYLE — Krishna Mystical Forest",
        "trigger": "KRSTYLE",
        "card_id": "614342735335368774",
        "ver_id": "614367323301661726",
        "description": "Deep blue-teal mystical forests, smooth digital painting, matte quality",
    },
    "lkstyle": {
        "name": "LKSTYLE — Arcane / Fortiche",
        "trigger": "LKSTYLE",
        "card_id": "614122746506477325",
        "ver_id": "614247929653245457",
        "description": "Arcane animated series style, cinematic moody lighting",
    },
    "hantex": {
        "name": "Hantex — Cinematic Concept Art",
        "trigger": "",
        "card_id": "614078666753801530",
        "ver_id": "614103179759504414",
        "description": "Cinematic matte painting, concept art quality",
    },
    "none": {
        "name": "None — Base Seedream 4.0",
        "trigger": "",
        "card_id": "",
        "ver_id": "",
        "description": "No LoRA, pure Seedream 4.0 base model",
    },
}


# ============================================================
# API HELPERS
# ============================================================
def _signed_post(action, body):
    ak = os.environ.get("VOLC_ACCESS_KEY", "")
    sk = os.environ.get("VOLC_SECRET_KEY", "")
    req_body = json.dumps(body, ensure_ascii=False)
    query = {"Action": action, "Version": API_VERSION}
    headers = {"Content-Type": "application/json", "Host": HOST}
    SignerV4.sign(path="/", method="POST", headers=headers, body=req_body,
                  post_params=None, query=query, ak=ak, sk=sk,
                  region=REGION, service=SERVICE)
    url = f"{ENDPOINT}?Action={action}&Version={API_VERSION}"
    resp = requests.post(url, headers=headers, data=req_body.encode("utf-8"), timeout=120)
    if not resp.text:
        return {"code": -1, "message": f"Empty response, HTTP {resp.status_code}"}
    return resp.json()


def _signed_get(action, extra=None):
    ak = os.environ.get("VOLC_ACCESS_KEY", "")
    sk = os.environ.get("VOLC_SECRET_KEY", "")
    query = {"Action": action, "Version": API_VERSION}
    if extra:
        query.update(extra)
    headers = {"Content-Type": "application/json", "Host": HOST}
    SignerV4.sign(path="/", method="GET", headers=headers, body="",
                  post_params=None, query=query, ak=ak, sk=sk,
                  region=REGION, service=SERVICE)
    from urllib.parse import urlencode
    url = f"{ENDPOINT}?{urlencode(query)}"
    resp = requests.get(url, headers=headers)
    return resp.json()


def _poll(task_id, timeout=300, interval=8):
    start = time.time()
    while time.time() - start < timeout:
        result = _signed_get("GetLumiInferenceTask", {"id": str(task_id)})
        if not result or result.get("code") != 0:
            time.sleep(interval)
            continue
        status = result.get("data", {}).get("status", "")
        if status == "complete":
            return result
        elif status in ("failed", "cancel", "risk", "timeout", "invalid_param"):
            return result
        time.sleep(interval)
    return None


def _upload_to_s3(filepath):
    fpath = Path(filepath)
    with open(fpath, 'rb') as f:
        data = f.read()
    ext = fpath.suffix.lower()
    mime = 'image/png' if ext == '.png' else 'image/jpeg'
    files = {"file": (fpath.name, data, mime)}
    headers = {"x-api-key": "juniordevKey@9911"}
    resp = requests.post(
        "https://tinify-backend-dev-868570596092.asia-south1.run.app/api/upload-file",
        files=files, headers=headers, timeout=120
    )
    if resp.status_code == 200:
        result = resp.json()
        return result.get("url") or result.get("fileUrl") or result.get("s3_url")
    return None


# ============================================================
# ENDPOINTS
# ============================================================
@app.get("/api/presets")
def get_presets():
    return [{"id": k, **v} for k, v in LORA_PRESETS.items()]


@app.post("/api/generate")
def generate(
    prompt: str = Form(...),
    lora_id: str = Form("none"),
    lora_weight: float = Form(0.8),
    denoise: float = Form(0.5),
    guidance: float = Form(3.0),
    steps: int = Form(20),
    use_4k: bool = Form(True),
    ref_image: Optional[UploadFile] = File(None),
):
    preset = LORA_PRESETS.get(lora_id, LORA_PRESETS["none"])
    trigger = preset.get("trigger", "")
    card_id = preset.get("card_id", "")
    ver_id = preset.get("ver_id", "")

    full_prompt = f"{trigger} {prompt}".strip() if trigger else prompt
    if len(full_prompt) > 1000:
        full_prompt = full_prompt[:997] + "..."

    w, h = (3840, 1646) if use_4k else (2688, 1152)

    # Upload ref image if provided
    ref_url = None
    if ref_image:
        tmp_path = os.path.join(OUTPUT_DIR, f"ref_{uuid.uuid4().hex[:8]}{Path(ref_image.filename).suffix}")
        with open(tmp_path, "wb") as f:
            f.write(ref_image.file.read())
        ref_url = _upload_to_s3(tmp_path)
        os.remove(tmp_path)

    inference_type = "i2i" if ref_url else "t2i"

    body = {
        "task_name": f"studio_{int(time.time())}",
        "template_id": TEMPLATE_ID,
        "mg_mv_id": MG_MV_ID,
        "inference_config": {
            "inference_pipeline": "ba_worker",
            "inference_id": BASE_MODEL_ID,
            "inference_ver_id": BASE_MODEL_VER_ID,
        },
        "request_source": 2,
        "count": 1,
        "inference_type": inference_type,
        "inputs": [
            {"name": "seed", "internal_name": "seed", "format": "input_seed", "value": "-1"},
            {"name": "denoise_strength", "internal_name": "denoise_strength", "format": "input_number", "value": str(denoise)},
            {"name": "guidance_scale", "internal_name": "guidance_scale", "format": "input_number", "value": str(guidance)},
            {"name": "steps", "internal_name": "steps", "format": "input_number", "value": str(steps)},
            {"name": "prompt", "internal_name": "prompt", "format": "text_area", "value": full_prompt},
            {"name": "width", "internal_name": "width", "format": "input_number", "value": str(w)},
            {"name": "height", "internal_name": "height", "format": "input_number", "value": str(h)},
            {"name": "process_type", "internal_name": "process_type", "format": "select", "value": "基础生成"},
        ],
    }

    if card_id and ver_id:
        lora_json = json.dumps([{
            "model_id": card_id,
            "model_version_id": ver_id,
            "configs": [{"name": "weight", "format": "input_number", "value": lora_weight}]
        }])
        body["inputs"].append(
            {"name": "lora", "internal_name": "lora", "format": "model_upload", "value": lora_json}
        )

    if ref_url:
        body["inputs"].append(
            {"name": "img0", "internal_name": "img0", "format": "image_upload", "value": ref_url}
        )

    result = _signed_post("CreateLumiInferenceTask", body)
    if not result or result.get("code") != 0:
        raise HTTPException(500, f"API error: {result}")

    task_id = (result.get("data", {}).get("parent_task_id")
               or result.get("data", {}).get("task_id"))
    if not task_id:
        raise HTTPException(500, f"No task_id: {result}")

    poll_result = _poll(task_id, timeout=300, interval=8)
    if not poll_result:
        raise HTTPException(504, "Timeout — task did not complete")

    status = poll_result.get("data", {}).get("status", "")
    if status != "complete":
        raise HTTPException(500, f"Task failed: {status}")

    children = poll_result.get("data", {}).get("children", [])
    if not children:
        raise HTTPException(500, "No output")

    img_url = children[0].get("output", {}).get("value", "")
    if not img_url or not img_url.startswith("http"):
        raise HTTPException(500, "No image URL in output")

    # Download and save
    resp = requests.get(img_url, verify=False, timeout=60)
    if resp.status_code != 200:
        raise HTTPException(500, f"Download failed: {resp.status_code}")

    img = Image.open(io.BytesIO(resp.content))
    filename = f"{uuid.uuid4().hex[:12]}.png"
    filepath = os.path.join(OUTPUT_DIR, filename)
    img.save(filepath, quality=100)

    return {
        "image_url": f"/outputs/{filename}",
        "width": img.size[0],
        "height": img.size[1],
        "task_id": task_id,
        "lora": preset["name"],
        "prompt": full_prompt,
    }


# Serve React frontend
FRONTEND_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "frontend")
if os.path.exists(FRONTEND_DIR):
    app.mount("/", StaticFiles(directory=FRONTEND_DIR, html=True), name="frontend")


if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)