import json
import aiohttp
import pprint
import decimal
import aiofiles
import sys
import asyncio
import aiosqlite
import base64
import logging
from collections import defaultdict
from pathlib import Path
from urllib.parse import urlparse
from typing import Any, Optional, List, Set
from aiolimiter import AsyncLimiter
from dataclasses import dataclass, field

log = logging.getLogger(__name__)


DEFAULTS = [
    "wd14-vit-v2-git",
    "wd14-vit",
    "wd14-swinv2-v2-git",
    "wd14-convnextv2-v2-git",
    "wd14-convnext-v2-git",
    "wd14-convnext",
]


@dataclass
class Interrogator:
    model_id: str
    address: str


class DDInterrogator(Interrogator):
    async def interrogate(self, session, path):
        async with session.post(
            f"{self.address}/",
            params={
                "threshold": "0.7",
            },
            headers={"Authorization": "Bearer sex"},
            data={"file": path.open("rb")},
        ) as resp:
            assert resp.status == 200
            tags = await resp.json()
            upstream_tags = [tag.replace(" ", "_") for tag in tags]
            return " ".join(upstream_tags)


class SDInterrogator(Interrogator):
    async def interrogate(self, session, path):
        async with aiofiles.open(path, "rb") as fd:
            as_base64 = base64.b64encode(await fd.read()).decode("utf-8")

        async with session.post(
            f"{self.address}/tagger/v1/interrogate",
            json={
                "model": self.model_id,
                "threshold": 0.7,
                "image": as_base64,
            },
        ) as resp:
            log.info("got %d", resp.status)
            assert resp.status == 200
            data = await resp.json()
            tags_with_scores = data["caption"]
            tags = list(tags_with_scores.keys())

            upstream_tags = [tag.replace(" ", "_") for tag in tags]
            return " ".join(upstream_tags)


@dataclass
class Config:
    sd_webui_address: str
    dd_address: str
    dd_model_name: str
    sd_webui_models: List[str] = field(default_factory=lambda: list(DEFAULTS))

    @property
    def all_available_models(self) -> List[Any]:
        return [
            SDInterrogator(sd_interrogator, self.sd_webui_address)
            for sd_interrogator in self.sd_webui_models
        ] + [DDInterrogator(self.dd_model_name, self.dd_address)]


@dataclass
class Context:
    db: Any
    session: aiohttp.ClientSession
    config: Config


@dataclass
class Booru:
    session: aiohttp.ClientSession
    limiter: AsyncLimiter
    tag_limiter: AsyncLimiter
    file_locks: dict = field(default_factory=lambda: defaultdict(asyncio.Lock))
    tag_locks: dict = field(default_factory=lambda: defaultdict(asyncio.Lock))

    fetch_type = "hash"

    @property
    def hash_style(self):
        return HashStyle.md5


class Danbooru(Booru):
    title = "Danbooru"
    base_url = "https://danbooru.donmai.us"

    async def posts(self, tag_query: str, limit):
        log.info("%s: submit %r", self.title, tag_query)
        async with self.limiter:
            log.info("%s: submit upstream %r", self.title, tag_query)
            async with self.session.get(
                f"{self.base_url}/posts.json",
                params={"tags": tag_query, "limit": limit},
            ) as resp:
                assert resp.status == 200
                rjson = await resp.json()
                return rjson


DOWNLOADS = Path.cwd() / "posts"
DOWNLOADS.mkdir(exist_ok=True)


ALLOWED_EXTENSIONS = (
    "jpg",
    "jpeg",
    "png",
)


async def download_images(ctx):
    try:
        tagquery = sys.argv[2]
    except IndexError:
        tagquery = ""

    try:
        limit = int(sys.argv[3])
    except IndexError:
        limit = 10

    danbooru = Danbooru(
        ctx.session,
        AsyncLimiter(1, 10),
        AsyncLimiter(1, 3),
    )

    posts = await danbooru.posts(tagquery, limit)
    for post in posts:
        if "md5" not in post:
            continue
        md5 = post["md5"]
        log.info("processing post %r", md5)
        existing_post = await fetch_post(ctx, md5)
        if existing_post:
            continue

        # download the post
        post_file_url = post["file_url"]
        post_url_path = Path(urlparse(post_file_url).path)
        good_extension = False
        for extension in ALLOWED_EXTENSIONS:
            if extension in post_url_path.stem:
                good_extension = True
        if not good_extension:
            log.info("ignoring %r, invalid extension (%r)", md5, post_url_path.stem)
            continue
        post_filename = post_url_path.name
        post_filepath = DOWNLOADS / post_filename
        if not post_filepath.exists():
            log.info("downloading %r to %r", post_file_url, post_filepath)
            async with ctx.session.get(post_file_url) as resp:
                assert resp.status == 200
                with post_filepath.open("wb") as fd:
                    async for chunk in resp.content.iter_chunked(1024):
                        fd.write(chunk)

        # when it's done successfully, insert it
        await insert_post(ctx, post)


async def fetch_post(ctx, md5) -> Optional[dict]:
    rows = await ctx.db.execute_fetchall("select data from posts where md5 = ?", (md5,))
    if not rows:
        return None
    assert len(rows) == 1
    return json.loads(rows[0][0])


async def insert_post(ctx, post):
    await ctx.db.execute_insert(
        "insert into posts (md5, filepath ,data) values (?, ?,?)",
        (
            post["md5"],
            Path(urlparse(post["file_url"]).path).name,
            json.dumps(post),
        ),
    )
    await ctx.db.commit()


async def insert_interrogated_result(
    ctx, interrogator: Interrogator, md5: str, tag_string: str
):
    await ctx.db.execute_insert(
        "insert into interrogated_posts (md5, model_name, output_tag_string) values (?,?,?)",
        (md5, interrogator.model_id, tag_string),
    )
    await ctx.db.commit()


async def fight(ctx):
    interrogators = ctx.config.all_available_models

    all_rows = await ctx.db.execute_fetchall("select md5 from posts")
    all_hashes = set(r[0] for r in all_rows)

    for interrogator in interrogators:
        log.info("processing for %r", interrogator)
        # calculate set of images we didn't interrogate yet
        interrogated_rows = await ctx.db.execute_fetchall(
            "select md5 from interrogated_posts where model_name = ?",
            (interrogator.model_id,),
        )
        interrogated_hashes = set(row[0] for row in interrogated_rows)
        missing_hashes = all_hashes - interrogated_hashes

        log.info("missing %d hashes", len(missing_hashes))

        for index, missing_hash in enumerate(missing_hashes):
            log.info(
                "interrogating %r (%d/%d)", missing_hash, index, len(missing_hashes)
            )
            post_filepath = next(DOWNLOADS.glob(f"{missing_hash}*"))
            tag_string = await interrogator.interrogate(ctx.session, post_filepath)
            log.info("got %r", tag_string)
            await insert_interrogated_result(
                ctx, interrogator, missing_hash, tag_string
            )


def score(danbooru_tags: Set[str], interrogator_tags: Set[str]) -> decimal.Decimal:
    tags_in_danbooru = danbooru_tags | interrogator_tags
    tags_not_in_danbooru = danbooru_tags - interrogator_tags
    return decimal.Decimal(
        len(tags_in_danbooru) - len(tags_not_in_danbooru)
    ) / decimal.Decimal(len(danbooru_tags))


async def scores(ctx):
    interrogators = ctx.config.all_available_models

    all_rows = await ctx.db.execute_fetchall("select md5 from posts")
    all_hashes = set(r[0] for r in all_rows)

    absolute_scores = defaultdict(decimal.Decimal)

    for md5_hash in all_hashes:
        log.info("processing for %r", md5_hash)
        post = await fetch_post(ctx, md5_hash)
        danbooru_tags = set(post["tag_string"].split())
        for interrogator in interrogators:
            rows = await ctx.db.execute_fetchall(
                "select output_tag_string from interrogated_posts where md5 = ? and model_name = ?",
                (md5_hash, interrogator.model_id),
            )
            assert len(rows) == 1  # run 'fight' mode, there's missing posts
            tag_string = rows[0][0]
            interrogator_tags = set(tag_string.split())
            absolute_scores[interrogator.model_id] += score(
                danbooru_tags, interrogator_tags
            )

    normalized_scores = {
        model: absolute_scores[model] / len(all_hashes) for model in absolute_scores
    }

    for model in sorted(
        normalized_scores.keys(),
        key=lambda model: normalized_scores[model],
        reverse=True,
    ):
        print(model, normalized_scores[model])


async def realmain(ctx):
    await ctx.db.executescript(
        """
        create table if not exists posts (
         md5 text primary key,
         filepath text,
         data text
        );
        create table if not exists interrogated_posts (
         md5 text,
         model_name text not null,
         output_tag_string text not null,
         primary key (md5, model_name)
        );
        """
    )

    try:
        mode = sys.argv[1]
    except IndexError:
        raise Exception("must have mode")

    if mode == "download_images":
        await download_images(ctx)
    elif mode == "fight":
        await fight(ctx)
    elif mode == "scores":
        await scores(ctx)
    else:
        raise AssertionError(f"invalid mode {mode}")


async def main():
    CONFIG_PATH = Path.cwd() / "config.json"
    log.info("hewo")
    async with aiosqlite.connect("./data.db") as db:
        async with aiohttp.ClientSession() as session:
            with CONFIG_PATH.open() as config_fd:
                config_json = json.load(config_fd)

            config = Config(**config_json)
            ctx = Context(db, session, config)
            await realmain(ctx)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    asyncio.run(main())