from __future__ import annotations
import argparse
import datetime
import json
import os
from operator import itemgetter
from pathlib import Path
from typing import Callable

import requests
import pynvml
import PIL.Image
import torch
from pydantic import BaseSettings, BaseModel


class Args(BaseSettings):

    @classmethod
    def from_args(cls):
        parser = argparse.ArgumentParser()
        for field in cls.__fields__.values():
            if issubclass(field.type_, BaseModel):
                prefix = field.type_.__name__.lower()
                for subfield in field.type_.__fields__.values():
                    short = "".join([x[0] for x in subfield.name.split("_")])
                    parser.add_argument(f"--{prefix}.{subfield.name}", default=subfield.default, required=subfield.required)
            else:
                short = "".join([x[0] for x in field.name.split("_")])
                parser.add_argument(f"-{short}", f"--{field.name}", default=field.default, required=field.required)
        args = vars(parser.parse_known_args()[0])
        to_delete = set()
        for field in cls.__fields__.values():
            if issubclass(field.type_, BaseModel):
                prefix = field.type_.__name__.lower()
                sub_args = {}
                for k, v in args.items():
                    if k.startswith(prefix):
                        to_delete.add(k)
                        sub_args[k.replace(f"{prefix}.", "")] = v
                args[field.name] = sub_args
        args = {k: v for k, v in args.items() if k not in to_delete}
        return cls(**args)

    class Config:
        env_file = ".env"
        env_file_encoding = "utf-8"
        env_prefix = "ARG_"


def inject_args(func: Callable) -> Callable:
    """Decorates a function to inject the arguments."""
    
    injected = None
    for type_ in func.__annotations__.values():
        if issubclass(type_, Args):
            injected = type_.from_args()
            break

    if injected is None:
        raise ValueError(f"Function {func.__name__} is not annotated with an Args subclass.")

    def wrapper(*args, **kwargs):
        return func(injected, *args, **kwargs)

    return wrapper


def get_free_gpu() -> int:
    pynvml.nvmlInit()
    total = torch.cuda.device_count()
    gpus = []
    for i in range(total):
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
        info = pynvml.nvmlDeviceGetMemoryInfo(handle)
        gpus.append((i, info.free))
    gpus = sorted(gpus, key=itemgetter(1), reverse=True)
    return gpus[0][0]


def get_user_name() -> str:
    return Path(os.environ["HOME"]).stem


def get_storage_dir() -> Path:
    return Path(f"/fsx/{get_user_name()}")


def get_checkpoints_dir(*, timestamp: bool) -> Path:
    base_dir = get_storage_dir()/"checkpoints"
    return Path(f"{base_dir}/{now()}") if timestamp else base_dir


def now() -> str:
    return datetime.datetime.utcnow().strftime("%Y%m%d_%H%M%S")


def read_jsonl(path: Path) -> list:
    return [json.loads(x) for x in Path(path).read_text().split("\n") if x]


def write_jsonl(path: Path, data: list):
    with Path(path).open("w") as f:
        for x in data:
            f.write(json.dumps(x) + "\n")


def get_image(url: str, filename: Path | None = None):
    if filename is None: filename = Path(f"{url.split('/')[-1]}.jpg")
    filename = Path(filename)
    if filename.exists(): return filename
    PIL.Image.open(requests.get(url, stream=True).raw).save(filename)
    return filename


def find_latest_checkpoint(dirname: Path) -> Path:
    dirname = Path(dirname)
    checkpoints = list(dirname.glob("*.ckpt"))
    if not checkpoints:
        return None
    latest = max(checkpoints, key=lambda path: path.stat().st_mtime)
    return latest


def list_files(dirname: Path, exts: list[str] | None = None) -> list:
    files = Path(dirname).iterdir()
    if not exts:
        return list(files)
    return [fn for fn in files for ext in exts if fn.match(f"*.{ext}")]