import json
import logging
import os
import re
import subprocess
import sys
import tempfile
import time
import asyncio
import pathlib
from pathlib import Path
from datetime import datetime

from .models import Workflow, SystemInfo, CodeResult
from . import co_datascientist_api
from .kpi_extractor import extract_kpi_from_stdout
from .settings import settings

OUTPUT_FOLDER = "co_datascientist_output"
CHECKPOINTS_FOLDER = "co_datascientist_checkpoints"


def print_workflow_info(message: str):
    """Print workflow info with consistent formatting"""
    print(f"   {message}")


def print_workflow_step(message: str):
    """Print workflow step with consistent formatting"""
    print(f"   🔄 {message}")


def print_workflow_success(message: str):
    """Print workflow success with consistent formatting"""
    print(f"   ✅ {message}")


def print_workflow_error(message: str):
    """Print workflow error with consistent formatting"""
    print(f"   ❌ {message}")


class _WorkflowRunner:
    def __init__(self):
        self.workflow: Workflow | None = None
        self.start_timestamp = 0
        self.should_stop_workflow = False
        self.debug_mode = True
        # Track best KPI seen when polling backend
        self._checkpoint_counter: int = 0

    async def run_workflow(self, code: str, python_path: str, project_absolute_path: str, config: dict, spinner=None, debug: bool = True):
        """Run a complete code evolution workflow.

        - Sequential mode (default): run one code version at a time.
        - Parallel mode: set config['parallel']=N (>1) to run batches of up to N in parallel.
        """
        self.should_stop_workflow = False
        
        # Set debug mode for the class instance
        self.debug_mode = debug
        
        try:
            if spinner:
                spinner.text = "Waking up the Co-DataScientist"
            self.start_timestamp = time.time()
            self.should_stop_workflow = False
            self.workflow = Workflow(status_text="Workflow started", user_id="")

            system_info = get_system_info(python_path)
            logging.info(f"user system info: {system_info}")
            
            response = await co_datascientist_api.start_workflow(code, system_info)
            self.workflow = response.workflow
            if spinner:
                spinner.stop()  # stop spinner without emoji
            print("Lets start exploring your problem space!")

            parallel_n = int(config.get('parallel', 1) or 1)
            if parallel_n > 1:
                await self._run_parallel_mode(response, python_path, project_absolute_path, config, spinner, parallel_n)
            else:
                await self._run_sequential_mode(response, python_path, project_absolute_path, config, spinner)

            if self.should_stop_workflow:
                await co_datascientist_api.stop_workflow(self.workflow.workflow_id)
                print_workflow_info("Workflow stopped by user.")
                if spinner:
                    spinner.text = "Workflow stopped"
            else:
                # Check if workflow finished due to baseline failure or successful completion
                if (hasattr(self.workflow, 'baseline_code') and 
                    self.workflow.baseline_code.result is not None and 
                    self.workflow.baseline_code.result.return_code != 0):
                    print_workflow_error("Workflow terminated due to baseline code failure!")
                    print("   📄 Review the error details above and fix your script.")
                    if spinner:
                        spinner.text = "Workflow failed"
                else:
                    print_workflow_success("Workflow completed successfully!")
                    if spinner:
                        spinner.text = "Workflow completed"
        
        except Exception as e:
            if spinner:
                spinner.stop()

            err_msg = str(e)
            # Detect user-facing validation errors coming from backend (prefixed with ❌)
            if err_msg.startswith("❌") and not self.debug_mode:
                # Show concise guidance without stack trace
                print_workflow_error(err_msg)
                return  # Do not re-raise, end gracefully

            # Otherwise, show generic workflow error and re-raise for full trace
            print_workflow_error(f"Workflow error: {err_msg}")
            raise

    async def _run_parallel_mode(self, initial_response, python_path: str,
                                 project_absolute_path: str, config: dict,
                                 spinner=None, parallel_n: int = 2):
        """Run workflow in parallel mode using batch endpoints.

        Requests batches of up to parallel_n programs from backend, executes them concurrently,
        sends results back in a single batch, and repeats until finished.
        """
        if spinner:
            spinner.text = f"Running up to {parallel_n} programs in parallel..."
            spinner.start()

        # Use the batch API regardless of initial single-item response
        batch_resp = await co_datascientist_api.get_batch_to_run(self.workflow.workflow_id, batch_size=parallel_n)

        while (batch_resp.batch_to_run is not None
               and not self.workflow.finished
               and not self.should_stop_workflow):

            code_versions = batch_resp.batch_to_run
            batch_id = batch_resp.batch_id

            # Execute all code versions concurrently (CPU-bound/blocking -> run in threads)
            def _execute(cv):
                if config.get('databricks'):
                    return _databricks_run_python_code(cv.code, python_path, config)
                else:
                    return _run_python_code(cv.code, python_path)

            tasks = [asyncio.to_thread(_execute, cv) for cv in code_versions]
            results: list[CodeResult] = await asyncio.gather(*tasks, return_exceptions=False)

            # Attach KPI values and prepare payload tuples
            tuples: list[tuple[str, CodeResult]] = []
            for cv, res in zip(code_versions, results):
                kpi_value = extract_kpi_from_stdout(res.stdout)
                res.kpi = kpi_value
                tuples.append((cv.code_version_id, res))

            # Submit batch results and get next batch
            batch_resp = await co_datascientist_api.finished_running_batch(
                self.workflow.workflow_id,
                batch_id,
                tuples,
            )
            # Update workflow state
            self.workflow = batch_resp.workflow

            # Optionally display live best KPI
            try:
                best_info = await co_datascientist_api.get_workflow_population_best(self.workflow.workflow_id)
                best_kpi = best_info.get("best_kpi") if best_info else None
                if best_kpi is not None and spinner:
                    spinner.write(f"🚀 Current best KPI: {best_kpi}")
            except Exception:
                pass

    async def _run_sequential_mode(self, response, python_path: str, 
                                  project_absolute_path: str, config: dict, spinner=None, poll_interval: int = 1):
        """Run workflow in sequential mode"""
        
        iteration_count = 0

        # Ensure spinner is active throughout the loop
        if spinner:
            spinner.text = "Thinking..."
            spinner.start()

        while not self.workflow.finished and response.code_to_run is not None and not self.should_stop_workflow:
            # No minimal/best-only modes – always standard
            
            # keep spinner running; write a blank line for spacing
            # if spinner:
            #     spinner.write("")
            # else:
            #     print()

            ###Remember to start the code is there... we need to grab it first! then evenything can be unchanged...
            # Section where we run the code with particulal cloud integration.
            if config.get('databricks'):
                result = _databricks_run_python_code(response.code_to_run.code, python_path, config)
            else:
                result = _run_python_code(response.code_to_run.code, python_path)

            # Log only for baseline; skip verbose output for other ideas
            if response.code_to_run.name == "baseline":
                await self._handle_baseline_result(result, response, spinner)
            # # Extra space before the next spinner line
            # if spinner:
            #     spinner.write("")
            # else:
            #     print()
            # Restart spinner while waiting for next idea
            if spinner:
                spinner.text = "Dont worry, I'm just thinking..."
                spinner.start()

            # Prepare objects for backend
            kpi_value = extract_kpi_from_stdout(result.stdout)
            result.kpi = kpi_value
            code_version = response.code_to_run
            code_version.result = result

            response = await co_datascientist_api.finished_running_code(
                self.workflow.workflow_id,
                code_version,
                result,
                kpi_value,
            )
            self.workflow = response.workflow

            # Poll backend for best KPI every poll_interval iterations
            iteration_count += 1
            if iteration_count % poll_interval == 0:
                try:
                    best_info = await co_datascientist_api.get_workflow_population_best(self.workflow.workflow_id)
                    best_kpi = best_info.get("best_kpi") if best_info else None
                    if best_kpi is not None and spinner:
                        spinner.write(f"🚀 Current best KPI: {best_kpi}")

                    # Always save checkpoint snapshot of current best
                    best_cv = best_info.get("best_code_version") if best_info else None
                    if best_cv and best_kpi is not None:
                        await self._save_population_best_checkpoint(best_cv, best_kpi, project_absolute_path, config)
                        if spinner:
                            spinner.write(f"💾 Checkpoint saved (KPI: {best_kpi})")
                    elif best_kpi is not None:
                        if spinner:
                            spinner.write(f"⚠️  No code version available for checkpoint (KPI: {best_kpi})")
                except Exception as e:
                    # Non-fatal: just log and continue
                    logging.warning(f"Failed fetching best KPI code: {e}")

    async def _handle_baseline_result(self, result: CodeResult, response, spinner=None):
        """Handle result in standard mode (original behavior)"""
        # Check if code execution failed and provide clear feedback
        if result.return_code != 0:
            # Code failed - show error details
            print_workflow_error(f"'{response.code_to_run.name}' failed with exit code {result.return_code}")
            if result.stderr:
                print("   📄 Error details:")
                # Print each line of stderr with proper indentation
                for line in result.stderr.strip().split('\n'):
                    if spinner:
                        spinner.write(f"      {line}")
                    else:
                        print(f"      {line}")
            
            # For baseline failures, give specific guidance
            if response.code_to_run.name == "baseline":
                print("   💡 The baseline code failed to run. This will stop the workflow.")
                print("   💡 Check the error above and fix your script before running again.")
                if "ModuleNotFoundError" in (result.stderr or ""):
                    print("   💡 Missing dependencies? Try: pip install <missing-package>")
        else:
            # Code succeeded - show success message
            kpi_value = extract_kpi_from_stdout(result.stdout)
            if kpi_value is not None:
                msg = f"Completed '{response.code_to_run.name}' | KPI = {kpi_value}"
                if spinner:
                    spinner.write("   ✅ " + msg)
                else:
                    print_workflow_success(msg)
            else:
                msg = f"Completed '{response.code_to_run.name}'"
                if spinner:
                    spinner.write("   ✅ " + msg)
                else:
                    print_workflow_success(msg)

    async def _save_population_best_checkpoint(self, best_cv, best_kpi: float, project_absolute_path: str, config: dict):
        """Persist best code/KPI - to Databricks volume if using Databricks, locally otherwise."""
        try:
            if not best_cv or best_kpi is None:
                return

            # Convert best_cv to CodeVersion model if it is raw dict
            from .models import CodeVersion, CodeResult
            if isinstance(best_cv, dict):
                try:
                    # Nested result may also be dict – handle gracefully
                    if isinstance(best_cv.get("result"), dict):
                        # Ensure runtime_ms field may be missing; allow extra
                        best_cv["result"] = CodeResult.model_validate(best_cv["result"])  # type: ignore
                    best_cv = CodeVersion.model_validate(best_cv)  # type: ignore
                except Exception as e:
                    logging.warning(f"Cannot parse best_code_version payload: {e}")
                    return

            safe_name = _make_filesystem_safe(best_cv.name or "best")
            base_filename = f"best_{self._checkpoint_counter}_{safe_name}"

            # Prepare metadata
            meta = {
                "code_version_id": best_cv.code_version_id,
                "name": best_cv.name,
                "kpi": best_kpi,
                "stdout": getattr(best_cv.result, "stdout", None) if best_cv.result else None,
            }

            # Check if using Databricks
            is_databricks = config and config.get('databricks')
            if is_databricks:
                # Save directly to Databricks volume using CLI (no local storage)
                await self._save_checkpoint_to_databricks_volume(
                    best_cv.code, 
                    json.dumps(meta, indent=4), 
                    base_filename, 
                    config
                )
            else:
                # Original behavior for local runs
                checkpoints_base = Path(project_absolute_path) / CHECKPOINTS_FOLDER
                checkpoints_base.mkdir(parents=True, exist_ok=True)
                
                code_path = checkpoints_base / f"{base_filename}.py"
                meta_path = checkpoints_base / f"{base_filename}.json"
                
                code_path.write_text(best_cv.code, encoding="utf-8")
                meta_path.write_text(json.dumps(meta, indent=4))

            self._checkpoint_counter += 1
        except Exception as e:
            logging.warning(f"Failed saving best checkpoint: {e}")

    async def _save_checkpoint_to_databricks_volume(self, code_content: str, meta_content: str, base_filename: str, config: dict):
        """Save checkpoint files directly to Databricks volume using CLI (following existing upload pattern)."""
        try:
            # Extract databricks configuration (same pattern as _databricks_run_python_code)
            if isinstance(config['databricks'], dict):
                databricks_config = config['databricks']
            else:
                databricks_config = config
            
            CLI = databricks_config.get('cli', "databricks")
            VOLUME_URI = databricks_config.get('volume_uri', "dbfs:/Volumes/workspace/default/volume")
            
            # Ensure checkpoints directory exists
            checkpoints_dir = f"{VOLUME_URI}/{CHECKPOINTS_FOLDER}"
            mkdir_result = subprocess.run([CLI, "fs", "mkdir", checkpoints_dir], 
                                        capture_output=True, text=True)
            # mkdir is okay to fail if directory already exists
            
            # Create remote paths
            remote_code_path = f"{checkpoints_dir}/{base_filename}.py"
            remote_meta_path = f"{checkpoints_dir}/{base_filename}.json"
            
            # Save code file using temp file + CLI upload pattern (following existing code)
            with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as f:
                f.write(code_content.encode())
                local_tmp_code = pathlib.Path(f.name)
            
            # Try uploading code file
            result = subprocess.run([CLI, "fs", "cp", str(local_tmp_code), remote_code_path,
                           "--overwrite", "--output", "json"], capture_output=True, text=True)
            if result.returncode != 0:
                print(f"❌ Failed to upload checkpoint code: {result.stderr}")
                return
            os.unlink(local_tmp_code)
            
            # Save metadata file using temp file + CLI upload pattern
            with tempfile.NamedTemporaryFile(delete=False, suffix=".json") as f:
                f.write(meta_content.encode())
                local_tmp_meta = pathlib.Path(f.name)
            
            # Try uploading metadata file
            result = subprocess.run([CLI, "fs", "cp", str(local_tmp_meta), remote_meta_path,
                           "--overwrite", "--output", "json"], capture_output=True, text=True)
            if result.returncode != 0:
                print(f"❌ Failed to upload checkpoint metadata: {result.stderr}")
                return
            os.unlink(local_tmp_meta)
            
            print(f"✅ Checkpoint uploaded to: {VOLUME_URI}/{CHECKPOINTS_FOLDER}/{base_filename}.*")
            
        except Exception as e:
            print(f"❌ Checkpoint upload error: {e}")


def _make_filesystem_safe(name):
    return re.sub(r"[/\\?%*:|\"<>\x7F\x00-\x1F]", '_', name)


def _run_python_code(code: str, python_path: str) -> CodeResult:
    import subprocess  # Ensure subprocess is imported
    start_time = time.time()
    # write the code to a temporary file
    with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as temp_file:
        temp_file.write(code.encode('utf-8'))
        temp_file_path = temp_file.name

    command = [python_path, temp_file_path]

    # run the command
    logging.info("running command: " + str(command))
    try:
        output = subprocess.run(command,
                                capture_output=True,
                                text=True,
                                input="",  # prevents it from blocking on stdin
                                timeout=settings.script_execution_timeout)  # Use centralized timeout
        return_code = output.returncode
        out = output.stdout
        err = output.stderr
    except subprocess.TimeoutExpired as e:
        # Handle timeout gracefully - mark as failed case and continue
        return_code = -9  # Standard timeout return code
        out = None
        err = f"Process timed out after {settings.script_execution_timeout} seconds"
        logging.info(f"Process timed out after {settings.script_execution_timeout} seconds")
    
    if isinstance(out, str) and out.strip() == "":
        out = None
    if isinstance(err, str) and err.strip() == "":
        err = None

    logging.info("stdout: " + str(out))
    logging.info("stderr: " + str(err))

    # delete the temporary file
    os.remove(temp_file_path)
    runtime_ms = int((time.time() - start_time) * 1000)
    return CodeResult(stdout=out, stderr=err, return_code=return_code, runtime_ms=runtime_ms)


def _databricks_run_python_code(code: str,
                                python_path: str = "python", 
                                config: dict = None) -> CodeResult:
    """
    Execute `code` on Databricks Serverless Jobs and collect stdout/err.
    
    Args:
        code: Python code to execute
        python_path: Python path (accepted for API parity, currently ignored)
        config: Configuration dict containing databricks settings. Supports two formats:
                Format 1 (nested): databricks: {cli: "databricks", volume_uri: "...", job: {...}}
                Format 2 (flat): databricks: true, cli: "databricks", volume_uri: "...", job: {...}
                Settings:
                - cli: Databricks CLI command (default: "databricks")
                - volume_uri: DBFS volume URI (default: "dbfs:/Volumes/workspace/default/volume")
                - code_path: Specific DBFS path for code file (optional, overrides volume_uri)
                - timeout: Job timeout (default: "30m")
                - job: Job configuration with name, tasks, environments
    
    Returns:
        CodeResult with stdout, stderr, return_code, and runtime_ms
    """
    # Extract databricks configuration with sensible defaults
    # Support both nested structure (databricks: {config}) and flat structure (databricks: true + top-level config)
    if config and config.get('databricks'):
        if isinstance(config['databricks'], dict):
            # New nested structure: databricks: {cli: ..., volume_uri: ..., etc}
            databricks_config = config['databricks']
            job_config = databricks_config.get('job', {})
        else:
            # Original flat structure: databricks: true + top-level fields
            databricks_config = config
            job_config = config.get('job', {})
    else:
        databricks_config = {}
        job_config = {}
    
    CLI = databricks_config.get('cli', "databricks")
    VOLUME_URI = databricks_config.get('volume_uri', "dbfs:/Volumes/workspace/default/volume")
    CODE_PATH = databricks_config.get('code_path', None)  # Specific code path if provided
    TIMEOUT = databricks_config.get('timeout', "30m")
    logging.info(f"CLI: {CLI}")
    logging.info(f"VOLUME_URI: {VOLUME_URI}")
    logging.info(f"CODE_PATH: {CODE_PATH}")
    logging.info(f"TIMEOUT: {TIMEOUT}")
    t0 = time.time()

    # ------------------------------------------------------------------ 1. write & upload
    with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as f:
        f.write(code.encode())
        local_tmp = pathlib.Path(f.name)

    # Always use temp filename to avoid overwriting original code_path
    # The code_path is only for reading the original script, never for writing iterations
    remote_uri = f"{VOLUME_URI}/{local_tmp.name}"
    subprocess.run([CLI, "fs", "cp", str(local_tmp), remote_uri,
                    "--overwrite", "--output", "json"], check=True)  # :contentReference[oaicite:7]{index=7}
    os.unlink(local_tmp) #does this delete the local temp file? 

    # ------------------------------------------------------------------ 2. build job JSON
    # job_config is already extracted above based on config structure
    
    # Build job name with template substitution
    job_name_template = job_config.get('name', 'run-<script-stem>-<timestamp>')
    job_name = job_name_template.replace('<script-stem>', local_tmp.stem).replace('<timestamp>', str(int(t0)))
    
    # Build tasks configuration
    default_task = {
        "task_key": "t",
        "spark_python_task": {"python_file": remote_uri},
        "environment_key": "default"
    }
    tasks = job_config.get('tasks', [default_task])
    # Update the python_file path in the first task
    if tasks and 'spark_python_task' in tasks[0]:
        tasks[0]['spark_python_task']['python_file'] = remote_uri
    
    # Build environments configuration
    default_environments = [{
        "environment_key": "default",
        "spec": {"client": "1"}
    }]
    environments = job_config.get('environments', default_environments)
    
    job_json = {
        "name": job_name,
        "tasks": tasks,
        "environments": environments
    }
    with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as spec:
        json.dump(job_json, spec)
        spec_path = spec.name

    # ------------------------------------------------------------------ 3. create job
    create = subprocess.run([CLI, "jobs", "create",
                             "--json", f"@{spec_path}", "--output", "json"],
                            text=True, capture_output=True, check=True)
    job_id = json.loads(create.stdout)["job_id"]          # :contentReference[oaicite:9]{index=9}
    
    # Clean up the job spec file
    os.unlink(spec_path)

    # ------------------------------------------------------------------ 4. run job (blocks)
    try:
        output = subprocess.run([CLI, "jobs", "run-now",
                                str(job_id),
                                "--timeout", TIMEOUT,
                                "--output", "json"],
                                text=True, input="", capture_output=True)
        
        # Check if job submission failed
        if output.returncode != 0 or not output.stdout.strip():
            # Job submission failed - return error result
            runtime_ms = int((time.time() - t0) * 1000)
            error_msg = "Databricks job submission failed"
            if output.stderr:
                error_msg += f": {output.stderr}"
            elif not output.stdout.strip():
                error_msg += ": Empty response from Databricks"
            
            logging.error(f"Job submission failed: stdout='{output.stdout}', stderr='{output.stderr}', returncode={output.returncode}")
            return CodeResult(stdout=None, stderr=error_msg, return_code=1, runtime_ms=runtime_ms)
        
        # Parse the job run response
        try:
            run_response = json.loads(output.stdout)
            run_id = run_response["run_id"]
        except (json.JSONDecodeError, KeyError) as e:
            runtime_ms = int((time.time() - t0) * 1000)
            logging.error(f"Failed to parse job run response: {e}, stdout='{output.stdout}'")
            return CodeResult(stdout=None, stderr=f"Failed to parse Databricks response: {e}", 
                            return_code=1, runtime_ms=runtime_ms)

    except subprocess.TimeoutExpired:
        # Handle timeout gracefully - mark as failed case and continue
        runtime_ms = int((time.time() - t0) * 1000)
        error_msg = f"Databricks job submission timed out after {TIMEOUT}"
        logging.info(error_msg)
        return CodeResult(stdout=None, stderr=error_msg, return_code=-9, runtime_ms=runtime_ms)

    # ------------------------------------------------------------------ 5. find child task-run
    try:
        get_run = subprocess.run([CLI, "jobs", "get-run", str(run_id),
                                  "--output", "json"],
                                 text=True, capture_output=True, check=True)
        task_run_id = json.loads(get_run.stdout)["tasks"][0]["run_id"]

        # ------------------------------------------------------------------ 6. fetch logs
        out_json = subprocess.run([CLI, "jobs", "get-run-output",
                                   str(task_run_id), "--output", "json"],
                                  text=True, capture_output=True, check=True)
        out = json.loads(out_json.stdout)
        logs = out.get("logs")
        result_state = out.get("metadata", {}).get("state", {}).get("result_state", "FAILED")

        rc = 0 if result_state == "SUCCESS" else 1
        runtime_ms = int((time.time() - t0) * 1000)

        # Extract error details if available
        stderr_msg = None if rc == 0 else "Task failed"
        if rc != 0 and out.get("error"):
            stderr_msg = f"Task failed: {out['error']}"

        return CodeResult(stdout=logs, stderr=stderr_msg, return_code=rc, runtime_ms=runtime_ms)
        
    except Exception as e:
        # Handle any errors in fetching run details/logs
        runtime_ms = int((time.time() - t0) * 1000)
        logging.error(f"Failed to fetch Databricks job results: {e}")
        return CodeResult(stdout=None, stderr=f"Failed to fetch job results: {e}", 
                        return_code=1, runtime_ms=runtime_ms)




def get_system_info(python_path: str) -> SystemInfo:
    return SystemInfo(
        python_libraries=_get_python_libraries(python_path),
        python_version=_get_python_version(python_path),
        os=sys.platform
    )


def _get_python_libraries(python_path: str) -> list[str]:
    try:
        # Use importlib.metadata to get installed packages (works in all Python 3.8+ environments)
        python_code = """
import importlib.metadata
for dist in importlib.metadata.distributions():
    print(f"{dist.metadata['Name']}=={dist.version}")
"""
        installed_libraries = subprocess.check_output(
            [python_path, "-c", python_code],
            universal_newlines=True
        ).strip()
        return [lib.strip() for lib in installed_libraries.split("\n") if lib.strip()]
    except subprocess.CalledProcessError:
        # If that fails, return empty list
        return []


def _get_python_version(python_path: str) -> str:
    return subprocess.check_output(
        [python_path, "--version"],
        universal_newlines=True
    ).strip()


workflow_runner = _WorkflowRunner()

    
    
