import sys
import os
import subprocess # For calling generate.py
import tempfile # For handling temporary image files
from typing import Optional
from PIL import Image as PILImage
import gradio as gr
import time # For timing

# Add the cloned nanoVLM directory to Python's system path
NANOVLM_REPO_PATH = "/app/nanoVLM"
if NANOVLM_REPO_PATH not in sys.path:
    print(f"DEBUG: Adding {NANOVLM_REPO_PATH} to sys.path")
    sys.path.insert(0, NANOVLM_REPO_PATH)

print(f"DEBUG: Python sys.path: {sys.path}")
print(f"DEBUG: Gradio version: {gr.__version__}") # Log Gradio version

GENERATE_SCRIPT_PATH = "/app/nanoVLM/generate.py"
MODEL_REPO_ID = "lusxvr/nanoVLM-222M"

print(f"DEBUG: Using generate.py script at: {GENERATE_SCRIPT_PATH}")
print(f"DEBUG: Using model repo ID: {MODEL_REPO_ID}")

def call_generate_script(image_path: str, prompt_text: str) -> str:
    print(f"\n--- DEBUG (call_generate_script) ---")
    print(f"Timestamp: {time.strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"Calling with image_path='{image_path}', prompt='{prompt_text}'")
    
    # Arguments for nanoVLM's generate.py
    # Using low max_new_tokens for CPU testing.
    cmd_args = [
        "python", "-u", GENERATE_SCRIPT_PATH, # -u for unbuffered output
        "--hf_model", MODEL_REPO_ID,
        "--image_path", image_path,      # Corrected: nanoVLM generate.py uses --image_path
        "--prompt", prompt_text,
        "--num_samples", "1",           # Corrected: Corresponds to --generations
        "--max_new_tokens", "30",       # Keep it low for testing
        "--device", "cpu"               # Explicitly set device for generate.py
        # Optional args for generate.py:
        # "--temperature", "0.7",
        # "--top_k", "50"
    ]
    
    print(f"Executing command: {' '.join(cmd_args)}")
    
    # Realistic timeout for the subprocess. HF Spaces free tier usually times out requests around 60s.
    # Set this shorter to catch issues within app.py.
    SCRIPT_TIMEOUT_SECONDS = 55 
    start_time = time.time()
    
    process_details = "Process details not available." # Placeholder
    try:
        process = subprocess.run(
            cmd_args,
            capture_output=True,
            text=True,
            check=False,  # Set to False to manually check returncode and log output
            timeout=SCRIPT_TIMEOUT_SECONDS
        )
        process_details = f"PID {process.pid if hasattr(process, 'pid') else 'N/A'}"

        duration = time.time() - start_time
        print(f"Subprocess ({process_details}) finished in {duration:.2f} seconds.")
        print(f"generate.py RETURN CODE: {process.returncode}")
        
        stdout = process.stdout.strip() if process.stdout else "[No STDOUT from generate.py]"
        stderr = process.stderr.strip() if process.stderr else "[No STDERR from generate.py]"

        print(f"---------- generate.py STDOUT ({process_details}) START ----------\n{stdout}\n---------- generate.py STDOUT ({process_details}) END ----------")
        if stderr or process.returncode != 0:
            print(f"---------- generate.py STDERR ({process_details}) START ----------\n{stderr}\n---------- generate.py STDERR ({process_details}) END ----------")

        if process.returncode != 0:
            error_message = f"Error: Generation script failed (code {process.returncode})."
            if "out of memory" in stderr.lower(): error_message += " Potential OOM in script."
            print(error_message) # Log it before returning
            return error_message + f" See Space logs for full STDOUT/STDERR from script ({process_details})."

        # --- Parse the output from nanoVLM's generate.py ---
        # Expected format:
        # Outputs:
        # > Sample 1: <generated text>
        output_lines = stdout.splitlines()
        generated_text = "[No parsable output from generate.py]" # Default
        
        found_output_line = False
        for line_idx, line in enumerate(output_lines):
            stripped_line = line.strip()
            # print(f"Parsing STDOUT line {line_idx}: '{stripped_line}'") # Can be very verbose
            if stripped_line.startswith("> Sample 1:") or stripped_line.startswith(">> Generation 1:"):
                prefix_to_remove = ""
                if stripped_line.startswith("> Sample 1:"): prefix_to_remove = "> Sample 1:"
                elif stripped_line.startswith(">> Generation 1:  "): prefix_to_remove = ">> Generation 1:  " # Note double space
                elif stripped_line.startswith(">> Generation 1: "): prefix_to_remove = ">> Generation 1: " # Note single space

                if prefix_to_remove:
                    generated_text = stripped_line.replace(prefix_to_remove, "", 1).strip()
                    found_output_line = True
                    print(f"Parsed generated text: '{generated_text}'")
                    break 
        
        if not found_output_line:
            print(f"Could not find 'Sample 1' or 'Generation 1' line in generate.py output.")
            # Return a snippet of STDOUT if parsing fails, to help debug output format
            generated_text = f"[Parsing failed] STDOUT (first 200 chars): {stdout[:200]}"


        print(f"Returning parsed text: '{generated_text}'")
        return generated_text

    except subprocess.TimeoutExpired as e:
        duration = time.time() - start_time
        print(f"ERROR: generate.py ({process_details}) timed out after {duration:.2f} seconds (limit: {SCRIPT_TIMEOUT_SECONDS}s).")
        stdout_on_timeout = e.stdout.strip() if e.stdout else "[No STDOUT on timeout]"
        stderr_on_timeout = e.stderr.strip() if e.stderr else "[No STDERR on timeout]"
        print(f"STDOUT on timeout:\n{stdout_on_timeout}")
        print(f"STDERR on timeout:\n{stderr_on_timeout}")
        return f"Error: Generation script timed out after {SCRIPT_TIMEOUT_SECONDS}s. Model loading and generation may be too slow for CPU."
    except Exception as e:
        duration = time.time() - start_time
        print(f"ERROR: An unexpected error occurred ({process_details}) after {duration:.2f}s: {type(e).__name__} - {e}")
        import traceback; traceback.print_exc()
        return f"Unexpected error calling script: {str(e)}"
    finally:
        print(f"--- END (call_generate_script) ---")


def gradio_interface_fn(image_input_pil: Optional[PILImage.Image], prompt_input_str: Optional[str]) -> str:
    print(f"\nDEBUG (gradio_interface_fn): Timestamp: {time.strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"Received prompt: '{prompt_input_str}', Image type: {type(image_input_pil)}")

    if image_input_pil is None:
        return "Please upload an image."
    
    cleaned_prompt = prompt_input_str.strip() if prompt_input_str else ""
    if not cleaned_prompt:
        return "Please provide a non-empty prompt."

    tmp_image_path = None 
    try:
        if image_input_pil.mode != "RGB":
            print(f"Converting image from {image_input_pil.mode} to RGB.")
            image_input_pil = image_input_pil.convert("RGB")

        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_image_file:
            image_input_pil.save(tmp_image_file, format="JPEG")
            tmp_image_path = tmp_image_file.name
        
        print(f"Temporary image saved to: {tmp_image_path}")
        
        result_text = call_generate_script(tmp_image_path, cleaned_prompt)
        
        print(f"Result from call_generate_script: '{result_text}'")
        return result_text
        
    except Exception as e:
        print(f"ERROR (gradio_interface_fn): Error processing image or calling script: {type(e).__name__} - {e}")
        import traceback; traceback.print_exc()
        return f"An error occurred in Gradio interface function: {str(e)}"
    finally:
        if tmp_image_path and os.path.exists(tmp_image_path):
            try:
                os.remove(tmp_image_path)
                print(f"Temporary image {tmp_image_path} removed.")
            except Exception as e_remove:
                print(f"WARN: Could not remove temporary image {tmp_image_path}: {e_remove}")
        print(f"DEBUG (gradio_interface_fn): Exiting.")


# --- Gradio Interface Definition ---
description_md = """
## nanoVLM-222M Interactive Demo (via generate.py)
Upload an image and type a prompt. This interface calls the `generate.py` script from
`huggingface/nanoVLM` under the hood to perform inference.
**Note:** Each request re-loads the model via the script, so it might be slow on CPU.
"""

print("DEBUG: Defining Gradio interface...")
iface = None
try:
    iface = gr.Interface(
        fn=gradio_interface_fn,
        inputs=[
            gr.Image(type="pil", label="Upload Image"),
            gr.Textbox(label="Your Prompt / Question", info="e.g., 'describe this image in detail'")
        ],
        outputs=gr.Textbox(label="Generated Text", show_copy_button=True, lines=5),
        title="nanoVLM-222M Demo (via Script)",
        description=description_md,
        allow_flagging="never"
    )
    print("DEBUG: Gradio interface defined successfully.")
except Exception as e:
    print(f"CRITICAL ERROR defining Gradio interface: {e}")
    import traceback; traceback.print_exc()

# --- Launch Gradio App ---
if __name__ == "__main__":
    print("DEBUG: Entered __main__ block for Gradio launch.")
    if not os.path.exists(GENERATE_SCRIPT_PATH):
        print(f"CRITICAL ERROR: The script {GENERATE_SCRIPT_PATH} was not found. Cannot launch app.")
        iface = None 

    if iface is not None:
        print("DEBUG: Attempting to launch Gradio interface...")
        try:
            iface.launch(server_name="0.0.0.0", server_port=7860)
            print("DEBUG: Gradio launch command issued. UI should be accessible.")
        except Exception as e:
            print(f"CRITICAL ERROR launching Gradio interface: {e}")
            import traceback; traceback.print_exc()
    else:
        print("CRITICAL ERROR: Gradio interface (iface) is None or not defined. Cannot launch.")