import base64
import os
import shutil
from collections import defaultdict
from datetime import date, datetime, timedelta
from io import BytesIO

import dotenv
from datasets import load_dataset
from dateutil.parser import parse
from dateutil.tz import tzutc
from fasthtml.common import *
from huggingface_hub import login, whoami

dotenv.load_dotenv()

style = Style("""
                .grid { margin-bottom: 1rem; }
                .card { display: flex; flex-direction: column; }
                .card img { margin-bottom: 0.5rem; }
                .card h5 { margin: 0; font-size: 0.9rem; line-height: 1.2; }
                .card a { color: inherit; text-decoration: none; }
                .card a:hover { text-decoration: underline; }
            """)

# delete data folder
if os.path.exists("data"):
    try:
        shutil.rmtree("data")
    except OSError as e:
        print("Error: %s : %s" % ("data", e.strerror))

app, rt = fast_app(html_style=(style,))

login(token=os.environ.get("HF_TOKEN"))

hf_user = whoami(os.environ.get("HF_TOKEN"))["name"]
HF_REPO_ID_TXT = f"{hf_user}/zotero-answer-ai-texts"
HF_REPO_ID_IMG = f"{hf_user}/zotero-answer-ai-images"

abstract_ds = load_dataset(HF_REPO_ID_TXT, "abstracts", split="train")
article_ds = load_dataset(HF_REPO_ID_TXT, "articles", split="train")

image_ds = load_dataset(HF_REPO_ID_IMG, "images_first_page", split="train")


def parse_date(date_string):
    try:
        return parse(date_string).astimezone(tzutc()).date()
    except ValueError:
        return date.today()


def get_week_start(date_obj):
    return date_obj - timedelta(days=date_obj.weekday())


week2articles = defaultdict(list)
for article in article_ds:
    date_added = parse_date(article["date_added"])
    week_start = get_week_start(date_added)
    week2articles[week_start].append(article["arxiv_id"])

weeks = sorted(week2articles.keys(), reverse=True)

arxiv2article = {article["arxiv_id"]: article for article in article_ds}
arxiv2abstract = {abstract["arxiv_id"]: abstract for abstract in abstract_ds}
arxiv2image = {image["arxiv_id"]: image for image in image_ds}


def get_article_details(arxiv_id):
    article = arxiv2article.get(arxiv_id, {})
    abstract = arxiv2abstract.get(arxiv_id, {})
    image = arxiv2image.get(arxiv_id, {})
    return article, abstract, image


def generate_week_content(current_week):
    week_index = weeks.index(current_week)
    prev_week = weeks[week_index + 1] if week_index < len(weeks) - 1 else None
    next_week = weeks[week_index - 1] if week_index > 0 else None

    nav_buttons = Group(
        Button(
            "← Previous Week",
            hx_get=f"/week/{prev_week}" if prev_week else "#",
            hx_target="#content",
            hx_swap="innerHTML",
            disabled=not prev_week,
        ),
        Button(
            "Next Week →",
            hx_get=f"/week/{next_week}" if next_week else "#",
            hx_target="#content",
            hx_swap="innerHTML",
            disabled=not next_week,
        ),
    )

    articles = week2articles[current_week]
    article_cards = []
    for arxiv_id in articles:
        article, abstract, image = get_article_details(arxiv_id)
        article_title = article["contents"][0].get("paper_title", "article") if article["contents"] else "article"

        card_content = [
            H5(
                A(
                    article_title,
                    href=f"https://arxiv.org/abs/{arxiv_id}",
                    target="_blank",
                )
            )
        ]

        if image:
            pil_image = image["image"]  # image[0]["image"]
            img_byte_arr = BytesIO()
            pil_image.save(img_byte_arr, format="JPEG")
            img_byte_arr = img_byte_arr.getvalue()
            image_url = f"data:image/jpeg;base64,{base64.b64encode(img_byte_arr).decode('utf-8')}"
            card_content.insert(
                0,
                Img(
                    src=image_url,
                    alt="Article image",
                    style="max-width: 100%; height: auto; margin-bottom: 15px;",
                ),
            )

        article_cards.append(Card(*card_content, cls="mb-4"))

    grid = Grid(
        *article_cards,
        style="display: grid; grid-template-columns: repeat(3, 1fr); gap: 1rem;",
    )

    week_end = current_week + timedelta(days=6)
    return Div(
        nav_buttons,
        H3(f"Week of {current_week.strftime('%B %d')} - {week_end.strftime('%B %d, %Y')} ({len(articles)} articles)"),
        grid,
        nav_buttons,
        id="content",
    )


@rt("/")
def get():
    return Titled("AnswerAI Zotero Weekly", generate_week_content(weeks[0]))


@rt("/week/{date}")
def get(date: str):
    try:
        current_week = datetime.strptime(date, "%Y-%m-%d").date()
        return generate_week_content(current_week)
    except Exception as e:
        return Div(f"Error displaying articles: {str(e)}")


# serve()

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))