import os
import uuid
import numpy as np
import itertools
import concurrent.futures
from tqdm import tqdm
import zipfile
import io
import time
from .core import TaskRequest, zip_dir
from .device import LocalDevice, CenterDevices


class BatchTask:
    def __init__(self, script, dependency, result_root):
        self.script_abs = os.path.abspath(script)
        self.dep_abs = os.path.abspath(dependency)
        self.result_root = result_root

        self.script_rel = os.path.relpath(self.script_abs, os.path.dirname(self.dep_abs))

        self.params = {}
        self.devices = []
        self.results = []

        self.project_zip = zip_dir(self.dep_abs)


    def add_parameter(self, key, value):
        self.params[key] = value

    def to(self, device_list):
        self.devices = []
        for dev in device_list:
            if isinstance(dev, CenterDevices):
                self.devices.extend(dev.expand())
            else:
                self.devices.append(dev)

    def _generate_tasks(self):
        keys = self.params.keys()
        values = self.params.values()
        combinations = list(itertools.product(*values))

        tasks = []
        for combo in combinations:
            param_dict = dict(zip(keys, combo))
            task_id = str(uuid.uuid4())[:8]

            cmd = ["python", "placeholder_script"]
            for k, v in param_dict.items():
                cmd.append(f"--{k}")
                cmd.append(str(v))

            task = TaskRequest(
                task_id=task_id,
                project_zip=self.project_zip,
                script_rel_path=self.script_rel,
                run_cmd=cmd,
                result_dir="result"
            )
            tasks.append(task)
        return tasks

    def run(self, wait=True):
        tasks = self._generate_tasks()
        total_tasks = len(tasks)
        print(f"Generated {total_tasks} tasks. Available devices: {len(self.devices)}")

        if not wait:
            print("Non-wait mode not fully implemented for demo (requires persistent server DB). Running sync mode.")

        with concurrent.futures.ThreadPoolExecutor(max_workers=len(self.devices)) as executor:
            import queue
            device_queue = queue.Queue()
            for d in self.devices:
                device_queue.put(d)

            def run_wrapper(task):
                device = device_queue.get()
                try:
                    start_time = time.time()
                    res = device.run_task(task)
                    res.duration = time.time() - start_time
                    return res
                finally:
                    device_queue.put(device)

            futures = [executor.submit(run_wrapper, t) for t in tasks]

            pbar = tqdm(total=total_tasks, desc="Training")
            completed_durations = []

            for future in concurrent.futures.as_completed(futures):
                res = future.result()
                self.results.append(res)

                completed_durations.append(getattr(res, 'duration', 0))
                avg_time = np.mean(completed_durations)
                remaining = total_tasks - len(self.results)
                est_time = (remaining / len(self.devices)) * avg_time
                pbar.set_postfix_str(f"Est: {est_time:.1f}s")
                pbar.update(1)

            pbar.close()

    def save_results(self):
        os.makedirs(self.result_root, exist_ok=True)
        for res in self.results:
            if not res.success:
                print(f"Task {res.task_id} failed: {res.stderr}")
                continue

            task_res_dir = os.path.join(self.result_root, f"task_{res.task_id}")
            os.makedirs(task_res_dir, exist_ok=True)

            if res.result_zip:
                with zipfile.ZipFile(io.BytesIO(res.result_zip)) as zf:
                    zf.extractall(task_res_dir)

            with open(os.path.join(task_res_dir, "stdout.log"), "w") as f:
                f.write(res.stdout)

        print(f"All results saved to {self.result_root}")