#!/usr/bin/env python3
"""
Plot unit for both AP and Modem logs.

Runs AFTER:
- archive/reorganize
- aplog filter (produces per-log *_filter structure)
- mdlog filter (produces auto_analysis/*.csv)

Outputs go under:
  <session_root>/<layout.analysis>/plots/{aplog,mdlog}/...
"""

from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, List, Optional

from fs_utils import ensure_dir, run
from workflow_layout import SessionLayout


def _safe_symlink(src: Path, dst: Path) -> bool:
    try:
        if dst.exists() or dst.is_symlink():
            return True
        os.symlink(str(src), str(dst))
        return True
    except Exception:
        return False


def _find_dirs(root: Path, name: str) -> List[Path]:
    if not root.exists():
        return []
    return [p for p in root.rglob("*") if p.is_dir() and p.name == name]


@dataclass(frozen=True)
class PlotUnit:
    layout: SessionLayout

    # mdlog plot tool
    md_logplot_script: Path
    md_subids: str = "1"

    # aplog plot tools
    ap_reg_signal_plot_script: Path = Path()
    ap_tput_plot_script: Path = Path()
    ap_uid_tput_plot_script: Path = Path()
    ap_modem_activity_plot_script: Path = Path()
    ap_window_overview_plot_script: Path = Path()
    ap_serving_cell_plot_script: Path = Path()

    def _plots_root(self, session_root: Path) -> Path:
        return ensure_dir(session_root / self.layout.analysis / "plots")

    def _run_mdlog_plots(self, session_root: Path) -> List[Path]:
        out_root = ensure_dir(self._plots_root(session_root) / "mdlog")
        mdlog_filter = session_root / self.layout.mdlog_filter
        generated: List[Path] = []

        # mdlog_filter_unit moves auto_analysis directories into mdlog_filter/
        auto_dirs = _find_dirs(mdlog_filter, "auto_analysis")
        for auto_dir in auto_dirs:
            csvs = list(auto_dir.glob("*.csv"))
            if not csvs:
                continue

            outdir = ensure_dir(out_root / auto_dir.parent.name)
            cmd = [
                "python3",
                str(self.md_logplot_script),
                str(auto_dir),
                "--subids",
                self.md_subids,
                "--outdir",
                str(outdir),
                "--legend-mode",
                "global",
                "--with-combined",
                "--multi-page",
                "--page-layout",
                "grid",
                "--export-single-png",
            ]
            try:
                run(cmd)
                generated.append(outdir)
            except Exception:
                pass

        return generated

    def _run_aplog_plots(self, session_root: Path) -> List[Path]:
        out_root = ensure_dir(self._plots_root(session_root) / "aplog")
        aplog_filter = session_root / self.layout.aplog_filter
        generated: List[Path] = []

        if not aplog_filter.exists():
            return generated

        # Each *_filter dir is treated as one AP plotting unit.
        for filt in sorted([p for p in aplog_filter.iterdir() if p.is_dir()]):
            plot_dir = ensure_dir(out_root / filt.name)

            # 1) Registration & signal plot (expects sim1/sim2 dirs under --log-dir).
            sim1 = filt / "sim1"
            sim2 = filt / "sim2"
            if sim1.exists() or sim2.exists():
                _safe_symlink(sim1, plot_dir / "sim1") if sim1.exists() else None
                _safe_symlink(sim2, plot_dir / "sim2") if sim2.exists() else None
                try:
                    run(
                        [
                            "python3",
                            str(self.ap_reg_signal_plot_script),
                            "--log-dir",
                            str(plot_dir),
                            "--sim",
                            "all",
                            "--dds",
                            "sim1",
                        ]
                    )
                    generated.append(plot_dir)
                except Exception:
                    pass

            # 2) Throughput plot
            for f in filt.rglob("data_net_netifstats.txt"):
                try:
                    run(
                        [
                            "python3",
                            str(self.ap_tput_plot_script),
                            "--input",
                            str(f),
                            "--outdir",
                            str(plot_dir),
                        ]
                    )
                    generated.append(plot_dir)
                except Exception:
                    pass

            # 3) Per-UID throughput and traffic plots
            for pattern in ("data_net_netifstats.txt", "data_net_uidstats.txt"):
                for f in filt.rglob(pattern):
                    try:
                        run(
                            [
                                "python3",
                                str(self.ap_uid_tput_plot_script),
                                "--input",
                                str(f),
                                "--outdir",
                                str(plot_dir),
                            ]
                        )
                        generated.append(plot_dir)
                    except Exception:
                        pass

            # 4) ModemActivityInfo plot
            modem_inputs = sorted(filt.rglob("modem_activity_info.csv"))
            if not modem_inputs:
                modem_inputs = sorted(filt.rglob("power_modem_activityinfo.txt"))
            for f in modem_inputs:
                try:
                    run(
                        [
                            "python3",
                            str(self.ap_modem_activity_plot_script),
                            "--input",
                            str(f),
                            "--outdir",
                            str(plot_dir),
                        ]
                    )
                    generated.append(plot_dir)
                except Exception:
                    pass

            # 5) NetworkBrain serving cell timeline.
            if self.ap_serving_cell_plot_script and self.ap_serving_cell_plot_script.exists():
                for f in filt.rglob("reg_data_sub_serv_cell.txt"):
                    try:
                        run(
                            [
                                "python3",
                                str(self.ap_serving_cell_plot_script),
                                "--input",
                                str(f),
                                "--outdir",
                                str(plot_dir),
                            ]
                        )
                        generated.append(plot_dir)
                    except Exception:
                        pass

            # 6) Fixed AP window overview template.
            overview_inputs = [
                filt / "modem_activity_info.csv",
                filt / "rf_sim1_signal_strength.txt",
                filt / "rf_sim2_signal_strength.txt",
                filt / "reg_sim1_data_reg_rat.txt",
                filt / "reg_sim2_data_reg_rat.txt",
                filt / "data_net_netifstats.txt",
                filt / "data_net_uidstats.txt",
            ]
            if self.ap_window_overview_plot_script and self.ap_window_overview_plot_script.exists():
                if any(path.exists() for path in overview_inputs):
                    try:
                        run(
                            [
                                "python3",
                                str(self.ap_window_overview_plot_script),
                                "--filter-dir",
                                str(filt),
                                "--outdir",
                                str(plot_dir),
                            ]
                        )
                        generated.append(plot_dir)
                    except Exception:
                        pass

        # dedupe dirs
        seen = set()
        uniq: List[Path] = []
        for d in generated:
            key = str(d.resolve())
            if key in seen:
                continue
            uniq.append(d)
            seen.add(key)
        return uniq

    def run(self, session_root: Path) -> None:
        self._run_aplog_plots(session_root)
        self._run_mdlog_plots(session_root)