import gradio as gr
from gradio_client import Client, handle_file
from google import genai
from google.genai import types
import os
from typing import Optional, List, Tuple
import requests
from PIL import Image
from io import BytesIO
import tempfile
import ffmpeg
import sqlite3
from datetime import datetime, date
from pathlib import Path
from threading import Lock

# --- Database Setup ---
DATA_DIR = Path("/data")
DATA_DIR.mkdir(exist_ok=True)
DB_PATH = DATA_DIR / "usage_limits.db"

DAILY_LIMIT = 40
EXEMPTED_USERS = ["multimodalart"]
db_lock = Lock()

def init_db():
    """Initialize the SQLite database (single global usage pool)."""
    print(f"Initializing database at: {DB_PATH}")
    try:
        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='usage'")
            table_exists = cursor.fetchone()

            if not table_exists:
                cursor.execute('''
                    CREATE TABLE usage (
                        username TEXT PRIMARY KEY,
                        date TEXT NOT NULL,
                        count INTEGER NOT NULL DEFAULT 0
                    )
                ''')
                conn.commit()
                print("Database initialized successfully")
                return

            cursor.execute("PRAGMA table_info(usage)")
            columns = [col[1] for col in cursor.fetchall()]

            if 'count_standard' in columns and 'count_pro' in columns:
                print("Migrating to single-pool schema (count_standard + count_pro → count)...")
                cursor.execute('''
                    CREATE TABLE usage_new (
                        username TEXT PRIMARY KEY,
                        date TEXT NOT NULL,
                        count INTEGER NOT NULL DEFAULT 0
                    )
                ''')
                cursor.execute('''
                    INSERT INTO usage_new (username, date, count)
                    SELECT username, date, count_standard + count_pro FROM usage
                ''')
                cursor.execute("DROP TABLE usage")
                cursor.execute("ALTER TABLE usage_new RENAME TO usage")
                conn.commit()
                print("Migration completed successfully")
            elif 'count' in columns:
                print("Database schema is already up to date")
            else:
                print("Unknown schema — recreating table...")
                cursor.execute("DROP TABLE usage")
                cursor.execute('''
                    CREATE TABLE usage (
                        username TEXT PRIMARY KEY,
                        date TEXT NOT NULL,
                        count INTEGER NOT NULL DEFAULT 0
                    )
                ''')
                conn.commit()
    except Exception as e:
        print(f"Error initializing database: {e}")
        import traceback
        traceback.print_exc()

def check_and_update_usage(username: str, credits_to_use: int = 1) -> bool:
    """
    Check if user has enough credits for today and update usage.
    Returns True if user can generate, False if limit reached.
    """
    if username in EXEMPTED_USERS:
        print(f"User {username} is exempted from rate limits")
        return True

    with db_lock:
        try:
            with sqlite3.connect(DB_PATH) as conn:
                today = str(date.today())
                cursor = conn.cursor()

                cursor.execute("SELECT date, count FROM usage WHERE username = ?", (username,))
                result = cursor.fetchone()

                if result is None:
                    cursor.execute("INSERT INTO usage (username, date, count) VALUES (?, ?, ?)",
                                   (username, today, credits_to_use))
                    conn.commit()
                    print(f"New user {username}: {credits_to_use}/{DAILY_LIMIT}")
                    return True

                user_date, user_count = result

                if user_date != today:
                    cursor.execute("UPDATE usage SET date = ?, count = ? WHERE username = ?",
                                   (today, credits_to_use, username))
                    conn.commit()
                    print(f"User {username} reset for new day: {credits_to_use}/{DAILY_LIMIT}")
                    return True

                if user_count + credits_to_use > DAILY_LIMIT:
                    print(f"User {username} insufficient credits: needs {credits_to_use}, has {DAILY_LIMIT - user_count}/{DAILY_LIMIT} remaining")
                    return False

                new_count = user_count + credits_to_use
                cursor.execute("UPDATE usage SET count = ? WHERE username = ?", (new_count, username))
                conn.commit()
                print(f"User {username} usage: {new_count}/{DAILY_LIMIT} (used {credits_to_use} credits)")
                return True

        except Exception as e:
            print(f"Error checking usage for {username}: {e}")
            import traceback
            traceback.print_exc()
            # On error, allow the request (fail open)
            return True

def get_remaining_generations(username: str) -> int:
    """Get the number of remaining credits for today."""
    if username in EXEMPTED_USERS:
        return 999999

    with db_lock:
        try:
            with sqlite3.connect(DB_PATH) as conn:
                today = str(date.today())
                cursor = conn.cursor()

                cursor.execute("SELECT date, count FROM usage WHERE username = ?", (username,))
                result = cursor.fetchone()

                if result is None:
                    return DAILY_LIMIT

                user_date, user_count = result

                if user_date != today:
                    return DAILY_LIMIT

                return max(0, DAILY_LIMIT - user_count)
        except Exception as e:
            print(f"Error getting remaining generations for {username}: {e}")
            return DAILY_LIMIT

# Initialize database on module load
init_db()

# --- Google Gemini API Configuration ---
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")
if not GOOGLE_API_KEY:
    raise ValueError("GOOGLE_API_KEY environment variable not set.")
client = genai.Client(api_key=os.environ.get("GOOGLE_API_KEY"))

# Added Nano Banana 2 (3.1 flash) model definition
GEMINI_MODEL_NAME = 'gemini-2.5-flash-image'
GEMINI_PRO_MODEL_NAME = 'gemini-3-pro-image-preview'
GEMINI_NANO_BANANA_2_MODEL_NAME = 'gemini-3.1-flash-image-preview'

HF_USERINFO_URL = "https://huggingface.co/oauth/userinfo"

def get_user_info(token: Optional[gr.OAuthToken]) -> Optional[dict]:
    """Fetch {username, is_pro} from the OAuth userinfo endpoint in a single call.
    Returns None if the token is missing/invalid or the request fails."""
    if not token:
        return None
    token_str = token.token if isinstance(token, gr.OAuthToken) else token
    try:
        resp = requests.get(
            HF_USERINFO_URL,
            headers={"Authorization": f"Bearer {token_str}"},
            timeout=10,
        )
        resp.raise_for_status()
        data = resp.json()
        is_pro = (
            data.get("isPro", False)
            or any(org.get("plan") in ("team", "enterprise") for org in data.get("orgs", []))
        )
        return {
            "username": data.get("preferred_username"),
            "is_pro": is_pro,
        }
    except Exception as e:
        print(f"Could not fetch OAuth userinfo: {e}")
        return None

def get_credit_cost(resolution: str) -> int:
    """Get the credit cost for a given resolution."""
    if "4K" in resolution:
        return 4
    elif "2K" in resolution:
        return 2
    else:  # 1K
        return 1

def get_resolution_value(resolution: str) -> str:
    """Extract the resolution value from the dropdown selection."""
    if "4K" in resolution:
        return "4K"
    elif "2K" in resolution:
        return "2K"
    else:
        return "1K"

def _extract_image_data_from_response(response) -> Optional[bytes]:
    """Helper to extract image data from the model's response."""
    if hasattr(response, 'candidates') and response.candidates:
        for part in response.candidates[0].content.parts:
            if hasattr(part, 'inline_data') and hasattr(part.inline_data, 'data'):
                return part.inline_data.data
    return None

def _get_video_info(video_path: str) -> Tuple[float, Tuple[int, int]]:
    """Instantly gets the framerate and (width, height) of a video using ffprobe."""
    probe = ffmpeg.probe(video_path)
    video_stream = next((s for s in probe['streams'] if s['codec_type'] == 'video'), None)
    if not video_stream:
        raise ValueError("No video stream found in the file.")
    framerate = eval(video_stream['avg_frame_rate'])
    resolution = (int(video_stream['width']), int(video_stream['height']))
    return framerate, resolution

def _resize_image(image_path: str, target_size: Tuple[int, int]) -> str:
    """Resizes an image to a target size and saves it to a new temp file."""
    with Image.open(image_path) as img:
        if img.size == target_size:
            return image_path
        resized_img = img.resize(target_size, Image.Resampling.LANCZOS)
        suffix = os.path.splitext(image_path)[1] or ".png"
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
            resized_img.save(tmp_file.name)
            return tmp_file.name

def _trim_first_frame_fast(video_path: str) -> str:
    """Removes exactly the first frame of a video without re-encoding."""
    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_output_file:
        output_path = tmp_output_file.name
    try:
        framerate, _ = _get_video_info(video_path)
        if framerate == 0: raise ValueError("Framerate cannot be zero.")
        start_time = 1 / framerate
        (
            ffmpeg
            .input(video_path, ss=start_time)
            .output(output_path, c='copy', avoid_negative_ts='make_zero')
            .run(overwrite_output=True, quiet=True)
        )
        return output_path
    except Exception as e:
        raise RuntimeError(f"FFmpeg trim error: {e}")

def _combine_videos_simple(video1_path: str, video2_path: str) -> str:
    """Combines two videos using the fast concat demuxer."""
    with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix=".txt") as tmp_list_file:
        tmp_list_file.write(f"file '{os.path.abspath(video1_path)}'\n")
        tmp_list_file.write(f"file '{os.path.abspath(video2_path)}'\n")
        list_file_path = tmp_list_file.name
    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_output_file:
        output_path = tmp_output_file.name
    try:
        (
            ffmpeg
            .input(list_file_path, format='concat', safe=0)
            .output(output_path, c='copy')
            .run(overwrite_output=True, quiet=True)
        )
        return output_path
    except ffmpeg.Error as e:
        raise RuntimeError(f"FFmpeg combine error: {e.stderr.decode()}")
    finally:
        if os.path.exists(list_file_path):
            os.remove(list_file_path)

def _generate_video_segment(input_image_path: str, output_image_path: str, prompt: str, token: str) -> str:
    """Generates a single video segment using the external service."""
    video_client = Client("multimodalart/wan-2-2-first-last-frame", token=token)
    result = video_client.predict(
        start_image_pil=handle_file(input_image_path),
        end_image_pil=handle_file(output_image_path),
        prompt=prompt, api_name="/generate_video"
    )
    return result[0]["video"]

def unified_image_generator(prompt: str, images: Optional[List[str]], previous_video_path: Optional[str], last_frame_path: Optional[str], aspect_ratio: str, model_selection: str, resolution: str, oauth_token: Optional[gr.OAuthToken]) -> tuple:
    user_info = get_user_info(oauth_token)
    if not user_info or not user_info["is_pro"]:
        raise gr.Error("Access Denied.")
    username = user_info["username"]
    if not username:
        raise gr.Error("Could not identify user.")

    use_pro_model = (model_selection == "Nano Banana PRO")
    use_banana_2 = (model_selection == "Nano Banana 2")

    # Credit cost: Banana = 1, Banana 2 = 1/2/4 by resolution, PRO = 2× resolution (2/4/8)
    if use_pro_model:
        credits_to_use = get_credit_cost(resolution) * 2
    elif use_banana_2:
        credits_to_use = get_credit_cost(resolution)
    else:
        credits_to_use = 1

    can_generate = check_and_update_usage(username, credits_to_use)

    if not can_generate:
        remaining = get_remaining_generations(username)
        if (use_pro_model or use_banana_2) and remaining > 0 and remaining < credits_to_use:
            gr.Info(f"You need {credits_to_use} credits for {get_resolution_value(resolution)} but only have {remaining} remaining. Try a lower resolution or use Nano Banana (1 credit/gen).")
        raise gr.Error(f"Insufficient credits. You need {credits_to_use} but have {remaining}/{DAILY_LIMIT} remaining today.")
    
    try:
        contents = [Image.open(image_path[0]) for image_path in images] if images else []
        contents.append(prompt)
        
        # Select API model identifier based on selection
        if use_pro_model:
            model_name_api = GEMINI_PRO_MODEL_NAME
        elif use_banana_2:
            model_name_api = GEMINI_NANO_BANANA_2_MODEL_NAME
        else:
            model_name_api = GEMINI_MODEL_NAME
        
        # Create config with aspect ratio and resolution
        if use_pro_model or use_banana_2:
            resolution_value = get_resolution_value(resolution)
            
            img_config_kwargs = {"image_size": resolution_value}
            if aspect_ratio != "Auto":
                img_config_kwargs["aspect_ratio"] = aspect_ratio
                
            gen_config_kwargs = {
                "response_modalities": ["IMAGE", "TEXT"],
                "image_config": types.ImageConfig(**img_config_kwargs)
            }
            
            # Apply Thinking Config uniquely for Nano Banana 2
            if use_banana_2:
                gen_config_kwargs["thinking_config"] = types.ThinkingConfig(thinking_level="MINIMAL")
                
            generate_content_config = types.GenerateContentConfig(**gen_config_kwargs)
        else:
            # Standard model: only aspect_ratio
            if aspect_ratio == "Auto":
                generate_content_config = types.GenerateContentConfig(
                    response_modalities=["IMAGE", "TEXT"],
                )
            else:
                generate_content_config = types.GenerateContentConfig(
                    response_modalities=["IMAGE", "TEXT"],
                    image_config=types.ImageConfig(
                        aspect_ratio=aspect_ratio,
                    ),
                )
                
        print(f"Generating image for user {username} with prompt {prompt} using {model_name_api}")
        response = client.models.generate_content(
            model=model_name_api, 
            contents=contents,
            config=generate_content_config
        )
        image_data = _extract_image_data_from_response(response)
        if not image_data: raise gr.Error("No image data in response")
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
            Image.open(BytesIO(image_data)).save(tmp.name)
            output_path = tmp.name
            
        can_create_video = bool(images and len(images) == 1)
        can_extend_video = False
        if can_create_video and previous_video_path and last_frame_path:
            # The crucial check for continuity
            if images[0][0] == last_frame_path:
                can_extend_video = True
                
        print(f"Image generated at {output_path}")
        return (output_path, gr.update(visible=can_create_video), gr.update(visible=can_extend_video), gr.update(visible=False))
    except Exception as e:
        raise gr.Error(f"Image generation failed: {e}. Rephrase your prompt to make image generation explicit and try again")

def create_new_video(input_image_gallery: List[str], prompt_input: str, output_image: str, oauth_token: Optional[gr.OAuthToken]) -> tuple:
    user_info = get_user_info(oauth_token)
    if not user_info or not user_info["is_pro"]: raise gr.Error("Access Denied.")
    if not input_image_gallery or not output_image: raise gr.Error("Input/output images required.")
    try:
        new_segment_path = _generate_video_segment(input_image_gallery[0][0], output_image, prompt_input, oauth_token.token)
        return new_segment_path, new_segment_path, output_image
    except Exception as e:
        raise gr.Error(f"Video creation failed: {e}")

def extend_existing_video(input_image_gallery: List[str], prompt_input: str, output_image: str, previous_video_path: str, oauth_token: Optional[gr.OAuthToken]) -> tuple:
    user_info = get_user_info(oauth_token)
    if not user_info or not user_info["is_pro"]: raise gr.Error("Access Denied.")
    if not previous_video_path: raise gr.Error("No previous video to extend.")
    if not input_image_gallery or not output_image: raise gr.Error("Input/output images required.")
    try:
        _, target_resolution = _get_video_info(previous_video_path)
        resized_input_path = _resize_image(input_image_gallery[0][0], target_resolution)
        resized_output_path = _resize_image(output_image, target_resolution)
        new_segment_path = _generate_video_segment(resized_input_path, resized_output_path, prompt_input, oauth_token.token)
        trimmed_segment_path = _trim_first_frame_fast(new_segment_path)
        final_video_path = _combine_videos_simple(previous_video_path, trimmed_segment_path)
        return final_video_path, final_video_path, output_image
    except Exception as e:
        raise gr.Error(f"Video extension failed: {e}")

css = '''
#sub_title{margin-top: -15px !important}
.tab-wrapper{margin-bottom: -33px !important}
.tabitem{padding: 0px !important}
.fillable{max-width: 980px !important}
.dark .progress-text {color: white}
.logo-dark{display: none}
.dark .logo-dark{display: block !important}
.dark .logo-light{display: none}
.grid-container img{object-fit: contain}
.grid-container {display: grid;grid-template-columns: repeat(2, 1fr)}
.grid-container:has(> .gallery-item:only-child) {grid-template-columns: 1fr}
#wan_ad p{text-align: center;padding: .5em}
'''

with gr.Blocks() as demo:
    gr.HTML('''
    <img class="logo-dark" src='https://huggingface.co/spaces/multimodalart/nano-banana/resolve/main/nano_banana_pros.png' style='margin: 0 auto; max-width: 650px' />
    <img class="logo-light" src='https://huggingface.co/spaces/multimodalart/nano-banana/resolve/main/nano_banana_pros_light.png' style='margin: 0 auto; max-width: 650px' />
    ''')
    gr.HTML("<h3 style='text-align:center'>Hugging Face PRO users can use Google's Nano Banana and Nano Banana PRO on this Space. <a href='http://huggingface.co/subscribe/pro?source=nana_banana' target='_blank'>Subscribe to PRO</a></h3>", elem_id="sub_title")
    pro_message = gr.Markdown(visible=False)
    main_interface = gr.Column(visible=False)
    
    previous_video_state = gr.State(None)
    last_frame_of_video_state = gr.State(None)

    with main_interface:
        with gr.Row():
            with gr.Column(scale=1):
                image_input_gallery = gr.Gallery(label="Upload one or more images here. Leave empty for text-to-image", file_types=["image"], height="auto")
                prompt_input = gr.Textbox(label="Prompt", placeholder="Turns this photo into a masterpiece")
                
                # Added Nano Banana 2 choice here
                model_radio = gr.Radio(
                    choices=["Nano Banana", "Nano Banana 2", "Nano Banana PRO"],
                    value="Nano Banana 2",
                    label="Model",
                )
                
                with gr.Row():
                    aspect_ratio_dropdown = gr.Dropdown(
                        label="Aspect Ratio",
                        choices=["Auto", "1:1", "9:16", "16:9", "3:4", "4:3", "3:2", "2:3", "5:4", "4:5", "21:9"],
                        value="Auto",
                        interactive=True
                    )
                    
                    resolution_dropdown = gr.Dropdown(
                        label="Resolution",
                        choices=["1K", "2K", "4K"],
                        value="1K",
                        interactive=True,
                        visible=True
                    )
                
                generate_button = gr.Button("Generate", variant="primary")
            with gr.Column(scale=1):
                output_image = gr.Image(label="Output", interactive=False, elem_id="output", type="filepath")
                use_image_button = gr.Button("♻️ Use this Image for Next Edit", variant="primary")
                with gr.Row():
                    create_video_button = gr.Button("Create video between the two images 🎥", variant="secondary", visible=False)
                    extend_video_button = gr.Button("Extend existing video with new scene 🎞️", variant="secondary", visible=False)
                with gr.Group(visible=False) as video_group:
                    video_output = gr.Video(label="Generated Video", buttons=["download"], autoplay=True)
                    gr.Markdown("Generate more with [Wan 2.2 first-last-frame](https://huggingface.co/spaces/multimodalart/wan-2-2-first-last-frame)", elem_id="wan_ad")
        gr.Markdown("<h2 style='text-align: center'>Thank you for being a PRO! 🤗</h2>")

    login_button = gr.LoginButton()

    # Show/hide resolution dropdown based on model selection (visible for PRO and Nano Banana 2)
    def update_resolution_visibility(model_selection):
        return gr.update(visible=(model_selection in ["Nano Banana PRO", "Nano Banana 2"]))
    
    model_radio.change(
        fn=update_resolution_visibility,
        inputs=[model_radio],
        outputs=[resolution_dropdown],
        api_visibility="private"
    )

    gr.on(
        triggers=[generate_button.click, prompt_input.submit],
        fn=unified_image_generator,
        inputs=[prompt_input, image_input_gallery, previous_video_state, last_frame_of_video_state, aspect_ratio_dropdown, model_radio, resolution_dropdown],
        outputs=[output_image, create_video_button, extend_video_button, video_group],
        api_visibility="private"
    )
    use_image_button.click(
        fn=lambda img: (
            [img] if img else None, None, gr.update(visible=False),
            gr.update(visible=False), gr.update(visible=False)
        ),
        inputs=[output_image],
        outputs=[image_input_gallery, output_image, create_video_button, extend_video_button, video_group],
        api_visibility="private"
    )
    create_video_button.click(
        fn=lambda: gr.update(visible=True), outputs=[video_group],
        api_visibility="private"
    ).then(
        fn=create_new_video,
        inputs=[image_input_gallery, prompt_input, output_image],
        outputs=[video_output, previous_video_state, last_frame_of_video_state],
        api_visibility="private"
    )
    extend_video_button.click(
        fn=lambda: gr.update(visible=True), outputs=[video_group],
        api_visibility="private"
    ).then(
        fn=extend_existing_video,
        inputs=[image_input_gallery, prompt_input, output_image, previous_video_state],
        outputs=[video_output, previous_video_state, last_frame_of_video_state],
        api_visibility="private"
    )

    def control_access(profile: Optional[gr.OAuthProfile] = None, oauth_token: Optional[gr.OAuthToken] = None):
        if not profile: return gr.update(visible=False), gr.update(visible=False)
        user_info = get_user_info(oauth_token)
        if user_info and user_info["is_pro"]:
            return gr.update(visible=True), gr.update(visible=False)
        else:
            message = (
                "## ✨ Exclusive Access for PRO Users\n\n"
                "Thank you for your interest! This app is available exclusively for our Hugging Face **PRO** members.\n\n"
                "To unlock this and many other cool stuff, please consider upgrading your account.\n\n"
                "### [**Become a PRO Today!**](http://huggingface.co/subscribe/pro?source=nana_banana)"
            )
            return gr.update(visible=False), gr.update(visible=True, value=message)
    demo.load(control_access, inputs=None, outputs=[main_interface, pro_message], api_visibility="private")

if __name__ == "__main__":
    demo.queue(max_size=None, default_concurrency_limit=None, api_open=False).launch(
        show_error=True,
        theme=gr.themes.Citrus(),
        css=css
    )