import gradio as gr
import os
import random
import modules.constants as constants
import modules.version_info as version_info
import modules.storage as storage

import urllib.request
from urllib.parse import urlparse, parse_qs, urlencode

user_dir = constants.TMPDIR
default_folder = "saved_models/3d_model_" + format(random.randint(1, 999999), "06d")

def _resolve_short_id_to_query_params(query_params: dict) -> dict:
    """
    Checks for a 'sid' (short ID) in query_params. If found, attempts to resolve
    it to a full URL and updates query_params with the parameters ('3d', 'hm', 'image')
    from the resolved URL.
    """
    short_id = query_params.get("sid")
    if short_id:
        status, full_url_from_shortener = storage.gen_full_url(
            short_url=short_id,
            repo_id=constants.HF_REPO_ID,
            json_file=constants.SHORTENER_JSON_FILE
        )
        if status == "success_retrieved_full" and full_url_from_shortener:
            print(f"Retrieved full URL from short ID '{short_id}': {full_url_from_shortener}")
            try:
                parsed_full_url = urlparse(full_url_from_shortener)
                retrieved_params = parse_qs(parsed_full_url.query)
                
                # Update query_params with those from the full_url_from_shortener
                # The parse_qs function returns lists for values, so get the first element.
                # If a param is not in the resolved URL, it will be set to None.
                query_params["3d"] = retrieved_params.get("3d", [None])[0]
                query_params["hm"] = retrieved_params.get("hm", [None])[0]
                query_params["image"] = retrieved_params.get("image", [None])[0]
            except Exception as e:
                print(f"Error parsing full URL from shortener: {e}")
                # Proceed with original query_params if parsing fails (i.e., don't overwrite them with Nones here)
        else:
            print(f"Failed to retrieve full URL for short ID '{short_id}': {status}")
            # If sid resolution fails, original query_params (including potentially 3d, hm, image if passed alongside sid) remain.
    return query_params

def getVersions():
    #return html_versions
    return version_info.versions_html()
    # Process URLs and download files if needed.

def process_url(url, default_ext=".png"):
    """Download file from URL if it's a remote URL and return its local path.
    
    Performs HuggingFace authentication if the URL requires it.
    The caller can pass an appropriate default_ext (e.g. ".glb" for models).
    Uses huggingface_hub library for HuggingFace URLs for better authentication.
    """
    if not url:
        return None
        
    # If it's already a local file, return it.
    if os.path.exists(url) or not (url.startswith('http://') or url.startswith('https://')):
        return url
    
    # Parse URL to get components
    try:
        
        # Create filename from URL.
        parsed_url = urlparse(url)
        filename = os.path.basename(parsed_url.path)
        if not filename:
            filename = f"downloaded_{hash(url) % 10000}.file"
            
        # Add extension if missing.
        ext = os.path.splitext(filename)[1].lower()
        if not ext:
            filename += default_ext
            
        # Create local path.
        local_path = os.path.join(constants.TMPDIR, filename)
        
        # If the file is hosted on HuggingFace, use huggingface_hub
        if 'huggingface.co' in url or 'hf.co' in url:
            try:
                from huggingface_hub import login, hf_hub_download

                # Log in to HuggingFace
                login(token=constants.HF_API_TOKEN)

                # Extract repo information from URL
                # Format: https://huggingface.co/datasets/{repo_id}/resolve/main/{path}
                if '/datasets/' in url and '/resolve/main/' in url:
                    parts = url.split('/datasets/')[1].split('/resolve/main/')
                    repo_id = parts[0]
                    
                    # The remaining path may contain subfolders and filename
                    full_path = parts[1]
                    
                    # Extract the filename and subfolder
                    if '/' in full_path:
                        subfolder, filename = full_path.rsplit('/', 1)
                    else:
                        subfolder = None
                        filename = full_path
                    
                    print(f"Downloading from HF repo '{repo_id}', filename '{filename}', subfolder '{subfolder}'")
                    
                    # Download using huggingface_hub
                    local_path = hf_hub_download(
                        repo_id=repo_id,
                        filename=filename,
                        subfolder=subfolder,
                        repo_type="dataset",
                        local_dir=constants.TMPDIR,
                        local_dir_use_symlinks=False
                    )
                    return local_path
                else:
                    # Fall back to standard download for other HF URLs
                    print("URL format not recognized for huggingface_hub download, falling back to standard method")
            except Exception as e:
                print(f"Error using huggingface_hub download: {e}, falling back to standard method")
        
        # Standard download for non-HF URLs or as fallback
        print(f"Downloading {url} to {local_path}")
        urllib.request.urlretrieve(url, local_path)
        return local_path
            
    except Exception as e:
        print(f"Error downloading file {url}: {e}")
        return url  # Return original URL if download fails

def load_data(request: gr.Request, model_3d, image_slider):
    """
    Load data from query parameters, download files if needed,
    and use current component values as defaults if no query parameters are provided.
    
    If query parameters are provided, generate a permalink using storage.generate_permalink_from_urls.
    
    Parameters:
        request: Gradio request object containing query parameters.
        model_3d: Current value or component for the 3D model.
        image_slider: Current value or component for the image slider.
    
    Returns:
        tuple: (model_url, slider_images, permalink)
               - model_url: processed URL for the 3D model.
               - slider_images: processed list of image URLs.
               - permalink: a generated permalink if query parameters were provided,
                 or an empty string if not.
    """
    # Parse query parameters.
    query_params = dict(request.query_params) if request is not None else {}
    
    # Resolve short ID if present
    query_params = _resolve_short_id_to_query_params(query_params)

    # Extract URLs from query parameters (which may have been updated)
    short_id = query_params.get("sid", None)
    model_url = query_params.get("3d", None)
    hm_url = query_params.get("hm", None)
    img_url = query_params.get("image", None)

    # If 'sid' was passed but didn't resolve, or no params were passed, then model_url, hm_url, img_url will be None.
    has_loadable_params = bool(model_url or hm_url or img_url)

    # Process the model URL if provided.
    if model_url:
        model_url = process_url(model_url, default_ext=".glb")
    
    # Process image URLs if provided.
    slider_images = []
    if img_url:
        local_img = process_url(img_url, default_ext=".png")
        if local_img:
            slider_images.append(local_img)
    if hm_url:
        local_hm = process_url(hm_url, default_ext=".png")
        if local_hm:
            if not slider_images or local_hm != slider_images[0]:
                if len(slider_images) == 1 and img_url:
                    slider_images.append(local_hm)
                else:
                    slider_images = [local_hm] + slider_images
                    slider_images = slider_images[:2]

    
    # Set default values if no URLs provided or processed:
    default_model_val = getattr(model_3d, "value", model_3d)
    default_images_val = getattr(image_slider, "value", image_slider)
    
    if not slider_images:
        slider_images = default_images_val if default_images_val and default_images_val != (None, None) else constants.default_slider_images
        
    if not model_url:
        model_url = default_model_val if default_model_val else constants.default_model_3d

    # If any loadable query parameters were effectively present (either directly or via sid), generate a permalink.
    permalink = ""
    permalink_short = ""
    if has_loadable_params:
        # Generate permalink using the processed URLs.
        try:
            permalink_model_url = query_params.get("3d", model_url) 
            permalink_hm_url = query_params.get("hm", hm_url if len(slider_images) > 1 and hm_url else (slider_images[1] if len(slider_images) > 1 else None) )
            permalink_img_url = query_params.get("image", img_url if slider_images and img_url else (slider_images[0] if slider_images else None) )
            permalink = storage.generate_permalink_from_urls(permalink_model_url, permalink_hm_url, permalink_img_url)
            if not short_id:
                # If no short ID was provided, generate a new one.
                result, short_id = storage.gen_full_url(full_url=permalink, repo_id=constants.HF_REPO_ID, json_file=constants.SHORTENER_JSON_FILE)
            permalink_short = f"{constants.APP_BASE_URL}/?sid={short_id}"
            print(f"Generated permalink: {result} (short ID: {short_id})")
        except Exception as e:
            print(f"Error generating permalink: {e}")


    
    return model_url, slider_images, permalink, permalink_short

def process_upload(files, current_model, current_images):
    """
    Process uploaded files and assign them to the appropriate component based on file extension.
    
    Files with extensions in [".glb", ".gltf", ".obj", ".ply"] are sent to the Model3D component.
    Files with extensions in [".png", ".jpg", ".jpeg"] are sent to the ImageSlider component.
    
    The function merges the uploaded files with current data. If a file for a component is not 
    provided in the upload (i.e. not exactly 1 model file or not exactly 2 image files), then the
    original data will be retained for that component. If an upload is provided, it will replace 
    the corresponding value.
    
    For the ImageSlider, if a single image is provided in the upload, it will update only the first
    image slot, leaving the second slot unchanged.
    """
    extracted_model = None
    extracted_images = []
    
    # Ensure files is a list.
    if not isinstance(files, list):
        files = [files]
        
    for f in files:
        # f can be a file path (string) or an object with attribute `name`
        file_name = f.name if hasattr(f, "name") else f
        ext = os.path.splitext(file_name)[1].lower()
        
        if ext in constants.model_extensions:
            if extracted_model is None:
                extracted_model = file_name
        elif ext in constants.image_extensions:
            if len(extracted_images) < 2:
                extracted_images.append(file_name)
                
    # Merge results with current data.
    updated_model = extracted_model if extracted_model is not None else current_model
    
    # Convert current_images if it's a tuple or a single item.
    if isinstance(current_images, tuple):
        current_images = list(current_images)
    elif current_images is not None and not isinstance(current_images, list):
        current_images = [current_images]
    
    # For the image slider, we expect a list of exactly 2 images.
    # Start with current images (or use defaults if None).
    if current_images is None or not isinstance(current_images, list) or len(current_images) == 0:
        new_images = [None, None]
    else:
        new_images = current_images[:2]
        if len(new_images) < 2:
            new_images.append(None)

    if len(extracted_images) == 1:
        new_images[0] = extracted_images[0]
    elif len(extracted_images) == 2:
        new_images[0] = extracted_images[0]
        new_images[1] = extracted_images[1]
        
    return updated_model, new_images

def get_open_graph_meta_tags(query_params):
    """Generates Open Graph meta tags based on query parameters."""
    og_title = constants.DEFAULT_OG_TITLE
    og_description = constants.DEFAULT_OG_DESCRIPTION
    og_image = constants.DEFAULT_OG_IMAGE_URL
    og_type = constants.DEFAULT_OG_TYPE
    og_url = constants.APP_BASE_URL

    img_url_from_query = query_params.get("image")
    if img_url_from_query:
        og_image = img_url_from_query
        
        model_url_from_query = query_params.get("3d")
        if model_url_from_query:
            try:
                model_filename = os.path.basename(urlparse(model_url_from_query).path)
                if model_filename:
                     og_title = f"3D Viewer: {model_filename}"
                else:
                    og_title = f"Shared 3D Model" 
            except Exception:
                og_title = "Shared 3D Model"
        else:
            og_title = "Shared Image"
    
    if query_params:
        filtered_query_params = {k: v for k, v in query_params.items() if v is not None}
        if filtered_query_params:
             og_url += "?" + urlencode(filtered_query_params)

    meta_tags = f'''
        <meta property="og:title" content="{og_title}" />
        <meta property="og:description" content="{og_description}" />
        <meta property="og:image" content="{og_image}" />
        <meta property="og:url" content="{og_url}" />
        <meta property="og:type" content="{og_type}" />
        <meta name="twitter:card" content="summary_large_image">
        <meta name="twitter:title" content="{og_title}">
        <meta name="twitter:description" content="{og_description}">
        <meta name="twitter:image" content="{og_image}">
    '''
    return meta_tags

placeholder_initial_query_params = {} 
processed_placeholder_params = _resolve_short_id_to_query_params(placeholder_initial_query_params)
initial_og_tags = get_open_graph_meta_tags(processed_placeholder_params)

gr.set_static_paths(paths=["images/", "models/", "assets/"])
with gr.Blocks(css_paths="style_20250503.css", title="3D viewer", theme='Surn/beeuty',delete_cache=(21600,86400), fill_width=True, head=initial_og_tags) as viewer3d:
    gr.Markdown("# 3D Model Viewer")
    
    with gr.Row():
        with gr.Column():
            model_3d = gr.Model3D(
                label="3D Model",
                value=None,
                elem_id="model_3d", key="model_3d", clear_color=[1.0, 1.0, 1.0, 0.1],
                elem_classes="centered solid imgcontainer", interactive=True
                
            )
            image_slider = gr.ImageSlider(
                label="2D Images",
                value=None,
                height="100%",
                elem_id="image_slider", key="image_slider",
                type="filepath"
            )
        
            with gr.Row():
                gr.Markdown("## Upload your own files")
                gr.Markdown("### Supported formats: " + ", ".join([f"`{ext}`" for ext in constants.upload_file_types]))
            with gr.Row():
                upload_btn = gr.UploadButton(
                    "Upload 3D Files", elem_id="upload_btn", key="upload_btn",
                    file_count="multiple",
                    file_types=constants.upload_file_types
                )
        
            with gr.Row():
                folder_name_box = gr.Textbox(
                    label="Folder Name",
                    value=default_folder,
                    elem_id="folder_name",
                    key="folder_name",
                    placeholder="Enter folder name...",
                    elem_classes="solid centered"
                )
                permalink_button = gr.Button("Generate Permalink", elem_id="permalink_button", key="permalink_button", elem_classes="solid small centered")

            with gr.Row(visible=False, elem_id="permalink_row") as permalink_row:
                permalink = gr.Textbox(
                    show_copy_button=True,
                    label="Permalink",
                    elem_id="permalink",
                    key="permalink",
                    elem_classes="solid centered",
                    max_lines=5,
                    lines=4
                )
                gr.Markdown("### Copy the permalink to share your model and images.", elem_classes="solid centered",)
                permalink_short = gr.Textbox(
                    show_copy_button=True,
                    label="Shortened Permalink",
                    elem_id="short_permalink",
                    key="permalink",
                    elem_classes="solid centered",
                    max_lines=5,
                    lines=3
                )
    with gr.Row():
        gr.HTML(value=getVersions(), visible=True, elem_id="versions")
    
    viewer3d.load(
        load_data,
        inputs=[model_3d, image_slider],
        outputs=[model_3d, image_slider, permalink, permalink_short],
        scroll_to_output=True
    ).then(
        lambda link: (gr.update(visible=True), gr.update(interactive=False))
            if link and len(link) > 0
            else (gr.update(visible=False), gr.update(interactive=True)),
        inputs=[permalink],
        outputs=[permalink_row, permalink_button]
    )

    upload_btn.upload(
        process_upload, 
        inputs=[upload_btn, model_3d, image_slider],
        outputs=[model_3d, image_slider],
        scroll_to_output=True,
        api_name="process_upload",
        show_progress=True
    ).then(
        lambda m, i: gr.update(interactive=True),
        inputs=[model_3d, image_slider],
        outputs=[permalink_button]
    )
    permalink_button.click(
        lambda model, images, folder: storage.upload_files_to_repo(
            files=[model] + list(images if images else []),
            repo_id=constants.HF_REPO_ID,
            folder_name=folder,
            create_permalink=True,
            repo_type="dataset"
        )[1],
        inputs=[model_3d, image_slider, folder_name_box],
        outputs=[permalink],
        scroll_to_output=True
    ).then(
        lambda link: gr.update(visible=True) if link and len(link) > 0 else gr.update(visible=False),
        inputs=[permalink],
        outputs=[permalink_row]
    )

if __name__ == "__main__":
    viewer3d.launch(
        allowed_paths=["assets", "assets/", "./assets", "images/", "./images", 'e:/TMP', 'models/', '3d_model_viewer/'],
        favicon_path="./assets/favicon.ico", show_api=True, strict_cors=False
    )