"""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",
        "",
        "Environment Variables:",
        "    QUICKETL_BASE_PATH: Base path for pipeline files (default: current directory)",
        '"""',
        "",
        "import os",
        "from datetime import datetime, timedelta",
        "from pathlib import Path",
        "",
        "from airflow import DAG",
        "from airflow.operators.python import PythonOperator",
        "from airflow.utils.task_group import TaskGroup",
        "",
        "from quicketl.pipeline import run_pipeline",
        "",
        "",
        "# =============================================================================",
        "# Configuration",
        "# =============================================================================",
        "",
        "# Base path for pipeline files - configurable via environment variable",
        f'BASE_PATH = os.environ.get("QUICKETL_BASE_PATH", "{base_path}")',
        "",
        "",
        "# =============================================================================",
        "# Pipeline Tasks",
        "# =============================================================================",
        "",
    ]

    # Generate task functions
    for stage in config.stages:
        for pipeline_ref in stage.pipelines:
            func_name = _sanitize_name(pipeline_ref.resolved_name)
            # Use relative path from base
            pipeline_path = pipeline_ref.path

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

            # 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",',
        '    "depends_on_past": False,',
        "    \"retries\": 1,",
        "    \"retry_delay\": timedelta(minutes=5),",
        "}",
        "",
        "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"],',
        '    doc_md=__doc__,',
        ") as dag:",
        "",
    ])

    # Generate task groups and tasks
    task_groups: dict[str, str] = {}  # stage_name -> task_group_var

    for stage in config.stages:
        group_var = f"tg_{stage.name}"
        task_groups[stage.name] = group_var

        lines.extend([
            f'    with TaskGroup(group_id="{stage.name}") as {group_var}:',
        ])

        for pipeline_ref in stage.pipelines:
            func_name = _sanitize_name(pipeline_ref.resolved_name)
            task_id = func_name

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

        lines.append("")

    # Generate stage-level dependencies (cleaner than task-to-task)
    lines.append("    # Stage dependencies")
    for stage in config.stages:
        if stage.depends_on:
            current_group = task_groups[stage.name]
            for dep_stage in stage.depends_on:
                upstream_group = task_groups.get(dep_stage)
                if upstream_group:
                    lines.append(f"    {upstream_group} >> {current_group}")

    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",
        "",
        "Environment Variables:",
        "    QUICKETL_BASE_PATH: Base path for pipeline files (default: current directory)",
        "",
        "Usage:",
        "    # Run locally",
        "    python flows/medallion_flow.py",
        "",
        "    # Deploy to Prefect Cloud",
        "    python flows/medallion_flow.py --deploy",
        "",
        "    # Serve locally (for development)",
        "    python flows/medallion_flow.py --serve",
        '"""',
        "",
        "import argparse",
        "import os",
        "from pathlib import Path",
        "",
        "from prefect import flow, task",
        "from prefect.futures import wait",
        "",
        "from quicketl.pipeline import run_pipeline",
        "",
        "",
        "# =============================================================================",
        "# Configuration",
        "# =============================================================================",
        "",
        "# Base path for pipeline files - configurable via environment variable",
        f'BASE_PATH = os.environ.get("QUICKETL_BASE_PATH", "{base_path}")',
        "",
        "",
        "# =============================================================================",
        "# 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 = pipeline_ref.path

            lines.extend([
                f'@task(name="{pipeline_ref.resolved_name}", retries=2, retry_delay_seconds=60)',
                f"def run_{func_name}():",
                f'    """Run {pipeline_ref.resolved_name} pipeline."""',
                f'    pipeline_file = str(Path(BASE_PATH) / "{pipeline_path}")',
                "    result = run_pipeline(",
                "        pipeline_file,",
            ])

            # 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 flow
    lines.extend([
        "# =============================================================================",
        "# Flow Definition",
        "# =============================================================================",
        "",
        f'@flow(name="{flow_name}", log_prints=True)',
        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

        # Submit tasks first
        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())")
            # Wait for parallel tasks to complete before next stage
            lines.append(f"    wait({stage_var})")
        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([
        "    }",
        "",
        "",
        "# =============================================================================",
        "# Deployment & CLI",
        "# =============================================================================",
        "",
        'if __name__ == "__main__":',
        '    parser = argparse.ArgumentParser(description="Run QuickETL workflow")',
        '    parser.add_argument("--deploy", action="store_true", help="Deploy to Prefect Cloud")',
        '    parser.add_argument("--serve", action="store_true", help="Serve locally for development")',
        '    parser.add_argument("--cron", type=str, default=None, help="Cron schedule for deployment")',
        "    args = parser.parse_args()",
        "",
        "    if args.deploy:",
        "        # Deploy to Prefect Cloud",
        f"        {flow_name}.deploy(",
        f'            name="{flow_name}-deployment",',
        '            work_pool_name=os.environ.get("PREFECT_WORK_POOL", "default-agent-pool"),',
        "            cron=args.cron,",
        "        )",
        "    elif args.serve:",
        "        # Serve locally for development",
        f"        {flow_name}.serve(",
        f'            name="{flow_name}-dev",',
        "            cron=args.cron,",
        "        )",
        "    else:",
        "        # Run immediately",
        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
