"""DAG generators for workflow orchestration.

Generates Airflow and Prefect code from workflow configurations.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from quicketl.config.workflow import WorkflowConfig


def generate_airflow_dag(
    config: WorkflowConfig,
    dag_id: str | None = None,
    schedule: str | None = None,
    base_path: str = ".",
) -> str:
    """Generate an Airflow DAG from a workflow configuration.

    Args:
        config: Workflow configuration
        dag_id: DAG ID (defaults to workflow name)
        schedule: Cron schedule (defaults to None/manual)
        base_path: Base path for pipeline files

    Returns:
        Python code for an Airflow DAG
    """
    dag_id = dag_id or config.name.replace(" ", "_").replace("-", "_")

    # Build imports
    lines = [
        '"""',
        f"Airflow DAG: {config.name}",
        "",
        f"{config.description}" if config.description else "Auto-generated from QuickETL workflow.",
        "",
        "Generated by: quicketl workflow generate --target airflow",
        '"""',
        "",
        "from datetime import datetime",
        "",
        "from airflow import DAG",
        "from airflow.operators.python import PythonOperator",
        "",
        "from quicketl.pipeline import run_pipeline",
        "",
        "",
        "# =============================================================================",
        "# Pipeline Tasks",
        "# =============================================================================",
        "",
    ]

    # Generate task functions
    for stage in config.stages:
        for pipeline_ref in stage.pipelines:
            func_name = _sanitize_name(pipeline_ref.resolved_name)
            pipeline_path = f"{base_path}/{pipeline_ref.path}"

            lines.extend([
                f"def run_{func_name}(**context):",
                f'    """Run {pipeline_ref.resolved_name} pipeline."""',
                "    result = run_pipeline(",
                f'        "{pipeline_path}",',
            ])

            # Add variables
            if config.variables or pipeline_ref.variables:
                merged = {**config.variables, **pipeline_ref.variables}
                lines.append("        variables={")
                for key, value in merged.items():
                    lines.append(f'            "{key}": "{value}",')
                lines.append("        },")

            lines.extend([
                "    )",
                "    if result.failed:",
                f'        raise Exception(f"Pipeline {pipeline_ref.resolved_name} failed: {{result.error}}")',
                "    return result.to_dict()",
                "",
                "",
            ])

    # Generate DAG
    schedule_str = f'"{schedule}"' if schedule else "None"

    lines.extend([
        "# =============================================================================",
        "# DAG Definition",
        "# =============================================================================",
        "",
        "default_args = {",
        '    "owner": "airflow",',
        "    \"retries\": 1,",
        "}",
        "",
        "with DAG(",
        f'    dag_id="{dag_id}",',
        f'    description="{config.description}",',
        f"    schedule={schedule_str},",
        "    start_date=datetime(2025, 1, 1),",
        "    catchup=False,",
        "    default_args=default_args,",
        '    tags=["quicketl"],',
        ") as dag:",
        "",
    ])

    # Generate task definitions
    task_ids: dict[str, list[str]] = {}  # stage_name -> [task_ids]

    for stage in config.stages:
        task_ids[stage.name] = []

        for pipeline_ref in stage.pipelines:
            func_name = _sanitize_name(pipeline_ref.resolved_name)
            task_id = f"{stage.name}_{func_name}"
            task_ids[stage.name].append(task_id)

            lines.extend([
                f"    {task_id} = PythonOperator(",
                f'        task_id="{task_id}",',
                f"        python_callable=run_{func_name},",
                "    )",
                "",
            ])

    # Generate dependencies
    lines.append("    # Dependencies")
    for stage in config.stages:
        if stage.depends_on:
            current_tasks = task_ids[stage.name]
            for dep_stage in stage.depends_on:
                upstream_tasks = task_ids.get(dep_stage, [])
                for upstream in upstream_tasks:
                    for current in current_tasks:
                        lines.append(f"    {upstream} >> {current}")

    return "\n".join(lines)


def generate_prefect_flow(
    config: WorkflowConfig,
    flow_name: str | None = None,
    base_path: str = ".",
) -> str:
    """Generate a Prefect flow from a workflow configuration.

    Args:
        config: Workflow configuration
        flow_name: Flow name (defaults to workflow name)
        base_path: Base path for pipeline files

    Returns:
        Python code for a Prefect flow
    """
    flow_name = flow_name or config.name.replace(" ", "_").replace("-", "_")

    lines = [
        '"""',
        f"Prefect Flow: {config.name}",
        "",
        f"{config.description}" if config.description else "Auto-generated from QuickETL workflow.",
        "",
        "Generated by: quicketl workflow generate --target prefect",
        '"""',
        "",
        "from prefect import flow, task",
        "from prefect.futures import wait",
        "",
        "from quicketl.pipeline import run_pipeline",
        "",
        "",
        "# =============================================================================",
        "# Pipeline Tasks",
        "# =============================================================================",
        "",
    ]

    # Generate task functions
    for stage in config.stages:
        for pipeline_ref in stage.pipelines:
            func_name = _sanitize_name(pipeline_ref.resolved_name)
            pipeline_path = f"{base_path}/{pipeline_ref.path}"

            lines.extend([
                "@task",
                f"def run_{func_name}():",
                f'    """Run {pipeline_ref.resolved_name} pipeline."""',
                "    result = run_pipeline(",
                f'        "{pipeline_path}",',
            ])

            # Add variables
            if config.variables or pipeline_ref.variables:
                merged = {**config.variables, **pipeline_ref.variables}
                lines.append("        variables={")
                for key, value in merged.items():
                    lines.append(f'            "{key}": "{value}",')
                lines.append("        },")

            lines.extend([
                "    )",
                "    if result.failed:",
                f'        raise Exception(f"Pipeline {pipeline_ref.resolved_name} failed: {{result.error}}")',
                "    return result",
                "",
                "",
            ])

    # Generate flow
    lines.extend([
        "# =============================================================================",
        "# Flow Definition",
        "# =============================================================================",
        "",
        "@flow(name=\"" + flow_name + "\")",
        f"def {flow_name}():",
        '    """',
        f"    {config.description}" if config.description else f"    Execute {config.name} workflow.",
        "",
        "    Stages:",
    ])

    for stage in config.stages:
        deps = f" (depends on: {', '.join(stage.depends_on)})" if stage.depends_on else ""
        lines.append(f"      - {stage.name}{deps}")

    lines.extend([
        '    """',
        "",
    ])

    # Generate stage execution with dependencies
    stage_results: dict[str, str] = {}  # stage_name -> result variable name

    for stage in config.stages:
        stage_var = f"{stage.name}_results"
        stage_results[stage.name] = stage_var

        # Wait for dependencies
        if stage.depends_on:
            for dep in stage.depends_on:
                dep_var = stage_results.get(dep)
                if dep_var:
                    lines.append(f"    # Wait for {dep} stage")
                    lines.append(f"    wait({dep_var})")
                    lines.append("")

        # Submit tasks
        lines.append(f"    # Stage: {stage.name}")

        if stage.parallel:
            lines.append(f"    {stage_var} = []")
            for pipeline_ref in stage.pipelines:
                func_name = _sanitize_name(pipeline_ref.resolved_name)
                lines.append(f"    {stage_var}.append(run_{func_name}.submit())")
        else:
            lines.append(f"    {stage_var} = []")
            for pipeline_ref in stage.pipelines:
                func_name = _sanitize_name(pipeline_ref.resolved_name)
                lines.append(f"    {stage_var}.append(run_{func_name}())")

        lines.append("")

    # Return statement
    lines.extend([
        "    return {",
    ])
    for stage_name, stage_var in stage_results.items():
        lines.append(f'        "{stage_name}": {stage_var},')
    lines.extend([
        "    }",
        "",
        "",
        'if __name__ == "__main__":',
        f"    {flow_name}()",
        "",
    ])

    return "\n".join(lines)


def _sanitize_name(name: str) -> str:
    """Convert a name to a valid Python identifier."""
    # Replace common separators with underscores
    result = name.replace("-", "_").replace(".", "_").replace(" ", "_")
    # Remove any remaining invalid characters
    result = "".join(c if c.isalnum() or c == "_" else "_" for c in result)
    # Ensure it doesn't start with a number
    if result and result[0].isdigit():
        result = "_" + result
    return result
