import os
import subprocess
import time
import threading
import sys
import requests

_sd_process = None

CONFIG_FILE = os.path.expanduser("~/.epub-to-manga")

BACKENDS = [
    {
        "name": "AUTOMATIC1111 Stable Diffusion WebUI",
        "repo": "https://github.com/AUTOMATIC1111/stable-diffusion-webui",
        "dir": "stable-diffusion-webui",
        "type": "a1111",
        "launch": ["bash", "webui.sh", "--api", "--nowebui"],
        "ready_url": "http://127.0.0.1:7860/sdapi/v1/options",
    },
    {
        "name": "ComfyUI  (lighter, faster on Apple Silicon)",
        "repo": "https://github.com/comfyanonymous/ComfyUI",
        "dir": "ComfyUI",
        "type": "comfyui",
        "launch": [sys.executable, "main.py", "--listen"],
        "ready_url": "http://127.0.0.1:8188/system_stats",
    },
    {
        "name": "InvokeAI",
        "repo": "https://github.com/invoke-ai/InvokeAI",
        "dir": "InvokeAI",
        "type": "invokeai",
        "launch": ["invokeai-web"],
        "ready_url": "http://127.0.0.1:9090/api/v1/app/version",
    },
]

KNOWN_ERRORS = [
    {
        "markers": ["No module named 'pkg_resources'", "Couldn't install clip"],
        "message": "AUTOMATIC1111 is not compatible with Python {python_version}. A1111 requires Python 3.10.",
        "solutions": [
            {
                "label": "Install Python 3.10 via pyenv and relaunch A1111",
                "action": "fix_a1111_python",
            },
            {
                "label": "Switch to ComfyUI (better Apple Silicon support)",
                "action": "switch_backend",
                "backend_index": 1,
            },
        ],
    },
    {
        "markers": ["CUDA out of memory", "out of memory"],
        "message": "GPU ran out of memory.",
        "solutions": [
            {
                "label": "Re-run with a smaller image size (--size 256x384)",
                "action": "suggest_flag",
                "flag": "--size 256x384",
            },
        ],
    },
]

def read_config():
    config = {}
    if os.path.exists(CONFIG_FILE):
        with open(CONFIG_FILE) as f:
            for line in f:
                line = line.strip()
                if "=" in line and not line.startswith("#"):
                    key, _, val = line.partition("=")
                    config[key.strip()] = val.strip()
    return config

def write_config(key, value):
    config = read_config()
    config[key] = value
    with open(CONFIG_FILE, "w") as f:
        for k, v in config.items():
            f.write(f"{k}={v}\n")

def get_backend_for_path(sd_path):
    if not sd_path:
        return BACKENDS[0]
    for b in BACKENDS:
        if b["dir"] in sd_path or os.path.exists(os.path.join(sd_path, b["dir"])):
            return b
        if b["type"] == "a1111" and os.path.exists(os.path.join(sd_path, "webui.sh")):
            return b
        if b["type"] == "comfyui" and os.path.exists(os.path.join(sd_path, "comfy_extras")):
            return b
    return BACKENDS[0]

def is_running(ready_url, timeout=3):
    try:
        requests.get(ready_url, timeout=timeout)
        return True
    except Exception:
        return False

def find_installed():
    home = os.path.expanduser("~")
    search_dirs = [home, os.path.join(home, "Downloads"), os.path.join(home, "Documents"), os.getcwd()]
    for backend in BACKENDS:
        for base in search_dirs:
            candidate = os.path.join(base, backend["dir"])
            if os.path.isdir(candidate):
                return candidate, backend
    return None, None

def get_python_version():
    return f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"

def handle_known_error(error_def, sd_path, backend):
    python_version = get_python_version()
    msg = error_def["message"].format(python_version=python_version)
    solutions = error_def["solutions"]

    print(f"\n\n  ⚠️  {msg}\n")
    print("  Solutions:\n")
    for i, s in enumerate(solutions, 1):
        print(f"    {i}) {s['label']}")
    print()

    while True:
        choice = input("  Choose a solution (or q to quit): ").strip().lower()
        if choice == "q":
            sys.exit(0)
        if choice.isdigit() and 1 <= int(choice) <= len(solutions):
            solution = solutions[int(choice) - 1]
            break
        print(f"  Enter a number between 1 and {len(solutions)}")

    action = solution["action"]

    if action == "fix_a1111_python":
        print("\n  Installing pyenv and Python 3.10.14...")
        subprocess.run(["brew", "install", "pyenv"], check=False)
        subprocess.run(["pyenv", "install", "3.10.14"], check=False)
        pyenv_python = os.path.expanduser("~/.pyenv/versions/3.10.14/bin/python3")
        if not os.path.exists(pyenv_python):
            print("\n  Error: pyenv install failed. See https://github.com/pyenv/pyenv")
            sys.exit(1)
        env = os.environ.copy()
        env["PYTHON"] = pyenv_python
        print(f"\n  Relaunching A1111 with Python 3.10.14...")
        global _sd_process
        _sd_process = subprocess.Popen(
            ["bash", "webui.sh", "--api", "--nowebui"],
            cwd=sd_path,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            env=env,
        )
        stream_and_wait(backend, sd_path)

    elif action == "switch_backend":
        new_backend = BACKENDS[solution["backend_index"]]
        dest = os.path.join(os.path.expanduser("~"), new_backend["dir"])
        if os.path.isdir(dest):
            print(f"\n  Found existing {new_backend['name']} at {dest} — using it.")
        else:
            confirm = input(f"\n  Install {new_backend['name']} to {dest}? [Y/n] ").strip().lower()
            if confirm in ("n", "no"):
                print("  Aborted.")
                sys.exit(0)
            print(f"\n  Cloning {new_backend['repo']}...")
            result = subprocess.run(["git", "clone", "--recursive", new_backend["repo"], dest], check=False)
            if result.returncode != 0:
                print("\n  Error: git clone failed.")
                sys.exit(1)
            print(f"\n  Installed to {dest}")
        write_config("sd-path", dest)
        write_config("sd-backend", new_backend["type"])
        print(f"  Re-run your original command — switching will take effect automatically.")
        sys.exit(0)

    elif action == "suggest_flag":
        print(f"\n  Re-run with: {solution['flag']}")
        sys.exit(0)

def stream_logs(proc, ready_event, error_detected, sd_path, backend):
    interesting = [
        "Loading", "Downloading", "Installing", "Running", "Creating",
        "Model loaded", "Starting", "Applying", "torch", "CUDA", "MPS",
        "checkpoint", "venv", "pip", "Startup", "listen",
    ]

    for raw in proc.stdout:
        if ready_event.is_set():
            break
        line = raw.decode("utf-8", errors="replace").rstrip()

        for err_def in KNOWN_ERRORS:
            if any(m in line for m in err_def["markers"]):
                ready_event.set()
                error_detected["def"] = err_def
                return

        if any(kw in line for kw in interesting):
            print(f"\r  > {line[:100]:<100}")
            print("  ", end="", flush=True)

def stream_and_wait(backend, sd_path):
    global _sd_process
    ready_event = threading.Event()
    error_detected = {}

    log_thread = threading.Thread(
        target=stream_logs,
        args=(_sd_process, ready_event, error_detected, sd_path, backend),
        daemon=True,
    )
    log_thread.start()

    start = time.time()
    while not ready_event.is_set():
        elapsed = int(time.time() - start)
        mins, secs = divmod(elapsed, 60)
        elapsed_str = f"{mins}m {secs:02d}s" if mins else f"{secs}s"
        print(f"\r  Waiting for API... {elapsed_str} elapsed  ", end="", flush=True)

        if is_running(backend["ready_url"]):
            ready_event.set()
            elapsed = int(time.time() - start)
            mins, secs = divmod(elapsed, 60)
            elapsed_str = f"{mins}m {secs:02d}s" if mins else f"{secs}s"
            print(f"\r  Ready in {elapsed_str}!                              \n")
            return

        if _sd_process.poll() is not None and not ready_event.is_set():
            ready_event.set()
            break

        time.sleep(2)

    log_thread.join(timeout=2)

    if "def" in error_detected:
        handle_known_error(error_detected["def"], sd_path, backend)
    elif not is_running(backend["ready_url"]):
        print("\n\n  Error: SD process exited unexpectedly.")
        sys.exit(1)

def launch(sd_path, backend):
    global _sd_process
    print(f"\nStarting {backend['name']} from {sd_path}...")

    _sd_process = subprocess.Popen(
        backend["launch"],
        cwd=sd_path,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    stream_and_wait(backend, sd_path)

def prompt_install():
    print("\nNo Stable Diffusion backend found on your system.")
    print("\nAvailable backends to install:\n")
    for i, b in enumerate(BACKENDS, 1):
        print(f"  {i}) {b['name']}")
        print(f"     {b['repo']}")
    print()

    while True:
        choice = input("Enter number to install (or q to quit): ").strip().lower()
        if choice == "q":
            sys.exit(0)
        if choice.isdigit() and 1 <= int(choice) <= len(BACKENDS):
            backend = BACKENDS[int(choice) - 1]
            break
        print(f"  Please enter a number between 1 and {len(BACKENDS)}")

    dest = os.path.join(os.path.expanduser("~"), backend["dir"])
    confirm = input(f"\nInstall {backend['name']} to {dest}? [Y/n] ").strip().lower()
    if confirm in ("n", "no"):
        print("Aborted.")
        sys.exit(0)

    if os.path.isdir(dest):
        print(f"\nFound existing {backend['name']} at {dest} — using it.")
    else:
        print(f"\nCloning {backend['repo']}...")
        result = subprocess.run(["git", "clone", "--recursive", backend["repo"], dest], check=False)
        if result.returncode != 0:
            print("\nError: git clone failed. Is git installed?")
            sys.exit(1)
        print(f"\nInstalled to {dest}")

    write_config("sd-path", dest)
    write_config("sd-backend", backend["type"])
    print(f"Saved to {CONFIG_FILE}")
    print(f"\nRe-run your original command — no --sd-path needed.")
    sys.exit(0)

def ensure_running(sd_url=None, sd_path=None, style="manga"):
    config = read_config()

    if not sd_path:
        sd_path = config.get("sd-path")

    backend = get_backend_for_path(sd_path)

    if is_running(backend["ready_url"]):
        return

    if sd_path and os.path.isdir(sd_path):
        from image import backend as backend_mod
        backend_mod.setup(sd_path)
        backend_mod.ensure_model(backend["type"], sd_path, style=style)
        launch(sd_path, backend)
        return

    installed_path, found_backend = find_installed()
    if installed_path:
        print(f"\nFound {found_backend['name']} at {installed_path}")
        write_config("sd-path", installed_path)
        write_config("sd-backend", found_backend["type"])
        print(f"Saved to {CONFIG_FILE}")
        from image import backend as backend_mod
        backend_mod.setup(installed_path)
        backend_mod.ensure_model(found_backend["type"], installed_path, style=style)
        launch(installed_path, found_backend)
        return

    prompt_install()

def shutdown():
    global _sd_process
    if _sd_process:
        _sd_process.terminate()
        _sd_process = None