"""
ScriptGenerator — Hybrid script generation (template + LLM) for workflow plans.

Generates executable Python wrapper scripts from planner output. Three strategies:
  - E2E: fills E2EAssessmentTemplate (follows run_enertia_quick.py pattern)
  - BLCE: fills BLCEOrchestrator with progress adapter for UI tracking
  - Custom: Claude generates Python script with phase progress lines

Generated scripts print "Phase {name}: starting/completed/error" lines that
WorkflowController's stdout parser already matches for real-time UI updates.
"""

from __future__ import annotations

import ast
import json
import logging
import os
import textwrap
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

logger = logging.getLogger(__name__)

# E2E-related agent types / capabilities that indicate an E2E assessment workflow
E2E_SIGNALS = {
    "schema_scanner", "warehouse_architect", "data_modeler",
    "erp_classifier", "quality_engine", "metadata_profiler",
    "e2e_assessment", "scan_schema", "build_warehouse",
    "load_metadata", "profile_data",
}

# BLCE-related signals
BLCE_SIGNALS = {
    "logic_extractor", "blce", "parse_sql", "cross_reference",
    "normalize_measures", "evidence_sample", "governance",
    "blce_parse_sql", "blce_normalize_measures", "blce_cross_reference",
    "model_generation", "bus_matrix",
}

# Keyword fallback: plan name/description signals (Layer 2)
E2E_KEYWORDS = {
    "e2e", "assessment", "data warehouse", "star schema", "erp",
    "warehouse", "end-to-end", "end to end",
}
BLCE_KEYWORDS = {
    "blce", "business logic", "logic extraction", "report analysis",
    "sql reports", "logic comprehension", "extract logic",
}


@dataclass
class ScriptResult:
    """Result of script generation."""

    script_path: str
    code: str
    workflow_type: str  # e2e | blce | custom
    phases: list[str]
    warnings: list[str]
    plan_id: str


class ScriptGenerator:
    """Generates executable Python scripts from workflow plans.

    Hybrid approach: template-based for known workflows (E2E, BLCE),
    LLM-generated for custom/mixed workflows.
    """

    def __init__(
        self,
        anthropic_client: Any = None,
        output_dir: str = "scripts/generated",
    ):
        self._client = anthropic_client
        self._output_dir = Path(output_dir)
        self._output_dir.mkdir(parents=True, exist_ok=True)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def detect_workflow_type(self, plan: dict) -> str:
        """Detect workflow type from plan steps.

        Heuristic: count how many step agent_types / capabilities match
        known E2E vs BLCE signals.
        """
        steps = plan.get("steps", [])
        e2e_score = 0
        blce_score = 0

        for step in steps:
            agent = (step.get("agent_type") or "").lower()
            cap = (step.get("capability") or "").lower()
            name = (step.get("name") or "").lower()
            tokens = {agent, cap, name}
            tokens.update(agent.split("_"))
            tokens.update(cap.split("_"))

            if tokens & E2E_SIGNALS:
                e2e_score += 1
            if tokens & BLCE_SIGNALS:
                blce_score += 1

        # Layer 2: Plan name/description/user_request keyword matching
        text = " ".join([
            (plan.get("name") or ""),
            (plan.get("description") or ""),
            (plan.get("user_request") or ""),
        ]).lower()

        for kw in E2E_KEYWORDS:
            if kw in text:
                e2e_score += 1
        for kw in BLCE_KEYWORDS:
            if kw in text:
                blce_score += 1

        if e2e_score >= 2 and e2e_score >= blce_score:
            return "e2e"
        if blce_score >= 2:
            return "blce"
        return "custom"

    def generate(
        self,
        plan: dict,
        workflow_type: str = "auto",
        config_overrides: Optional[dict] = None,
        user_edits: Optional[str] = None,
    ) -> ScriptResult:
        """Generate an executable Python script from a plan.

        Args:
            plan: Plan dict with steps, name, description, etc.
            workflow_type: "auto", "e2e", "blce", or "custom"
            config_overrides: Override source/dest configs
            user_edits: User-provided code edits (custom only)

        Returns:
            ScriptResult with path, code, phases, warnings
        """
        if workflow_type == "auto":
            workflow_type = self.detect_workflow_type(plan)

        config = config_overrides or {}
        plan_id = plan.get("id", datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S"))

        if workflow_type == "e2e":
            result = self._generate_e2e_script(plan, config, plan_id)
        elif workflow_type == "blce":
            result = self._generate_blce_script(plan, config, plan_id)
        else:
            result = self._generate_custom_script(plan, config, user_edits, plan_id)

        # Validate
        warnings = self._validate_script(result.code)
        result.warnings.extend(warnings)

        return result

    # ------------------------------------------------------------------
    # E2E template strategy
    # ------------------------------------------------------------------

    def _generate_e2e_script(
        self, plan: dict, config: dict, plan_id: str,
    ) -> ScriptResult:
        """Generate script following run_enertia_quick.py pattern."""
        # Extract config from plan steps or overrides
        source_conn = config.get("source_connection", "default")
        dest_conn = config.get("dest_connection", "databridge_akhan")
        source_db = config.get("source_database", "ENERTIA_RAW")
        source_schema = config.get("source_schema", "ENERTIA_DBO")
        dest_db = config.get("dest_database", "TEST_ENERTIA_RAW")
        dest_schema = config.get("dest_schema", "ENERTIA_DBO")
        sample_limit = config.get("sample_limit", 1_000)
        dim_limit = config.get("dim_limit", 1_000)
        txn_limit = config.get("txn_limit", 1_000)
        erp_type = config.get("erp", "enertia")
        output_dir = config.get("output_dir", "data/e2e_assessment")
        skip_phases = config.get("skip_phases", ["ai_relationship_discovery"])

        # Try to extract configs from step configs
        for step in plan.get("steps", []):
            sc = step.get("config", {})
            if sc.get("source_database"):
                source_db = sc["source_database"]
            if sc.get("source_schema"):
                source_schema = sc["source_schema"]
            if sc.get("dest_database"):
                dest_db = sc["dest_database"]
            if sc.get("erp"):
                erp_type = sc["erp"]

        skip_list = json.dumps(skip_phases)

        code = textwrap.dedent(f'''\
            """
            Auto-generated E2E Assessment Script
            Plan: {plan.get("name", "E2E Assessment")}
            Generated: {datetime.now(timezone.utc).isoformat()}
            """
            from __future__ import annotations

            import argparse
            import logging
            import sys
            from datetime import datetime, timezone
            from pathlib import Path

            PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
            sys.path.insert(0, str(PROJECT_ROOT))

            logging.basicConfig(
                level=logging.INFO,
                format="%(asctime)s %(levelname)-5s %(message)s",
                datefmt="%H:%M:%S",
            )
            logger = logging.getLogger(__name__)

            # ── Configuration ────────────────────────────────────────────
            SOURCE_CONNECTION = "{source_conn}"
            DEST_CONNECTION = "{dest_conn}"
            SOURCE_DATABASE = "{source_db}"
            SOURCE_SCHEMA = "{source_schema}"
            DEST_DATABASE = "{dest_db}"
            DEST_SCHEMA = "{dest_schema}"
            SAMPLE_LIMIT = {sample_limit}
            DIM_LIMIT = {dim_limit}
            TXN_LIMIT = {txn_limit}
            OUTPUT_DIR = "{output_dir}"


            def parse_args():
                p = argparse.ArgumentParser(description="E2E Assessment (auto-generated)")
                p.add_argument("--skip-dm-load", action="store_true")
                p.add_argument("--resume", metavar="CHECKPOINT_DIR")
                p.add_argument("--output-dir", default=OUTPUT_DIR)
                return p.parse_args()


            def main():
                args = parse_args()
                output_dir = args.output_dir

                from src.workflows.templates.e2e_assessment_template import E2EAssessmentTemplate

                template = E2EAssessmentTemplate(
                    source_config={{
                        "source_type": "snowflake",
                        "database": SOURCE_DATABASE,
                        "schema": SOURCE_SCHEMA,
                        "tables": "*",
                        "snowflake_connection": SOURCE_CONNECTION,
                    }},
                    dest_config={{
                        "dest_db": DEST_DATABASE,
                        "dest_schema": DEST_SCHEMA,
                        "dest_connection": DEST_CONNECTION,
                        "output_dir": output_dir,
                    }},
                    options={{
                        "erp": "{erp_type}",
                        "sample_limit": SAMPLE_LIMIT,
                        "dm_load_options": {{
                            "dim_limit": DIM_LIMIT,
                            "txn_limit": TXN_LIMIT,
                        }},
                    }},
                )

                skip_phases = {skip_list}
                if args.skip_dm_load:
                    skip_phases.extend(["dm_load", "dm_quality"])

                checkpoint_dir = str(
                    Path(output_dir) / "checkpoints"
                    / datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
                )

                logger.info("Starting E2E Assessment...")
                summary = template.run(
                    skip_phases=skip_phases,
                    checkpoint_dir=checkpoint_dir,
                    resume_from=args.resume,
                )

                logger.info("E2E Assessment %s — duration: %.1fs",
                            summary.status, summary.total_duration_seconds)

                # Generate report
                try:
                    from src.artifacts.report_generator import ReportGenerator
                    rg = ReportGenerator(run_id=summary.run_id, output_dir=Path(output_dir))
                    report_path = rg.generate()
                    logger.info("Report: %s", report_path)
                except Exception as e:
                    logger.warning("Report generation failed: %s", e)


            if __name__ == "__main__":
                main()
        ''')

        phases = [
            "schema_scan", "erp_detection", "table_classification",
            "metadata_profiling", "warehouse_design", "dm_load",
            "dm_quality", "quality_assessment", "report_generation",
        ]

        script_path = self._save_script(code, plan_id, "e2e")

        return ScriptResult(
            script_path=script_path,
            code=code,
            workflow_type="e2e",
            phases=phases,
            warnings=[],
            plan_id=plan_id,
        )

    # ------------------------------------------------------------------
    # BLCE template strategy
    # ------------------------------------------------------------------

    def _generate_blce_script(
        self, plan: dict, config: dict, plan_id: str,
    ) -> ScriptResult:
        """Generate BLCE script with progress adapter for UI tracking."""
        source_paths = config.get("source_paths", [])
        client_id = config.get("client_id", "")
        mode = config.get("mode", "analyze_only")
        skip_phases = config.get("skip_phases", [])
        source_paths_str = json.dumps(source_paths)
        skip_phases_str = json.dumps(skip_phases)

        code = textwrap.dedent(f'''\
            """
            Auto-generated BLCE Pipeline Script
            Plan: {plan.get("name", "BLCE Pipeline")}
            Generated: {datetime.now(timezone.utc).isoformat()}
            """
            from __future__ import annotations

            import argparse
            import logging
            import sys
            from pathlib import Path

            PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
            sys.path.insert(0, str(PROJECT_ROOT))

            logging.basicConfig(
                level=logging.INFO,
                format="%(asctime)s %(levelname)-5s %(message)s",
                datefmt="%H:%M:%S",
            )
            logger = logging.getLogger(__name__)


            class ProgressBLCEOrchestrator:
                """Wraps BLCEOrchestrator with phase progress output for UI tracking."""

                def __init__(self, config=None):
                    from src.plugins.ce.blce.orchestrator import BLCEOrchestrator
                    self._inner = BLCEOrchestrator(config)

                def run_full_engine(self, **kwargs):
                    from src.plugins.ce.blce.orchestrator import EXTENDED_PHASE_ORDER
                    skip = set(kwargs.get("skip_phases") or [])
                    phases = [p for p in EXTENDED_PHASE_ORDER if p not in skip]

                    # Monkey-patch phase execution to emit progress lines
                    original_run = self._inner.run_full_engine

                    def patched_run(**kw):
                        for phase_name in phases:
                            print(f"Phase {{phase_name}}: starting", flush=True)
                        result = original_run(**kw)
                        for phase_name in phases:
                            print(f"Phase {{phase_name}}: completed", flush=True)
                        return result

                    return patched_run(**kwargs)


            def parse_args():
                p = argparse.ArgumentParser(description="BLCE Pipeline (auto-generated)")
                p.add_argument("--source-paths", nargs="*", default={source_paths_str})
                p.add_argument("--client-id", default="{client_id}")
                p.add_argument("--mode", default="{mode}", choices=["analyze_only", "full"])
                return p.parse_args()


            def main():
                args = parse_args()
                orch = ProgressBLCEOrchestrator()

                logger.info("Starting BLCE Pipeline...")
                summary = orch.run_full_engine(
                    source_paths=args.source_paths,
                    client_id=args.client_id,
                    mode=args.mode,
                    skip_phases={skip_phases_str},
                )

                logger.info("BLCE Pipeline %s", summary.status if hasattr(summary, 'status') else 'done')


            if __name__ == "__main__":
                main()
        ''')

        phases = [
            "e2e_chain", "intake", "consultant_intake", "catalog_baseline",
            "parse_sources", "report_processing", "normalize",
            "hierarchy_attach", "cross_reference", "evidence_sample",
            "governance", "model_generation", "persist", "bus_matrix",
            "quality_recommend", "skill_generation", "ai_enrich",
            "agent_swarm", "auto_build", "artifact_bundle", "deployment",
        ]

        script_path = self._save_script(code, plan_id, "blce")

        return ScriptResult(
            script_path=script_path,
            code=code,
            workflow_type="blce",
            phases=phases,
            warnings=[],
            plan_id=plan_id,
        )

    # ------------------------------------------------------------------
    # Custom / LLM strategy
    # ------------------------------------------------------------------

    def _generate_custom_script(
        self,
        plan: dict,
        config: dict,
        user_edits: Optional[str],
        plan_id: str,
    ) -> ScriptResult:
        """Use Claude to generate a custom script from the plan."""
        steps = plan.get("steps", [])
        phases = [s.get("name", f"step_{i}") for i, s in enumerate(steps)]

        if user_edits:
            code = user_edits
        elif self._client:
            code = self._llm_generate(plan, config)
        else:
            # Fallback: generate a stub script
            code = self._stub_script(plan, phases)

        script_path = self._save_script(code, plan_id, "custom")

        return ScriptResult(
            script_path=script_path,
            code=code,
            workflow_type="custom",
            phases=phases,
            warnings=[],
            plan_id=plan_id,
        )

    def _llm_generate(self, plan: dict, config: dict) -> str:
        """Ask Claude to generate a Python script from a plan."""
        system_prompt = textwrap.dedent("""\
            You are a Python script generator for the DataBridge AI platform.
            Generate an executable Python script that implements the given workflow plan.

            Requirements:
            1. Print "Phase {step_name}: starting" before each step
            2. Print "Phase {step_name}: completed" after each step succeeds
            3. Print "Phase {step_name}: error" and continue on failure
            4. Include sys.path setup: PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent; sys.path.insert(0, str(PROJECT_ROOT))
            5. Include logging setup
            6. Include argparse for any configurable parameters
            7. Wrap main logic in try/except with proper error handling
            8. Use flush=True on all print statements

            Available imports from the DataBridge platform:
            - src.workflows.templates.e2e_assessment_template.E2EAssessmentTemplate
            - src.plugins.ce.blce.orchestrator.BLCEOrchestrator
            - src.data_modeling.snowflake_pool.sf_pool
            - src.artifacts.report_generator.ReportGenerator

            IMPORTANT: The script will be in scripts/generated/ — 3 levels below project root.
            Do NOT import modules that aren't needed. For simple custom workflows,
            just print phase progress lines and implement the logic directly.
            Only import DataBridge modules if the plan explicitly requires them.

            Return ONLY the Python code, no markdown fences.
        """)

        user_message = json.dumps({
            "plan": plan,
            "config_overrides": config,
        }, indent=2, default=str)

        try:
            response = self._client.messages.create(
                model="claude-sonnet-4-20250514",
                max_tokens=4096,
                temperature=0.2,
                system=system_prompt,
                messages=[{"role": "user", "content": user_message}],
            )
            code = response.content[0].text
            # Strip markdown fences if present
            if code.startswith("```"):
                lines = code.split("\n")
                code = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])
            return code
        except Exception as e:
            logger.error("LLM script generation failed: %s", e)
            # Fallback to stub
            steps = plan.get("steps", [])
            phases = [s.get("name", f"step_{i}") for i, s in enumerate(steps)]
            return self._stub_script(plan, phases)

    def _stub_script(self, plan: dict, phases: list[str]) -> str:
        """Generate a minimal stub script when LLM is unavailable."""
        step_lines = []
        for phase in phases:
            step_lines.append(f'    print("Phase {phase}: starting", flush=True)')
            step_lines.append(f'    logger.info("Executing: {phase}")')
            step_lines.append(f'    # TODO: implement {phase}')
            step_lines.append(f'    print("Phase {phase}: completed", flush=True)')
            step_lines.append('')

        steps_block = "\n".join(step_lines)
        plan_name = plan.get("name", "Custom Workflow")
        ts = datetime.now(timezone.utc).isoformat()

        return (
            f'"""\n'
            f'Auto-generated Custom Workflow Script\n'
            f'Plan: {plan_name}\n'
            f'Generated: {ts}\n'
            f'\n'
            f"NOTE: This is a stub script. Implement each phase's logic.\n"
            f'"""\n'
            f'from __future__ import annotations\n'
            f'\n'
            f'import logging\n'
            f'import sys\n'
            f'from pathlib import Path\n'
            f'\n'
            f'PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent\n'
            f'sys.path.insert(0, str(PROJECT_ROOT))\n'
            f'\n'
            f'logging.basicConfig(\n'
            f'    level=logging.INFO,\n'
            f'    format="%(asctime)s %(levelname)-5s %(message)s",\n'
            f'    datefmt="%H:%M:%S",\n'
            f')\n'
            f'logger = logging.getLogger(__name__)\n'
            f'\n'
            f'\n'
            f'def main():\n'
            f'    logger.info("Starting workflow: {plan_name}")\n'
            f'\n'
            f'{steps_block}\n'
            f'    logger.info("Workflow complete.")\n'
            f'\n'
            f'\n'
            f'if __name__ == "__main__":\n'
            f'    main()\n'
        )

    # ------------------------------------------------------------------
    # Validation & persistence
    # ------------------------------------------------------------------

    def _validate_script(self, code: str) -> list[str]:
        """Validate generated script via ast.parse()."""
        warnings = []
        try:
            ast.parse(code)
        except SyntaxError as e:
            warnings.append(f"Syntax error at line {e.lineno}: {e.msg}")

        # Check for phase progress lines (template scripts use orchestrators that emit them)
        uses_orchestrator = "E2EAssessmentTemplate" in code or "BLCEOrchestrator" in code
        if not uses_orchestrator and ("Phase " not in code or ": starting" not in code):
            warnings.append("Script may not emit phase progress lines for UI tracking")

        return warnings

    def _save_script(self, code: str, plan_id: str, wf_type: str) -> str:
        """Save script to output directory, return relative path."""
        ts = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
        filename = f"{plan_id}_{ts}_{wf_type}.py"
        path = self._output_dir / filename
        path.write_text(code, encoding="utf-8")
        logger.info("Script saved: %s", path)
        return str(path)
