from __future__ import annotations

import json
import os
from pathlib import Path

from .analyzer import SQL2Analysis
from .emitter import SQL2Emitter
from .ir import SQL2ProgramIR


class SQL2DebugDumper:
    """Env-guarded debug dumps for SQL2 graph + program artifacts."""

    def __init__(self, *, emitter: SQL2Emitter):
        self._emitter = emitter

    def dump_model(
        self,
        *,
        run_id: int,
        model_version: int | None,
        analysis: SQL2Analysis,
        raw_program_ir: SQL2ProgramIR,
        optimized_program_ir: SQL2ProgramIR,
    ) -> None:
        dump_root = os.getenv("PYREL_SQL2_DEBUG_DUMP_DIR")
        if not dump_root:
            return
        dump_dir = Path(dump_root)
        dump_dir.mkdir(parents=True, exist_ok=True)

        prefix = f"sql2_run_{run_id:04d}_v{model_version if model_version is not None else 'na'}"
        self._write_json(dump_dir / f"{prefix}_graph.json", self._graph_payload(analysis))
        (dump_dir / f"{prefix}_raw_program.sql").write_text(
            self._program_sql_dump(raw_program_ir),
            encoding="utf-8",
        )
        (dump_dir / f"{prefix}_optimized_program.sql").write_text(
            self._program_sql_dump(optimized_program_ir),
            encoding="utf-8",
        )

    def _graph_payload(self, analysis: SQL2Analysis) -> dict[str, object]:
        graph = analysis.graph
        nodes = [
            {
                "node_id": node.node_id,
                "update_task_id": node.update_task_id,
                "target_relation_name": node.target_relation_name,
                "target_table_name": node.target_table_name,
                "key_var_ids": list(node.key_var_ids),
                "dependency_node_ids": list(node.dependency_node_ids),
            }
            for node in sorted(graph.nodes, key=lambda n: n.node_id)
        ]
        edges = [
            {
                "source_node_id": edge.source_node_id,
                "target_node_id": edge.target_node_id,
                "edge_type": edge.edge_type.value,
                "via_var_ids": list(edge.via_var_ids),
            }
            for edge in sorted(
                graph.edges,
                key=lambda e: (e.source_node_id, e.target_node_id, e.edge_type.value, e.via_var_ids),
            )
        ]
        return {
            "nodes": nodes,
            "edges": edges,
            "sccs": [list(scc) for scc in graph.sccs],
            "stratum_by_node": {str(node_id): stratum for node_id, stratum in sorted(graph.stratum_by_node.items())},
            "recursive_node_ids": sorted(graph.recursive_node_ids),
        }

    def _program_sql_dump(self, program: SQL2ProgramIR) -> str:
        lines: list[str] = []
        for stratum in program.strata:
            lines.append(f"-- Stratum {stratum.stratum_id}")
            if not stratum.operations and not stratum.semi_naive_loops:
                lines.append("-- (no operations)")
                lines.append("")
                continue
            for op in stratum.operations:
                lines.append(self._emitter.emit(op))
                lines.append("")
            for loop_ix, loop in enumerate(stratum.semi_naive_loops, start=1):
                lines.append(f"-- Semi-naive loop {loop_ix}")
                for table in loop.tables:
                    lines.append(f"--   table: {table.table_name}")
                    lines.append("--   base_query")
                    lines.append(self._emitter.emit(table.base_query))
                    lines.append("--   delta_query")
                    lines.append(self._emitter.emit(table.delta_query))
                    lines.append("--   merge_query")
                    lines.append(self._emitter.emit(table.merge_query))
                    lines.append("--   final_query")
                    lines.append(self._emitter.emit(table.final_query))
                    lines.append("")
        return "\n".join(lines).rstrip() + "\n"

    def _write_json(self, path: Path, payload: dict[str, object]) -> None:
        path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
