import os
from typing import Any, List, Optional, Literal
import subprocess
from fastapi import UploadFile, File, Form
from fastapi.responses import FileResponse, JSONResponse
from fastapi import APIRouter, Depends, HTTPException
from fastapi import WebSocket, BackgroundTasks
from sqlalchemy.orm import Session
import asyncio
import time
import shutil
import app.crud as crud
import app.models as models
import app.schemas as schemas 
from app.api import deps
from datetime import datetime

from app.core.celery_app import celery_app
from app.core.config import settings
from pathlib import Path
from app.utils import random_name

from app.core.celery_app import celery_app
from gradio_client import Client

BACKEND_ZIP_STORAGE = Path("/app").joinpath(settings.BACKEND_ZIP_STORAGE)
LOCAL_ZIP_STORAGE = Path("/").joinpath(settings.LOCAL_ZIP_STORAGE)


router = APIRouter()
sr_clients = {}

@router.post("/sr")
def supser_resolution(
    *,
    db: Session = Depends(deps.get_db),
    current_user: models.User = Depends(deps.get_current_active_user),
    upload_files: List[UploadFile]=File(description="Multiple files as UploadFile"),
    background_tasks: BackgroundTasks,
) -> Any:
    """
    Super Resolution.
    """
    print("api start: "+str(datetime.now()))
    filenames = [random_name(20)+Path(file.filename).suffix for file in upload_files]
    stemnames = [Path(filename).stem for filename in filenames]
    new_dir = random_name(10)
    new_dir_path = Path(BACKEND_ZIP_STORAGE).joinpath(new_dir)
    new_dir_path.mkdir()
    for i in range(len(upload_files)):
      try:
          with open(str(Path(new_dir_path).joinpath(filenames[i])), 'wb') as f:
              while contents := upload_files[i].file.read(1024 * 1024):
                  f.write(contents)
      except Exception as e:
          print(e, type(e))
          return {"error": str(e)}
      finally:
          upload_files[i].file.close()
          
    background_tasks.add_task(wait_finish, new_dir, filenames)

    print(filenames)
    return JSONResponse({"filenames": filenames}, background=background_tasks)

async def wait_finish(dirname, filenames):
    process = await asyncio.create_subprocess_exec("sshpass", "-p", "choozmo9", 
                     "scp", "-P", "5722", "-o", "StrictHostKeyChecking=no", "-r", f"{str(BACKEND_ZIP_STORAGE/dirname)}", f"root@172.104.93.163:{str(LOCAL_ZIP_STORAGE)}")
    await process.wait()
    # r = subprocess.run(["sshpass", "-p", "choozmo9", 
    #                 "scp", "-P", "5722", "-o", "StrictHostKeyChecking=no", "-r", f"{str(new_dir_path)}", f"root@172.104.93.163:{str(LOCAL_ZIP_STORAGE)}"])
    res = celery_app.send_task("app.worker.super_resolution", args=[dirname, filenames])

    while True:
       await asyncio.sleep(0.5)
       if res.state == "SUCCESS":
           break
    restored_imgs = "restored_imgs"
    process = await asyncio.create_subprocess_exec("sshpass", "-p", "choozmo9", 
                     "scp", "-P", "5722", "-o", "StrictHostKeyChecking=no", "-r", f"root@172.104.93.163:{str(LOCAL_ZIP_STORAGE/dirname/restored_imgs)}/*", f"{str(BACKEND_ZIP_STORAGE)}")
    await process.wait()
    for filename in filenames:
        await publish(filename)

    process = await asyncio.create_subprocess_exec("sshpass", "-p", "choozmo9", 
                     "ssh", "-P", "5722", "-o", "StrictHostKeyChecking=no", "root@172.104.93.163", "rm", "-r", f"{str(LOCAL_ZIP_STORAGE/dirname)}")
    await process.wait()


async def publish(data):
    for sr_client in sr_clients.values():
        await sr_client.send_text(f"{data}")

@router.get("/sr")
def get_image(
    *,
    db: Session = Depends(deps.get_db),
    current_user: models.User = Depends(deps.get_current_active_user),
    stored_file_name: str,
    file_name: str,
    background_tasks: BackgroundTasks
) -> Any:
    """
    Download image
    """
    return_files = list(BACKEND_ZIP_STORAGE.glob(stored_file_name))
    if return_files:
        return_file = return_files[0]
        background_tasks.add_task(delete_img, return_file)
        return FileResponse(path=str(return_file), filename=file_name+"_hr"+return_file.suffix, background=background_tasks)
    else:
        print("non")

async def delete_img(filename:str):
    await asyncio.sleep(1)
    path = Path(filename)
    if path.exists():
        path.unlink()
  

@router.delete("/sr")
def del_image():
    pass

aspects = ['16:9', '3:2', '1:1', '2:3', '9:16']
flux_image_size = [
    (448, 800),
    (480, 720),
    (600, 600),
    (720, 480),
    (800, 448)
]

def remove_file(path: str) -> None:
    os.remove(path)

@router.post('/flux')
def get_image(
    *,
    db: Session = Depends(deps.get_db),
    current_user: models.User = Depends(deps.get_current_active_user),
    prompt:str, 
    aspect:Literal['16:9', '3:2', '1:1', '2:3', '9:16'],
    background_tasks: BackgroundTasks
) -> Any:
    width, height = flux_image_size[aspects.index(aspect)]
    client = Client("http://172.104.93.163:7860/")
    result = client.predict(
            model_id="models/FLUX.1-schnell",
            prompt=prompt,
            width=width,
            height=height,
            seed=-1,
            steps=4,
            guidance_scale=3.5,
            add_sampling_metadata=True,
            api_name="/generate"
    )   
    background_tasks.add_task(remove_file, result[0])
    return FileResponse(result[0], media_type='image/webp', background=background_tasks)

@router.websocket("/sr")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    key = websocket.headers.get('sec-websocket-key')
    sr_clients[key] = websocket
    print(f"new comer: {key}")
    try:
        while True:
            data = await websocket.receive_text()
            print(f"{key}:{data}")
            if data.startswith("unsubscribe"):
              del sr_clients[key]
              #for client in sr_clients.values():
              #      await client.send_text(f"ID: {key} | Message: {data}")

    except:
        # 接続が切れた場合、当該クライアントを削除する
        del sr_clients[key]
        
        
