#!/usr/bin/env python3
import os
import sys
from functools import lru_cache
from typing import Literal

import psycopg2
from pydantic import BaseModel

from alfred_utils import AlfredOutput, print_debug, print_items, print_out
from extract_uuid import extract_uuid
from hook_cache import HookCache
from rate_limiter import RateLimiter

# Initialize cache with permanent storage
cache = HookCache()

# Initialize rate limiter (1 second minimum between DB queries)
rate_limiter = RateLimiter(min_interval=1.0)


class HookLookup(BaseModel):
    hook_id: str
    original_clip_id: str
    video_upload_id: str
    file_ext: str
    env: Literal["prod", "staging"] = "staging"

    def get_base_url(self) -> str:
        if self.env == "prod":
            return "https://suno.com/"
        else:
            return "https://b.suno.fm/"


def make_raw_upload_result(upload_id: str, file_ext: str | None = None) -> AlfredOutput:
    subtitle = "View raw video upload in S3"
    if not file_ext:
        subtitle = (
            "View raw video upload in S3 (unknown file extension, defaulting to mp4)"
        )
        file_ext = "mp4"

    url = f"https://us-east-1.console.aws.amazon.com/s3/object/suno-uploads?region=us-east-1&bucketType=general&prefix=raw_uploads/{upload_id}.{file_ext}"
    return AlfredOutput(
        title="Open Raw Video Upload",
        subtitle=subtitle,
        arg=url,
        valid=True,
    )


def make_processed_video_result(upload_id: str) -> AlfredOutput:
    url = f"https://cdn1.suno.ai/video_upload_{upload_id}.mp4"
    return AlfredOutput(
        title="Open Processed Video",
        subtitle="View processed video in CDN",
        arg=url,
        valid=True,
    )


def make_rendered_hook_result(hook_id: str) -> AlfredOutput:
    url = f"https://cdn1.suno.ai/hook_{hook_id}.mp4"
    return AlfredOutput(
        title="Open Rendered Hook in CDN",
        subtitle="View rendered hook in CDN",
        arg=url,
        valid=True,
    )


def make_hook_app_result(hook_lookup: HookLookup) -> AlfredOutput:
    url = f"{hook_lookup.get_base_url()}/hook/{hook_lookup.hook_id}"
    return AlfredOutput(
        title="Open Hook on Web",
        subtitle="View the hook live in the app",
        arg=url,
        valid=True,
    )


def make_original_clip_result(hook_lookup: HookLookup) -> AlfredOutput:
    return AlfredOutput(
        title="Open Original Clip",
        subtitle="View the original clip",
        arg=f"{hook_lookup.get_base_url()}/song/{hook_lookup.original_clip_id}",
        valid=True,
    )


def make_original_clip_cdn_result(hook_lookup: HookLookup) -> AlfredOutput:
    return AlfredOutput(
        title="Download Original Clip from CDN",
        subtitle="View the original clip",
        arg=f"https://cdn1.suno.ai/{hook_lookup.original_clip_id}.mp4",
        valid=True,
    )


def make_copy_upload_id_result(hook_lookup: HookLookup) -> AlfredOutput:
    return AlfredOutput(
        title="Copy Upload ID",
        subtitle="Copy the upload ID to the clipboard",
        arg=f"{hook_lookup.video_upload_id}",
        valid=True,
        variables={"action": "copy"},
    )


def make_remove_cache_result(hook_id: str) -> AlfredOutput:
    return AlfredOutput(
        title="Remove from Cache",
        subtitle="Clear this hook ID from the cache",
        arg=f"{hook_id}",
        valid=True,
        variables={"action": "clear-cache"},
    )


class ReturnException(Exception):
    def __init__(self, title: str, subtitle: str, valid: bool, error: str):
        self.title = title
        self.subtitle = subtitle
        self.valid = valid
        self.error = error
        super().__init__(title)

    def print_out(self) -> None:
        print_out(self.title, self.subtitle, self.valid, self.error)


@lru_cache(maxsize=128)
def lookup_hook_video_upload(hook_id: str, database_url: str) -> HookLookup | None:
    conn = None
    cur = None
    try:
        # Connect to database
        conn = psycopg2.connect(database_url)
        cur = conn.cursor()

        # Query for hook data
        cur.execute(
            """
            SELECT vh.id, vh.original_clip, vu.id, vu.original_file_ext
            FROM video_videohook vh
            JOIN bots_videoupload vu ON vh.raw_video_upload_id = vu.id
            WHERE vh.id = %s
        """,
            (hook_id,),
        )

        result = cur.fetchone()

        if result:
            hook_id, original_clip_id, upload_id, file_ext = result
            return HookLookup(
                hook_id=hook_id,
                original_clip_id=original_clip_id,
                video_upload_id=upload_id,
                file_ext=file_ext,
            )
    except Exception as e:
        # Log error for debugging but still return None to try next database
        print_debug(f"Database error: {e}")
        return None
    finally:
        if cur:
            cur.close()
        if conn:
            conn.close()

    return None


def print_results(hook_lookup: HookLookup) -> None:
    raw_upload_result = make_raw_upload_result(
        hook_lookup.video_upload_id, hook_lookup.file_ext
    )
    processed_video_result = make_processed_video_result(hook_lookup.video_upload_id)
    rendered_hook_result = make_rendered_hook_result(hook_lookup.hook_id)
    hook_app_result = make_hook_app_result(hook_lookup)
    original_clip_result = make_original_clip_result(hook_lookup)
    original_clip_cdn_result = make_original_clip_cdn_result(hook_lookup)
    copy_upload_id_result = make_copy_upload_id_result(hook_lookup)
    remove_cache_result = make_remove_cache_result(hook_lookup.hook_id)

    print_items(
        [
            raw_upload_result,
            processed_video_result,
            rendered_hook_result,
            hook_app_result,
            original_clip_result,
            original_clip_cdn_result,
            copy_upload_id_result,
            remove_cache_result,
        ]
    )


def print_not_found_with_cache_option(hook_id: str, is_cached: bool = False) -> None:
    """Print not found message with option to remove from cache."""
    title = "No video upload found for this hook ID"
    if is_cached:
        title += " (cached)"

    print_items(
        [
            AlfredOutput(
                title=title,
                subtitle="Check the hook ID and try again",
                valid=False,
                arg=None,
            ),
            make_remove_cache_result(hook_id),
        ]
    )


def handle_hook_lookup(hook_id: str) -> None:
    uuids = extract_uuid(hook_id)
    if not uuids:
        print_out(
            title="No UUID found",
            subtitle="No valid UUID found in the input text",
            valid=False,
        )
        return
    hook_id = uuids[0]

    # Check cache first
    cache_key = f"hook_{hook_id}"
    cached_result = cache.get(cache_key)

    if cached_result is not None:
        try:
            # Cache hit - cached_result is a dict with 'found' key
            if cached_result.get("found"):
                # Reconstruct HookLookup from dict
                hook_data = cached_result["hook_lookup"]
                if hook_data:
                    hook_lookup = HookLookup(**hook_data)
                    print_results(hook_lookup)
            else:
                print_not_found_with_cache_option(hook_id, is_cached=True)
            return
        except Exception as e:
            # If deserialization fails, clear the cache entry and continue with fresh lookup
            print_debug(f"Cache deserialization failed: {e}, clearing cache entry")
            cache.remove(cache_key)

    # Not in cache, do database lookup
    # Apply rate limiting before database queries
    if rate_limiter.wait_if_needed():
        print_debug("Rate limit applied")

    hook_lookup = lookup_hook_video_upload(hook_id, os.getenv("DATABASE_URL_PROD"))
    print_debug(f"hook_lookup: {hook_lookup}")
    if not hook_lookup:
        print_debug("not found in prod, looking in staging")
        hook_lookup = lookup_hook_video_upload(
            hook_id, os.getenv("DATABASE_URL_STAGING")
        )
        if not hook_lookup:
            print_debug("not found in staging")
            print_not_found_with_cache_option(hook_id, is_cached=False)
            return
        hook_lookup.env = "staging"
    else:
        hook_lookup.env = "prod"

    # Record that we made database queries
    rate_limiter.record_operation()

    # Save to cache with metadata about whether it was found
    # Convert Pydantic model to dict for serialization
    cache_value = {
        "found": bool(hook_lookup),
        "hook_lookup": hook_lookup.model_dump() if hook_lookup else None,
    }
    cache.set(cache_key, cache_value)

    if hook_lookup:
        print_results(hook_lookup)
    else:
        print_not_found_with_cache_option(hook_id, is_cached=False)


def main():
    if len(sys.argv) < 2:
        print_out(
            title="Enter a hook ID",
            subtitle="Type a hook ID to lookup",
            valid=False,
        )
        return

    hook_id = sys.argv[1].strip()

    if not hook_id:
        print_out(
            title="Enter a hook ID",
            subtitle="Type a hook ID to lookup",
            valid=False,
        )
        return

    try:
        handle_hook_lookup(hook_id)
    except ReturnException as e:
        e.print_out()


if __name__ == "__main__":
    main()