#!/usr/bin/env python3
"""
Gemini API proxy through SSH tunnel.

Provides SSH tunnel management and API wrappers for NanoBanana and Veo3
to bypass geo-restrictions by routing through a US-based server.
"""

import subprocess
import time
import base64
import os
import signal
import requests
from typing import Optional

BASE_URL = "https://generativelanguage.googleapis.com/v1beta"

# Global tunnel process reference
_tunnel_process: Optional[subprocess.Popen] = None


def start_tunnel(host: str, port: int, user: str = "root", local_port: int = 1080) -> bool:
    """
    Start SSH dynamic port forwarding (SOCKS5 proxy).

    Args:
        host: Remote SSH host (e.g., RunPod IP)
        port: Remote SSH port
        user: SSH username
        local_port: Local SOCKS5 port to bind

    Returns:
        True if tunnel started successfully
    """
    global _tunnel_process

    # Check if port is already in use
    try:
        import socket
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        result = sock.connect_ex(('127.0.0.1', local_port))
        sock.close()
        if result == 0:
            print(f"Port {local_port} already in use, assuming tunnel exists")
            return True
    except Exception:
        pass

    cmd = [
        "ssh",
        "-D", str(local_port),
        "-N",  # No remote command
        "-o", "StrictHostKeyChecking=no",
        "-o", "UserKnownHostsFile=/dev/null",
        "-o", "ServerAliveInterval=30",
        "-o", "ServerAliveCountMax=3",
        "-p", str(port),
        f"{user}@{host}"
    ]

    try:
        _tunnel_process = subprocess.Popen(
            cmd,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE
        )
        # Wait a bit for tunnel to establish
        time.sleep(2)

        if _tunnel_process.poll() is not None:
            stderr = _tunnel_process.stderr.read().decode() if _tunnel_process.stderr else ""
            raise RuntimeError(f"SSH tunnel failed to start: {stderr}")

        print(f"SSH tunnel started on localhost:{local_port}")
        return True

    except Exception as e:
        print(f"Failed to start SSH tunnel: {e}")
        return False


def stop_tunnel():
    """Stop the SSH tunnel process."""
    global _tunnel_process

    if _tunnel_process is not None:
        _tunnel_process.terminate()
        try:
            _tunnel_process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            _tunnel_process.kill()
        _tunnel_process = None
        print("SSH tunnel stopped")


def get_proxied_session(proxy_port: int = 1080) -> requests.Session:
    """Create a requests session configured to use SOCKS5 proxy."""
    session = requests.Session()
    session.proxies = {
        'http': f'socks5h://127.0.0.1:{proxy_port}',
        'https': f'socks5h://127.0.0.1:{proxy_port}'
    }
    return session


def call_nanobanana(
    image_data: bytes,
    prompt: str,
    api_key: str,
    mime_type: str = "image/webp",
    proxy_port: int = 1080,
    aspect_ratio: str = "16:9"
) -> bytes:
    """
    Call NanoBanana (Gemini image generation) API to process an image.

    Args:
        image_data: Raw image bytes
        prompt: Processing instructions (cleanup, resize, etc.)
        api_key: Gemini API key
        mime_type: Image MIME type
        proxy_port: Local SOCKS5 proxy port
        aspect_ratio: Output aspect ratio ("16:9", "9:16", "1:1", "4:3", "3:4")

    Returns:
        Processed image as bytes
    """
    session = get_proxied_session(proxy_port)

    image_base64 = base64.b64encode(image_data).decode('utf-8')

    payload = {
        "contents": [{
            "parts": [
                {"inlineData": {"mimeType": mime_type, "data": image_base64}},
                {"text": prompt}
            ]
        }],
        "generationConfig": {
            "responseModalities": ["IMAGE"],
            "imageConfig": {
                "aspectRatio": aspect_ratio
            }
        }
    }

    response = session.post(
        f"{BASE_URL}/models/gemini-3-pro-image-preview:generateContent",
        headers={
            "x-goog-api-key": api_key,
            "Content-Type": "application/json"
        },
        json=payload,
        timeout=180
    )

    if response.status_code != 200:
        raise RuntimeError(f"NanoBanana API error: {response.status_code} - {response.text[:500]}")

    data = response.json()

    # Extract image from response
    try:
        parts = data["candidates"][0]["content"]["parts"]
        for part in parts:
            if "inlineData" in part:
                image_data = base64.b64decode(part["inlineData"]["data"])
                mime_type_returned = part["inlineData"].get("mimeType", "image/png")
                # Return tuple of (data, mime_type) for proper extension handling
                return image_data, mime_type_returned
        raise KeyError("No image in response")
    except (KeyError, IndexError) as e:
        raise RuntimeError(f"Failed to extract image from response: {e}\n{data}")


def call_veo3(
    prompt: str,
    api_key: str,
    image_data: Optional[bytes] = None,
    mime_type: str = "image/png",
    last_frame_data: Optional[bytes] = None,
    last_frame_mime_type: str = "image/png",
    aspect_ratio: str = "16:9",
    proxy_port: int = 1080
) -> str:
    """
    Start Veo3 video generation (text-to-video, image-to-video, or frame interpolation).

    Args:
        prompt: Video generation prompt
        api_key: Gemini API key
        image_data: Optional image bytes for first frame (image-to-video mode)
        mime_type: Image MIME type for first frame
        last_frame_data: Optional image bytes for last frame (frame interpolation)
        last_frame_mime_type: Image MIME type for last frame
        aspect_ratio: Video aspect ratio ("16:9" or "9:16")
        proxy_port: Local SOCKS5 proxy port

    Returns:
        Operation name for polling
    """
    session = get_proxied_session(proxy_port)

    # Build instance with prompt
    instance = {"prompt": prompt}

    # Add first frame image for image-to-video mode
    if image_data is not None:
        image_base64 = base64.b64encode(image_data).decode('utf-8')
        instance["image"] = {
            "bytesBase64Encoded": image_base64,
            "mimeType": mime_type
        }

    # Add last frame for frame interpolation mode
    if last_frame_data is not None:
        last_frame_base64 = base64.b64encode(last_frame_data).decode('utf-8')
        instance["lastFrame"] = {
            "bytesBase64Encoded": last_frame_base64,
            "mimeType": last_frame_mime_type
        }

    payload = {
        "instances": [instance],
        "parameters": {"aspectRatio": aspect_ratio}
    }

    response = session.post(
        f"{BASE_URL}/models/veo-3.1-fast-generate-preview:predictLongRunning",
        headers={
            "x-goog-api-key": api_key,
            "Content-Type": "application/json"
        },
        json=payload,
        timeout=60
    )

    if response.status_code != 200:
        raise RuntimeError(f"Veo3 API error: {response.status_code} - {response.text[:500]}")

    data = response.json()
    return data["name"]


def poll_veo3(
    operation_name: str,
    api_key: str,
    proxy_port: int = 1080,
    poll_interval: int = 10,
    timeout: int = 600
) -> dict:
    """
    Poll Veo3 operation until complete.

    Args:
        operation_name: Operation name from call_veo3()
        api_key: Gemini API key
        proxy_port: Local SOCKS5 proxy port
        poll_interval: Seconds between polls
        timeout: Maximum wait time in seconds

    Returns:
        Dict with video URI and metadata
    """
    session = get_proxied_session(proxy_port)
    start_time = time.time()

    while time.time() - start_time < timeout:
        response = session.get(
            f"{BASE_URL}/{operation_name}",
            headers={"x-goog-api-key": api_key},
            timeout=30
        )

        if response.status_code != 200:
            raise RuntimeError(f"Poll error: {response.status_code} - {response.text[:500]}")

        data = response.json()

        if data.get("done"):
            # Extract video info
            try:
                video_info = data["response"]["generateVideoResponse"]["generatedSamples"][0]["video"]
                return {
                    "uri": video_info["uri"],
                    "state": video_info.get("state", "COMPLETE")
                }
            except (KeyError, IndexError) as e:
                raise RuntimeError(f"Failed to extract video from response: {e}\n{data}")

        elapsed = int(time.time() - start_time)
        print(f"  Veo3 generating... ({elapsed}s elapsed)")
        time.sleep(poll_interval)

    raise TimeoutError(f"Veo3 generation timed out after {timeout}s")


def download_video(
    video_uri: str,
    api_key: str,
    output_path: str,
    proxy_port: int = 1080
) -> str:
    """
    Download generated video.

    Args:
        video_uri: Video URI from poll_veo3()
        api_key: Gemini API key
        output_path: Local path to save video
        proxy_port: Local SOCKS5 proxy port

    Returns:
        Path to downloaded video
    """
    session = get_proxied_session(proxy_port)

    response = session.get(
        video_uri,
        headers={"x-goog-api-key": api_key},
        stream=True,
        timeout=300
    )

    if response.status_code != 200:
        raise RuntimeError(f"Download error: {response.status_code}")

    with open(output_path, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)

    return output_path


def test_connection(api_key: str, proxy_port: int = 1080) -> dict:
    """
    Test API connection through proxy.

    Returns:
        Dict with test results for each API
    """
    results = {}
    session = get_proxied_session(proxy_port)

    # Test NanoBanana
    try:
        response = session.post(
            f"{BASE_URL}/models/gemini-3-pro-image-preview:generateContent",
            headers={"x-goog-api-key": api_key, "Content-Type": "application/json"},
            json={
                "contents": [{"parts": [{"text": "A red dot"}]}],
                "generationConfig": {"responseModalities": ["TEXT", "IMAGE"]}
            },
            timeout=30
        )
        if response.status_code == 200:
            results["nanobanana"] = {"status": "ok"}
        else:
            results["nanobanana"] = {"status": "error", "code": response.status_code}
    except Exception as e:
        results["nanobanana"] = {"status": "error", "message": str(e)}

    # Test Veo3
    try:
        response = session.post(
            f"{BASE_URL}/models/veo-3.1-fast-generate-preview:predictLongRunning",
            headers={"x-goog-api-key": api_key, "Content-Type": "application/json"},
            json={
                "instances": [{"prompt": "A red dot"}],
                "parameters": {"aspectRatio": "16:9"}
            },
            timeout=30
        )
        if response.status_code == 200:
            results["veo3"] = {"status": "ok", "operation": response.json().get("name")}
        else:
            results["veo3"] = {"status": "error", "code": response.status_code}
    except Exception as e:
        results["veo3"] = {"status": "error", "message": str(e)}

    return results


if __name__ == "__main__":
    import sys

    if len(sys.argv) < 5:
        print("Usage: python gemini_proxy.py <ssh_host> <ssh_port> <api_key> <test|stop>")
        print("Example: python gemini_proxy.py PROXY_HOST 14492 AIzaSy... test")
        sys.exit(1)

    ssh_host = sys.argv[1]
    ssh_port = int(sys.argv[2])
    api_key = sys.argv[3]
    command = sys.argv[4]

    if command == "test":
        print(f"Starting SSH tunnel to {ssh_host}:{ssh_port}...")
        if start_tunnel(ssh_host, ssh_port):
            print("Testing API connection...")
            results = test_connection(api_key)
            print(f"NanoBanana: {results['nanobanana']}")
            print(f"Veo3: {results['veo3']}")
            stop_tunnel()
    elif command == "stop":
        stop_tunnel()
