import sys

try:
    from pprint import pprint
    import re
    import os
    from pathlib import Path
    import tempfile
    import subprocess
    from typing import Optional, Union, Any, Tuple
    import shutil
    from importlib import import_module
    import json
    from types import ModuleType

    from colorama import init, Style
    import numpy as np
    import cv2
    from skimage import transform
    from PIL import Image, UnidentifiedImageError
    from rich.console import Console
    from rich.progress import SpinnerColumn, Progress, TextColumn, BarColumn, TaskProgressColumn, TimeElapsedColumn
    from rich.text import Text
    from beartype import beartype
except ModuleNotFoundError as e:
    print(f"Failed ot load module: {e}")
    sys.exit(1)

init()

def dier (msg: Any) -> None:
    pprint(msg)
    sys.exit(1)

console = Console()

@beartype
def print_predictions_line(predictions: np.ndarray, labels: list) -> None:
    vals = predictions[0]
    max_index = int(np.argmax(vals))  # Index des höchsten Werts

    parts = []
    for i, (label, value) in enumerate(zip(labels, vals)):
        if i == max_index:
            part = f"{Style.BRIGHT}{label}: {value:.10f}{Style.RESET_ALL}"
        else:
            part = f"{label}: {value:.10f}"
        parts.append(part)

    line = "  ".join(parts)
    sys.stdout.write("\r" + line + " " * 5)
    sys.stdout.flush()

def _pip_install(package: str) -> bool:
    cmd = [sys.executable, "-m", "pip", "install", package]
    try:
        with Progress(
            SpinnerColumn(),
            TextColumn(f"[cyan]Installing {package}...[/cyan]"),
            transient=True,
            console=console,
        ) as progress:
            task = progress.add_task("pip_install", start=False)
            progress.start_task(task)
            result = subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            if result.returncode != 0:
                console.print(f"[red]Failed to install {package}.[/red]")
                console.print(f"[red]{result.stderr.strip()}[/red]")
            return result.returncode == 0
    except FileNotFoundError:
        console.print(f"[red]Python executable not found: {sys.executable}[/red]")
    except subprocess.CalledProcessError as e:
        console.print(f"[red]Installation failed for {package} (non-zero exit).[/red]")
        console.print(f"[red]{e.stderr.strip()}[/red]")
    except subprocess.SubprocessError as e:
        console.print(f"[red]A subprocess error occurred during installation of {package}.[/red]")
        console.print(f"[red]{str(e).strip()}[/red]")
    except KeyboardInterrupt:
        console.print(f"[yellow]Installation of {package} interrupted by user.[/yellow]")

    return False

@beartype
def rule(msg) -> None:
    console.rule(f"{msg}")

@beartype
def install_tensorflow() -> Optional[ModuleType]:
    try:
        rule("[bold cyan]Checking for TensorFlow...[/]")

        try:
            tf = import_module("tensorflow")
            return tf
        except ModuleNotFoundError:
            console.print("[yellow]TensorFlow not found. Installation required.[/yellow]")

        # 2. Check virtual environment
        if sys.prefix == sys.base_prefix:
            console.print("[red]Error: You need to be in a virtual environment to install TensorFlow.[/red]")

        # 3. Try install tensorflow or tf_nightly
        if not _pip_install("tensorflow") and not _pip_install("tf_nightly"):
            console.print("[red]Cannot install TensorFlow via pip.[/red]")

            _pip_install("keras")
            _pip_install("tf_keras")

        # 4. Tell user to restart the script
        console.print("[green]TensorFlow installed successfully! Run the script again, it should work now.[/green]")

        sys.exit(0)
    except KeyboardInterrupt:
        console.print("[green]You cancelled loading or installing TensorFlow by pressing CTRL-C. The script generated by asanAI will not work without tensorflow being installed. Try running it again or install it manually to the virtualenv.[/green]")
        sys.exit(0)


@beartype
def _newest_match(directory: Union[Path, str], pattern: str) -> Optional[Path]:
    directory = Path(directory)

    candidates = [
        p for p in directory.iterdir()
        if re.fullmatch(pattern, p.name)
    ]

    if not candidates:
        return None

    def extract_number(p: Path) -> int:
        match = re.search(r"\((\d+)\)", p.name)
        if not match:
            raise ValueError(f"No number found in parentheses in: {p.name}")
        return int(match.group(1))

    candidates.sort(
        key=extract_number,
        reverse=True,
    )
    return candidates[0]

@beartype
def find_model_files(directory: Optional[Union[Path, str]] = ".") -> dict[str, Optional[Path]]:
    if directory is None:
        console.log("[red]No directory provided[/red]")
        return {}

    directory = Path(directory)

    jobs: tuple[tuple[str, str], ...] = (
        ("model.json",        r"model\((\d+)\)\.json"),
        ("model.weights.bin", r"model\.weights\((\d+)\)\.bin"),
    )

    found_files: dict[str, Optional[Path]] = {}

    with Progress(
        TextColumn("[bold]{task.description}"),
        BarColumn(bar_width=None),
        TaskProgressColumn(),
        TimeElapsedColumn(),
        console=console,
        transient=True
    ) as progress:
        task_ids = {
            canonical: progress.add_task(f"Checking {canonical}", total=1)
            for canonical, _ in jobs
        }

        for canonical, regex in jobs:
            progress.update(task_ids[canonical], advance=0)  # render row
            target = directory / canonical
            if target.exists():
                progress.update(task_ids[canonical], completed=1)
                console.log(f"[green]{canonical} found[/green]")
                found_files[canonical] = target
                continue

            newest = _newest_match(directory, regex)
            if newest:
                console.log(f"[yellow]Using[/yellow] {newest.name} instead of {canonical}")
                found_files[canonical] = newest
            else:
                console.log(f"[red]Missing:[/red] No match for {canonical}")
                found_files[canonical] = None
            progress.update(task_ids[canonical], completed=1)

    return found_files

@beartype
def _is_command_available(cmd: str) -> bool:
    return shutil.which(cmd) is not None

@beartype
def _pip_install_tensorflowjs_converter_and_run_it(conversion_args: list) -> bool:
    if  not _is_command_available('tensorflowjs_converter'):
        _pip_install("tensorflowjs")

    if  _is_command_available('tensorflowjs_converter'):
        if _is_command_available('tensorflowjs_converter'):
            with console.status("[bold green]Local tensorflowjs_converter found. Starting conversion..."):
                cmd = ['tensorflowjs_converter'] + conversion_args
                try:
                    completed_process = subprocess.run(
                        cmd,
                        check=True,
                        stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE,
                        text=True
                    )
                    console.print("[green]✔ Local conversion succeeded.[/]")
                    console.print(Text(completed_process.stdout.strip(), style="dim"))
                    return True
                except subprocess.CalledProcessError as e:
                    console.print("[red]✘ Local conversion failed:[/]")
                    console.print(Text(e.stderr.strip(), style="bold red"))
                    console.print("[yellow]➜ Falling back to Docker-based conversion...[/]")
                except KeyboardInterrupt:
                    console.print("[green]You cancelled the conversion progress by CTRL-C. You need to run this script again or do it manually for this program to work.[/green]")
                    sys.exit(0)
        else:
            console.print("[yellow]⚠ tensorflowjs_converter CLI not found locally.[/]")
    else:
        console.print("[yellow]⚠ Installing tensorflowjs module failed. Trying to fall back to docker.[/]")

    return False

@beartype
def copy_and_patch_tfjs(model_json_path: str, weights_bin_path: str, out_prefix: str = "tmp_model") -> Tuple[str, str]:
    json_out = f"{out_prefix}.json"
    bin_out  = f"{out_prefix}.bin"

    # --- patch JSON ---
    with open(model_json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    # Point every manifest entry to the newly created .bin
    for manifest in data.get("weightsManifest", []):
        manifest["paths"] = [f"./{Path(bin_out).name}"]

    with open(json_out, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)

    # --- copy BIN ---
    shutil.copyfile(weights_bin_path, bin_out)

    return json_out, bin_out

@beartype
def convert_to_keras_if_needed(directory: Optional[Union[Path, str]] = ".") -> bool:
    rule("[bold cyan]Trying to see if if it is neccessary to convert TFJS models to Keras[/]")

    keras_h5_file = 'model.h5'

    if os.path.exists(keras_h5_file):
        console.print(f"[green]✔ Conversion not needed:[/] '{keras_h5_file}' already exists.")
        return True

    files = find_model_files(directory)

    original_tfjs_model_json = str(files.get("model.json"))
    original_weights_bin = str(files.get("model.weights.bin"))

    if not os.path.exists(original_tfjs_model_json) or not os.path.exists(original_weights_bin):
        console.print("[red]No model.json and/or model.weights.bin found. Cannot continue. Have you downloaded the models from asanAI? If not, do so and put them in the same folder as your script.[/red]")
        sys.exit(1)

    tfjs_model_json, weights_bin = copy_and_patch_tfjs(original_tfjs_model_json, original_weights_bin)

    if not tfjs_model_json or not weights_bin:
        console.log("[red]Missing model files. Conversion aborted.[/red]")
        return False

    if not os.path.exists(original_tfjs_model_json):
        console.print(f"[yellow]⚠ Conversion not possible:[/] '{original_tfjs_model_json}' not found.")
        return False

    console.print(f"[cyan]Conversion needed:[/] '{keras_h5_file}' does not exist, but '{original_tfjs_model_json}' found.")

    # Helper function to check if a command exists in PATH

    conversion_args = [
        '--input_format=tfjs_layers_model',
        '--output_format=keras',
        tfjs_model_json,
        keras_h5_file
    ]

    if _pip_install_tensorflowjs_converter_and_run_it(conversion_args):
        return True

    if not _is_command_available('docker'):
        console.print("[red]✘ Docker is not installed or not found in PATH. Cannot perform fallback conversion. Please install docker.[/]")
        return False

    try:
        subprocess.run(['docker', 'info'], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except subprocess.CalledProcessError:
        console.print("[red]✘ Docker daemon not running or inaccessible. Cannot perform fallback conversion.[/]")
        return False

    with tempfile.TemporaryDirectory() as tmpdir:
        dockerfile_path = os.path.join(tmpdir, 'Dockerfile')

        dockerfile_content = '''FROM python:3.10-slim

RUN apt-get update && \\
    apt-get install -y --no-install-recommends build-essential curl && \\
    rm -rf /var/lib/apt/lists/*

RUN python -m pip install --upgrade pip

RUN python -m pip install \\
    tensorflow==2.12.0 \\
    tensorflowjs==4.7.0 \\
    jax==0.4.13 \\
    jaxlib==0.4.13

WORKDIR /app

CMD ["/bin/bash"]
'''
        with open(dockerfile_path, mode='w', encoding="utf-8") as f:
            f.write(dockerfile_content)

        image_name = 'tfjs_converter_py310_dynamic'

        console.print("[cyan]Building Docker image for fallback conversion...[/]")
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            transient=True,
            console=console
        ) as progress:
            build_task = progress.add_task("Building Docker image...", total=None)
            try:
                build_cmd = ['docker', 'build', '-t', image_name, '-f', dockerfile_path, tmpdir]
                subprocess.run(build_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                progress.update(build_task, description="Docker image built successfully.")
            except subprocess.CalledProcessError as e:
                progress.stop()
                console.print("[red]✘ Docker build failed with error:[/]")
                console.print(Text(e.stderr.strip(), style="bold red"))
                return False
            except KeyboardInterrupt:
                progress.stop()
                console.print("[red]✘ Docker build was cancelled by CTRL-C[/]")
                sys.exit(0)

        run_cmd = [
            'docker', 'run', '--rm',
            '-v', f"{os.path.abspath(os.getcwd())}:/app",
            image_name,
            'tensorflowjs_converter',
        ] + conversion_args

        with console.status("[bold green]Running conversion inside Docker container..."):
            try:
                run_process = subprocess.run(run_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                console.print("[green]✔ Conversion inside Docker container succeeded.[/]")
                console.print(Text(run_process.stdout.strip(), style="dim"))
            except subprocess.CalledProcessError as e:
                console.print("[red]✘ Conversion inside Docker container failed with error:[/]")
                console.print(Text(e.stderr.strip(), style="bold red"))
                return False
            except KeyboardInterrupt:
                console.print("[red]✘ Docker build was cancelled by CTRL-C[/]")
                sys.exit(0)

    return True

@beartype
def load(filename: Union[Path, str], height: int = 224, width: int = 224, divide_by: Union[int, float] = 255.0) -> Optional[np.ndarray]:
    rule(f"[bold cyan]Loading image {filename}[/]")
    try:
        if not os.path.exists(filename):
            console.print(f"[red]Error: The path '{filename}' could not be found![/red]")
            return None

        try:
            with console.status(f"Loading image {filename}"):
                image = Image.open(filename)

            with console.status(f"Converting image {filename} to numpy array and normalizing"):
                np_image: np.ndarray = np.array(image).astype('float32') / divide_by

            with console.status(f"Resizing image {filename} to (height = {height}, width = {width}, channels = 3)"):
                np_image = transform.resize(np_image, (height, width, 3))

            with console.status(f"Expanding numpy array dimensions from image {filename}"):
                np_image = np.expand_dims(np_image, axis=0)

            return np_image

        except PermissionError:
            console.print(f"[red]Error: Permission denied for file '{filename}'. Please check file permissions.[/red]")

        except UnidentifiedImageError:
            console.print(f"[red]Error: The file '{filename}' is not a valid image or is corrupted.[/red]")

        except ValueError as ve:
            console.print(f"[red]Error: ValueError encountered: {ve}. Possibly wrong image dimensions or resize parameters.[/red]")

        except TypeError as te:
            console.print(f"[red]Error: TypeError encountered: {te}. Check if 'divide_by' is a number (int or float).[/red]")

        except OSError as ose:
            console.print(f"[red]Error: OS error occurred: {ose}. Possible file system issue.[/red]")
    except KeyboardInterrupt:
        console.print(f"[green]You cancelled loading the image {filename} by pressing CTRL-C[/green]")
        sys.exit(0)

    return None

@beartype
def load_frame(frame: np.ndarray, height: int = 224, width: int = 224, divide_by: Union[int, float] = 255.0) -> Optional[np.ndarray]:
    try:
        np_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # pylint: disable=no-member
        np_image = np.array(np_image).astype('float32') / divide_by
        np_image = transform.resize(np_image, (height, width, 3))
        np_image = np.expand_dims(np_image, axis=0)
        return np_image

    except cv2.error as e: # pylint: disable=no-member
        console.print(f"[red]OpenCV error during color conversion: {e}[/red]")

    except ValueError as ve:
        console.print(f"[red]ValueError during resize or processing: {ve}[/red]")

    except TypeError as te:
        console.print(f"[red]TypeError encountered: {te}. Check input types.[/red]")

    except OSError as ose:
        console.print(f"[red]OS error occurred: {ose}.[/red]")

    except KeyboardInterrupt:
        console.print("[green]You cancelled loading the fame by pressing CTRL-C[/green]")
        sys.exit(0)

    return None

@beartype
def annotate_frame(frame: np.ndarray, predictions: np.ndarray, labels: list[str]) -> Optional[np.ndarray]:
    probs = predictions[0]
    best_idx = int(np.argmax(probs))

    def format_probabilities(values: np.ndarray) -> list[str]:
        for precision in range(3, 12):  # vernünftiger Bereich
            formatted = [f"{v:.{precision}f}" for v in values]
            if len(set(formatted)) == len(values):
                return formatted
        return [f"{v:.10f}" for v in values]

    if len(labels) != len(probs):
        console.print(
            f"[bold red]❌ Label count ({len(labels)}) does not match number of prediction probabilities ({len(probs)}).[/bold red]",
        )
        console.print("[yellow]Make sure the number of labels in your script is correct.[/yellow]")

        sys.exit(0)

    formatted_probs = format_probabilities(probs)

    for i, label in enumerate(labels):
        text = f"{label}: {formatted_probs[i]}"
        colour = (0, 255, 0) if i == best_idx else (255, 0, 0)
        cv2.putText( # pylint: disable=no-member
            frame,
            text,
            (10, 30 * (i + 1)),
            cv2.FONT_HERSHEY_SIMPLEX, # pylint: disable=no-member
            0.8,
            colour,
            2,
        )
    return frame
