import os
import sys
from argparse import Namespace
import docker
from docker import DockerClient, client
from docker import errors
from docker.models.containers import Container

from testgen.service.service import Service

AST_STRAT = 1
FUZZ_STRAT = 2
RANDOM_STRAT = 3

UNITTEST_FORMAT = 1
PYTEST_FORMAT = 2
DOCTEST_FORMAT = 3

class DockerController:
    def __init__(self):
        self.service = Service()
        self.args = None

    def run_in_docker(self, project_root: str, docker_client: DockerClient, args: Namespace) -> bool:
        self.args = args
        os.environ["RUNNING_IN_DOCKER"] = "1"

        # Check if Docker image exists, build it if not
        image_name = "testgen-runner"
        # If args.safe is set to false it means the image was not found and the system will try to run_locally
        self.get_image(docker_client, image_name, project_root)
        if not self.args.safe:
            print("Docker image not found. Running locally...")
            return False

        docker_args = [os.path.basename(args.file_path)] + [arg for arg in sys.argv[2:] if arg != "--safe"]

        # Run the container with the same arguments
        try:
            container = self.run_container(docker_client, image_name, docker_args, project_root)

            # Stream the logs to the console
            logs_output = self.get_logs(container)
            print(logs_output)

            try:
                # Create the target directory if it doesn't exist
                if args.output is None:
                    target_path = os.path.join(os.getcwd(), 'tests')
                else:
                    target_path = args.output
                os.makedirs(target_path, exist_ok=True)

                print(f"SERVICE target path after logs: {target_path}")

                test_cases = self.service.parse_test_cases_from_logs(logs_output)

                print(f"Extracted {len(test_cases)} test cases from container.")

                file_path = os.path.abspath(args.file_path)
                print(f"Filepath in CLI CONTROLLER: {file_path}")
                self.service.set_file_path(file_path)

                if args.test_format == "pytest":
                    self.service.set_test_generator_format(PYTEST_FORMAT)
                elif args.test_format == "doctest":
                    self.service.set_test_generator_format(DOCTEST_FORMAT)
                else:
                    self.service.set_test_generator_format(UNITTEST_FORMAT)

                test_file = self.service.generate_test_file(test_cases, target_path)
                print(f"Unit tests saved to: {test_file}")

                if not args.generate_only:
                    print("Running coverage...")
                    import traceback
                    print(traceback.format_exc())
                    self.service.run_coverage(test_file)

            except Exception as e:
                print(f"Error running container: {e}")
                sys.exit(1)

        except Exception as e:
            print(f"Error running container: {e}")
            sys.exit(1)

    def get_image(self, docker_client: DockerClient, image_name: str, project_root: str):
        try:
            docker_client.images.get(image_name)
            print(f"Using existing {image_name} Docker image")
        except docker.errors.ImageNotFound:
            print(f"Building {image_name} Docker image...")

            # Look for Dockerfile in the project root
            docker_dir = os.path.join(project_root, "testgen", "docker")
            dockerfile_path = os.path.join(docker_dir, "Dockerfile")
            if not os.path.exists(dockerfile_path):
                print(f"Dockerfile not found at {dockerfile_path}")
                sys.exit(1)

            print(f"Using Dockerfile at: {dockerfile_path}")

            if not self.build_docker_image(docker_client, image_name, dockerfile_path, project_root):
                print("Failed to build Docker image. Continuing without safe mode.")
                self.args.safe = False

    @staticmethod
    def get_logs(container) -> str:
        # Stream the logs to the console
        print("Running in Docker container...")
        logs = container.logs(stream=True)
        logs_output = ""
        for log in logs:
            log_line = log.decode()
            logs_output += log_line
            print(log_line, end="")
        return logs_output

    @staticmethod
    def run_container(docker_client: DockerClient, image_name: str, docker_args: list, project_root: str) -> Container:
        return docker_client.containers.run(
                image=image_name,
                command=["poetry", "run", "python", "-m", "testgen.main"] + docker_args,
                volumes={project_root: {"bind": "/controller", "mode": "rw"}},  # Mount current dir
                environment={"RUNNING_IN_DOCKER": "1"},
                detach=True,
                remove=True,
                stdout=True,
                stderr=True
            )

    @staticmethod
    def build_docker_image(docker_client, image_name, dockerfile_path, project_root):
        try:
            print(f"Starting Docker build for image: {image_name}")
            dockerfile_rel_path = os.path.relpath(dockerfile_path, project_root)
            print(f"Project root {project_root}")
            print(f"Docker directory: {os.path.dirname(dockerfile_path)}")
            print(f"Docker rel path: {dockerfile_rel_path}")
            build_progress = docker_client.api.build(
                path=os.path.join(project_root, "testgen", "docker"),
                dockerfile=os.path.join(project_root, "testgen", "docker", "Dockerfile"),
                tag=image_name,
                rm=True,
                decode=True
            )

            for chunk in build_progress:
                print(f"CHUNK: {chunk}")
                if 'stream' in chunk:
                    for line in chunk['stream'].splitlines():
                        if line.strip():
                            print(f"Docker: {line.strip()}")
                elif 'error' in chunk:
                    print(f"Docker build error: {chunk['error']}")
                    return False
            print(f"Docker image built successfully: {image_name}")
            return True

        except docker.errors.BuildError as e:
            print(f"Docker build error: {e}")
        except docker.errors.APIError as e:
            print(f"Docker API error: {e}")
        except Exception as e:
            print(f"Unexpected error during Docker build: {str(e)}")
            return False


