"""CTF submission bundle builder.

Generates all required and optional submission artifacts from a model script:

    ctf_submission/
    ├── model.py          Required — cleaned script (strips local __main__ block)
    ├── weights.csv       Required — id, eom, w output
    ├── requirements.txt  Required if non-standard deps — auto-detected
    └── methodology.pdf   Optional — generated from Markdown string input

Usage::

    bundle = build_submission(
        script_path="reference/benchmark-ipca-pf.py",
        weights_df=weights,
        output_dir="ctf_submission/",
        methodology_md="# Method\\nWe use IPCA...",
    )
    print(bundle.script_path)
    print(bundle.weights_path)
"""

from __future__ import annotations

import ast
import io
import os
import shutil
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path

import pandas as pd

from eapctf.ctf.interface import CTF_PREINSTALLED

# ---------------------------------------------------------------------------
# Data class
# ---------------------------------------------------------------------------


@dataclass
class SubmissionBundle:
    """Paths to all generated submission artifacts.

    Attributes
    ----------
    script_path : Path
        Cleaned ``model.py`` (local ``__main__`` block stripped).
    weights_path : Path
        ``weights.csv`` with columns id, eom, w.
    requirements_path : Path | None
        ``requirements.txt`` if non-pre-installed deps were detected, else None.
    pdf_path : Path | None
        ``methodology.pdf`` if methodology_md was provided, else None.
    """

    script_path: Path
    weights_path: Path
    requirements_path: Path | None
    pdf_path: Path | None

    def __str__(self) -> str:
        lines = [
            "SubmissionBundle",
            f"  model.py        : {self.script_path}",
            f"  weights.csv     : {self.weights_path}",
            f"  requirements    : {self.requirements_path or '(none — all pre-installed)'}",
            f"  methodology.pdf : {self.pdf_path or '(not generated)'}",
        ]
        return "\n".join(lines)


# ---------------------------------------------------------------------------
# Script cleaning
# ---------------------------------------------------------------------------


def _strip_main_block(source: str) -> str:
    """Remove ``if __name__ == '__main__':`` block and ``if False:`` blocks.

    The CTF runner calls ``main()`` directly; local test harnesses should be
    stripped to avoid hardcoded paths that don't exist in the CTF environment.
    """
    try:
        tree = ast.parse(source)
    except SyntaxError:
        return source  # Return unchanged if unparseable

    lines = source.splitlines(keepends=True)
    # Identify line ranges to remove (1-indexed, matching ast node lineno)
    remove_ranges: list[tuple[int, int]] = []

    for node in tree.body:
        if isinstance(node, ast.If):
            test = node.test
            # Pattern: if __name__ == "__main__":
            is_main = (
                isinstance(test, ast.Compare)
                and isinstance(test.left, ast.Name)
                and test.left.id == "__name__"
                and len(test.ops) == 1
                and isinstance(test.ops[0], ast.Eq)
                and len(test.comparators) == 1
                and isinstance(test.comparators[0], ast.Constant)
                and test.comparators[0].value == "__main__"
            )
            # Pattern: if False:
            is_false = isinstance(test, ast.Constant) and test.value is False

            if is_main or is_false:
                start = node.lineno - 1  # 0-indexed
                end = node.end_lineno    # exclusive upper bound (0-indexed end+1)
                remove_ranges.append((start, end))  # type: ignore[arg-type]

    if not remove_ranges:
        return source

    # Build output, skipping removed line ranges
    result: list[str] = []
    remove_set: set[int] = set()
    for start, end in remove_ranges:
        remove_set.update(range(start, end))

    for i, line in enumerate(lines):
        if i not in remove_set:
            result.append(line)

    # Strip trailing blank lines added by removal
    cleaned = "".join(result).rstrip() + "\n"
    return cleaned


# ---------------------------------------------------------------------------
# Dependency detection
# ---------------------------------------------------------------------------


def _detect_non_preinstalled_deps(source: str) -> list[str]:
    """Return sorted list of non-pre-installed top-level imports in source."""
    try:
        tree = ast.parse(source)
    except SyntaxError:
        return []

    imports: set[str] = set()
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                imports.add(alias.name.split(".")[0])
        elif isinstance(node, ast.ImportFrom):
            if node.module:
                imports.add(node.module.split(".")[0])

    non_pre = imports - CTF_PREINSTALLED
    non_pre.discard("")
    return sorted(non_pre)


# ---------------------------------------------------------------------------
# PDF generation
# ---------------------------------------------------------------------------


def _generate_pdf(markdown_text: str, output_path: Path) -> bool:
    """Attempt to generate a PDF from Markdown text.

    Tries (in order):
    1. pandoc (if installed) — markdown → PDF via LaTeX
    2. weasyprint (if importable) — markdown → HTML → PDF
    3. Falls back to writing .md only; returns False.

    Returns True if PDF was generated, False otherwise.
    """
    # Strategy 1: pandoc
    if shutil.which("pandoc"):
        with _temp_md_file(markdown_text) as md_path:
            result = subprocess.run(
                ["pandoc", str(md_path), "-o", str(output_path),
                 "--pdf-engine=pdflatex"],
                capture_output=True,
            )
        if result.returncode == 0 and output_path.exists():
            return True

    # Strategy 2: weasyprint
    try:
        import markdown as md_lib
        import weasyprint

        html_body = md_lib.markdown(markdown_text)
        html = f"<html><body>{html_body}</body></html>"
        weasyprint.HTML(string=html).write_pdf(str(output_path))
        if output_path.exists():
            return True
    except (ImportError, Exception):
        pass

    return False


class _temp_md_file:
    """Context manager for a temporary Markdown file."""

    def __init__(self, content: str) -> None:
        self._content = content
        self._path: Path | None = None

    def __enter__(self) -> Path:
        fd, name = tempfile.mkstemp(suffix=".md")
        self._path = Path(name)
        self._path.write_text(self._content, encoding="utf-8")
        os.close(fd)
        return self._path

    def __exit__(self, *args: object) -> None:
        if self._path and self._path.exists():
            self._path.unlink()


# ---------------------------------------------------------------------------
# Bundle verification
# ---------------------------------------------------------------------------


def _verify_bundle(bundle: SubmissionBundle) -> list[str]:
    """Return list of warning strings for any bundle constraint violations."""
    warnings: list[str] = []

    # model.py
    if not bundle.script_path.exists():
        warnings.append("model.py not found")
    else:
        size_kb = bundle.script_path.stat().st_size / 1024
        if size_kb > 1024:
            warnings.append(f"model.py size {size_kb:.0f} KB exceeds 1 MB limit")
        try:
            with bundle.script_path.open("rb") as fh:
                fh.read().decode("utf-8")
        except UnicodeDecodeError:
            warnings.append("model.py is not valid UTF-8")

    # weights.csv
    if not bundle.weights_path.exists():
        warnings.append("weights.csv not found")
    else:
        size_mb = bundle.weights_path.stat().st_size / (1024 ** 2)
        if size_mb > 50:
            warnings.append(f"weights.csv size {size_mb:.1f} MB exceeds 50 MB limit")

    return warnings


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------


def build_submission(
    script_path: str | Path,
    weights_df: pd.DataFrame,
    output_dir: str | Path = "ctf_submission",
    methodology_md: str | None = None,
    auto_detect_deps: bool = True,
) -> SubmissionBundle:
    """Build a complete CTF submission bundle.

    Parameters
    ----------
    script_path : str | Path
        Path to the model script (with CTF ``main()``).
    weights_df : pd.DataFrame
        Portfolio weights DataFrame with columns id, eom, w.
    output_dir : str | Path
        Directory to write all bundle artifacts.  Created if needed.
    methodology_md : str | None
        Markdown text for the methodology document.  If provided, generates
        ``methodology.pdf`` (via pandoc or weasyprint) or falls back to
        ``methodology.md`` with a warning.
    auto_detect_deps : bool
        If True, scan script for non-pre-installed imports and write
        ``requirements.txt`` if any are found.

    Returns
    -------
    SubmissionBundle
        Paths to all generated files.

    Raises
    ------
    FileNotFoundError
        If script_path does not exist.
    ValueError
        If weights_df is missing required columns.
    """
    script_path = Path(script_path).resolve()
    output_dir = Path(output_dir).resolve()

    if not script_path.exists():
        raise FileNotFoundError(f"Script not found: {script_path}")

    for col in ("id", "eom", "w"):
        if col not in weights_df.columns:
            raise ValueError(f"weights_df missing column '{col}'")

    output_dir.mkdir(parents=True, exist_ok=True)

    # --- Step 1: Clean and copy model.py ---
    source = script_path.read_text(encoding="utf-8")
    cleaned = _strip_main_block(source)
    model_out = output_dir / "model.py"
    model_out.write_text(cleaned, encoding="utf-8")

    # --- Step 2: Write weights.csv ---
    weights_out = output_dir / "weights.csv"
    weights_export = weights_df[["id", "eom", "w"]].copy()

    # Check size — if > 40 MB, filter to ctff_test period rows by proxy
    # (we can't know ctff_test here, so we warn and write all dates)
    buf = io.StringIO()
    weights_export.to_csv(buf, index=False)
    size_mb = buf.tell() / (1024 ** 2)

    if size_mb > 40:
        print(
            f"[WARNING] weights.csv estimated size {size_mb:.1f} MB > 40 MB. "
            f"Consider filtering to ctff_test=True rows before building the bundle "
            f"if the file exceeds 50 MB."
        )

    weights_export.to_csv(weights_out, index=False)

    # --- Step 3: Auto-detect and write requirements.txt ---
    requirements_out: Path | None = None
    if auto_detect_deps:
        deps = _detect_non_preinstalled_deps(cleaned)
        if deps:
            req_content = "\n".join(deps) + "\n"
            requirements_out = output_dir / "requirements.txt"
            requirements_out.write_text(req_content, encoding="utf-8")

    # --- Step 4: Generate methodology PDF ---
    pdf_out: Path | None = None
    if methodology_md is not None:
        pdf_path = output_dir / "methodology.pdf"
        success = _generate_pdf(methodology_md, pdf_path)
        if success:
            pdf_out = pdf_path
        else:
            # Fallback: write .md and warn
            md_path = output_dir / "methodology.md"
            md_path.write_text(methodology_md, encoding="utf-8")
            print(
                "[WARNING] PDF generation failed (pandoc and weasyprint not available). "
                f"Wrote {md_path} instead. Convert to PDF manually before submission."
            )

    bundle = SubmissionBundle(
        script_path=model_out,
        weights_path=weights_out,
        requirements_path=requirements_out,
        pdf_path=pdf_out,
    )

    # --- Step 5: Verify bundle ---
    issues = _verify_bundle(bundle)
    for issue in issues:
        print(f"[WARNING] Bundle verification: {issue}")

    return bundle
