import argparse
import base64
import csv
import io
import json
import math
import os
import struct
import sys
import tempfile
import types
import zipfile
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Iterable, List, Tuple, Dict, Any, Optional

# Some bundled tools (e.g., MySQL Workbench) inject their own Python stdlib zip into sys.path.
# That can ship an incompatible ctypes build and crash Flask/colorama imports. Strip it out early.
sys.path = [p for p in sys.path if "MySQL Workbench" not in p]


Vec3 = Tuple[float, float, float]


@dataclass
class PrinterSettings:
    layer_height_mm: float = 0.2
    nozzle_diameter_mm: float = 0.4
    perimeter_count: int = 2
    top_layers: int = 4
    bottom_layers: int = 4
    infill_density: float = 0.2  # 0.0 - 1.0
    perimeter_speed_mm_s: float = 40.0
    infill_speed_mm_s: float = 60.0
    travel_speed_mm_s: float = 120.0
    filament_diameter_mm: float = 1.75
    filament_density_g_cm3: float = 1.24  # PLA-ish
    travel_factor: float = 0.1  # additional travel distance as a fraction of extrusion distance


@dataclass
class StlStats:
    volume_mm3: float
    surface_area_mm2: float
    bounds_min: Vec3
    bounds_max: Vec3


@dataclass
class Estimate:
    file: str
    volume_mm3: float
    surface_area_mm2: float
    bounding_box_x_mm: float
    bounding_box_y_mm: float
    bounding_box_z_mm: float
    filament_length_mm: float
    filament_mass_g: float
    estimated_print_time_min: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Estimate filament usage and print time for STL files.")
    parser.add_argument("--config", "-c", default="printer-config.json", help="Path to printer configuration JSON.")
    parser.add_argument("--output-format", "-f", choices=["csv", "json", "xlsx"], default="csv", help="Output format.")
    parser.add_argument("--out", "-o", default=None, help="Output file path. Defaults to estimates.csv/json/xlsx.")
    parser.add_argument("--serve", action="store_true", help="Run web interface instead of CLI batch mode.")
    parser.add_argument("--port", type=int, default=5000, help="Port for the web interface.")
    return parser.parse_args()


def load_config(path: Path) -> PrinterSettings:
    if not path.exists():
        return PrinterSettings()
    with path.open("r", encoding="utf-8") as f:
        raw = json.load(f)
    return PrinterSettings(**raw)


def is_binary_stl(data: bytes) -> bool:
    if len(data) < 84:
        return True
    header = data[:80]
    tri_count = struct.unpack("<I", data[80:84])[0]
    expected_len = 84 + tri_count * 50
    if expected_len == len(data):
        return True
    # ASCII files start with "solid" but some binaries do too; check for non-text bytes
    if header[:5].lower() != b"solid":
        return True
    if b"\0" in header:
        return True
    return False


def accumulate_triangle_stats(v1: Vec3, v2: Vec3, v3: Vec3) -> Tuple[float, float]:
    # Signed volume of tetrahedron (origin, v1, v2, v3)
    cross_x = v2[1] * v3[2] - v2[2] * v3[1]
    cross_y = v2[2] * v3[0] - v2[0] * v3[2]
    cross_z = v2[0] * v3[1] - v2[1] * v3[0]
    signed_volume = (v1[0] * cross_x + v1[1] * cross_y + v1[2] * cross_z) / 6.0

    # Triangle area
    ab = (v2[0] - v1[0], v2[1] - v1[1], v2[2] - v1[2])
    ac = (v3[0] - v1[0], v3[1] - v1[1], v3[2] - v1[2])
    cross_area_x = ab[1] * ac[2] - ab[2] * ac[1]
    cross_area_y = ab[2] * ac[0] - ab[0] * ac[2]
    cross_area_z = ab[0] * ac[1] - ab[1] * ac[0]
    area = 0.5 * math.sqrt(
        cross_area_x * cross_area_x + cross_area_y * cross_area_y + cross_area_z * cross_area_z
    )
    return signed_volume, area


def read_stl(path: Path) -> StlStats:
    data = path.read_bytes()
    volume = 0.0
    area = 0.0
    min_pt = [math.inf, math.inf, math.inf]
    max_pt = [-math.inf, -math.inf, -math.inf]

    def update_bounds(pt: Vec3) -> None:
        min_pt[0] = min(min_pt[0], pt[0])
        min_pt[1] = min(min_pt[1], pt[1])
        min_pt[2] = min(min_pt[2], pt[2])
        max_pt[0] = max(max_pt[0], pt[0])
        max_pt[1] = max(max_pt[1], pt[1])
        max_pt[2] = max(max_pt[2], pt[2])

    if is_binary_stl(data):
        tri_count = struct.unpack("<I", data[80:84])[0] if len(data) >= 84 else 0
        offset = 84
        for _ in range(tri_count):
            chunk = data[offset : offset + 50]
            if len(chunk) < 50:
                break
            coords = struct.unpack("<12fH", chunk)
            v1 = (coords[3], coords[4], coords[5])
            v2 = (coords[6], coords[7], coords[8])
            v3 = (coords[9], coords[10], coords[11])
            update_bounds(v1)
            update_bounds(v2)
            update_bounds(v3)
            v, a = accumulate_triangle_stats(v1, v2, v3)
            volume += v
            area += a
            offset += 50
    else:
        lines = data.decode("utf-8", errors="ignore").splitlines()
        verts: List[Vec3] = []
        for line in lines:
            parts = line.strip().split()
            if len(parts) >= 4 and parts[0].lower() == "vertex":
                try:
                    v = (float(parts[1]), float(parts[2]), float(parts[3]))
                except ValueError:
                    continue
                verts.append(v)
                if len(verts) == 3:
                    v1, v2, v3 = verts
                    update_bounds(v1)
                    update_bounds(v2)
                    update_bounds(v3)
                    v, a = accumulate_triangle_stats(v1, v2, v3)
                    volume += v
                    area += a
                    verts = []
    if math.isinf(min_pt[0]):
        min_pt = [0.0, 0.0, 0.0]
        max_pt = [0.0, 0.0, 0.0]
    return StlStats(abs(volume), area, (min_pt[0], min_pt[1], min_pt[2]), (max_pt[0], max_pt[1], max_pt[2]))


def estimate(stats: StlStats, settings: PrinterSettings, filename: str) -> Estimate:
    line_area = settings.nozzle_diameter_mm * settings.layer_height_mm
    shell_thickness = settings.perimeter_count * settings.nozzle_diameter_mm

    # Shell volume approximated from surface area (rough but better than box volume)
    shell_volume = stats.surface_area_mm2 * shell_thickness

    # Infill volume limited to remaining interior
    internal_volume = max(stats.volume_mm3 - shell_volume, 0.0)
    infill_volume = internal_volume * settings.infill_density
    total_extrusion_volume = shell_volume + infill_volume

    # Filament calculations
    filament_cross_section = math.pi * (settings.filament_diameter_mm / 2.0) ** 2
    filament_length = total_extrusion_volume / filament_cross_section if filament_cross_section else 0.0
    filament_mass = (total_extrusion_volume / 1000.0) * settings.filament_density_g_cm3

    # Print time estimate
    perimeter_length = shell_volume / line_area if line_area else 0.0
    infill_length = infill_volume / line_area if line_area else 0.0
    travel_length = (perimeter_length + infill_length) * settings.travel_factor
    time_s = (
        perimeter_length / settings.perimeter_speed_mm_s
        + infill_length / settings.infill_speed_mm_s
        + travel_length / settings.travel_speed_mm_s
    )

    bbox = (
        stats.bounds_max[0] - stats.bounds_min[0],
        stats.bounds_max[1] - stats.bounds_min[1],
        stats.bounds_max[2] - stats.bounds_min[2],
    )
    return Estimate(
        file=filename,
        volume_mm3=round(stats.volume_mm3, 3),
        surface_area_mm2=round(stats.surface_area_mm2, 3),
        bounding_box_x_mm=round(bbox[0], 3),
        bounding_box_y_mm=round(bbox[1], 3),
        bounding_box_z_mm=round(bbox[2], 3),
        filament_length_mm=round(filament_length, 3),
        filament_mass_g=round(filament_mass, 3),
        estimated_print_time_min=round(time_s / 60.0, 2),
    )


def find_stl_files(folder: Path) -> List[Path]:
    return sorted([p for p in folder.iterdir() if p.is_file() and p.suffix.lower() == ".stl"])


def write_output(estimates: List[Estimate], fmt: str, out_path: Path) -> None:
    if fmt == "json":
        with out_path.open("w", encoding="utf-8") as f:
            json.dump([asdict(e) for e in estimates], f, indent=2)
    elif fmt == "xlsx":
        write_xlsx(estimates, out_path)
    else:
        with out_path.open("w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=list(asdict(estimates[0]).keys()))
            writer.writeheader()
            for e in estimates:
                writer.writerow(asdict(e))


def minutes_to_hours_formula(expr: str) -> str:
    return f'INT({expr}/60) & " hr" & IF(ROUND(MOD({expr},60),0)>0, " " & ROUND(MOD({expr},60),0) & " min","")'


def minutes_to_hours_formula_text(minutes: float) -> str:
    hours = int(minutes // 60)
    mins = int(round(minutes % 60))
    if mins > 0:
        return f"{hours} hr {mins} min" if hours else f"{mins} min"
    return f"{hours} hr"


def require_openpyxl():
    block_broken_numpy()
    try:
        import openpyxl  # type: ignore
    except Exception:
        print("XLSX output requires openpyxl. Install it with: pip install openpyxl", file=sys.stderr)
        sys.exit(1)
    return openpyxl


def block_broken_numpy():
    # Some environments ship a corrupted/bundled numpy (e.g., via other apps) that can crash ctypes.
    # openpyxl only needs numpy optionally; force it to skip by blocking the import.
    class _Blocker:
        def find_spec(self, fullname, path=None, target=None):
            if fullname.startswith("numpy"):
                raise ImportError("Blocked numpy to avoid crashing openpyxl")
            return None

    sys.meta_path.insert(0, _Blocker())
    for key in list(sys.modules.keys()):
        if key.startswith("numpy"):
            sys.modules.pop(key, None)


def block_colorama():
    # Avoid loading colorama (ctypes-dependent) which crashes in this environment.
    dummy = types.SimpleNamespace(
        Fore=types.SimpleNamespace(RESET=""),
        Style=types.SimpleNamespace(RESET_ALL=""),
        Back=types.SimpleNamespace(RESET=""),
    )

    def _init(*args, **kwargs):
        return None

    class _Ansi:
        def __init__(self, stream=None, convert=None, strip=None, autoreset=None):
            self.stream = stream

        def write(self, text):
            if self.stream:
                return self.stream.write(text)

        def flush(self):
            if self.stream:
                return self.stream.flush()

    mod = types.SimpleNamespace(
        init=_init,
        Fore=dummy.Fore,
        Back=dummy.Back,
        Style=dummy.Style,
        AnsiToWin32=_Ansi,
    )

    for name in ["colorama", "colorama.initialise", "colorama.ansitowin32", "colorama.winterm"]:
        sys.modules[name] = mod


def block_click_winconsole():
    # Avoid click's _winconsole ctypes usage on this system.
    def _get_windows_console_stream(stream, encoding=None, errors=None):
        return stream

    def get_best_encoding(stream):
        return "utf-8"

    def _default_text_stdout():  # noqa: N802
        return sys.stdout

    def _default_text_stdin():  # noqa: N802
        return sys.stdin

    def _default_text_stderr():  # noqa: N802
        return sys.stderr

    dummy = types.SimpleNamespace(
        _get_windows_console_stream=_get_windows_console_stream,
        get_best_encoding=get_best_encoding,
        _default_text_stdout=_default_text_stdout,
        _default_text_stdin=_default_text_stdin,
        _default_text_stderr=_default_text_stderr,
    )
    sys.modules["click._winconsole"] = dummy


def build_workbook(estimates: List[Estimate]):
    openpyxl = require_openpyxl()
    wb = openpyxl.Workbook()
    ws = wb.active
    ws.title = "Estimates"
    headers = [
        "File",
        "Filament Length (mm)",
        "Filament Mass (g)",
        "Estimated Print Time (min)",
        "Estimated Print Time (hrs)",
        "Estimated Print Time (hrs text)",
    ]
    ws.append(headers)
    ws.cell(row=1, column=7, value="Totals")

    start_row = 2
    for idx, est in enumerate(estimates, start=start_row):
        ws.cell(row=idx, column=1, value=est.file)
        ws.cell(row=idx, column=2, value=est.filament_length_mm)
        ws.cell(row=idx, column=3, value=est.filament_mass_g)
        ws.cell(row=idx, column=4, value=est.estimated_print_time_min)
        time_formula = minutes_to_hours_formula(f"D{idx}")
        ws.cell(row=idx, column=5, value=f"={time_formula}")
        ws.cell(row=idx, column=6, value=minutes_to_hours_formula_text(est.estimated_print_time_min))

    end_row = start_row + len(estimates) - 1
    if estimates:
        ws.cell(row=start_row, column=7, value="Filament (g)")
        ws.cell(row=start_row, column=8, value=f"=SUM(C{start_row}:C{end_row})")
        ws.cell(row=start_row + 1, column=7, value="Time")
        total_minutes_expr = f"SUM(D{start_row}:D{end_row})"
        total_time_formula = minutes_to_hours_formula(total_minutes_expr)
        ws.cell(row=start_row + 1, column=8, value=f"={total_time_formula}")
        ws.cell(row=start_row + 1, column=9, value=minutes_to_hours_formula_text(sum(e.estimated_print_time_min for e in estimates)))

    # Force Excel/Sheets to recalc on open so formulas show values immediately
    from openpyxl.workbook.properties import CalcProperties  # type: ignore
    wb.calculation_properties = CalcProperties(calcMode="auto", fullCalcOnLoad=True)
    return wb


def write_xlsx(estimates: List[Estimate], out_path: Path) -> None:
    wb = build_workbook(estimates)
    wb.save(out_path)


def xlsx_bytes(estimates: List[Estimate]) -> bytes:
    wb = build_workbook(estimates)
    buf = io.BytesIO()
    wb.save(buf)
    return buf.getvalue()


def batch_cli(settings: PrinterSettings, fmt: str, out_path: Optional[Path]) -> None:
    folder = Path.cwd()
    stl_files = find_stl_files(folder)
    if not stl_files:
        print("No STL files found in current directory.", file=sys.stderr)
        sys.exit(1)
    estimates: List[Estimate] = []
    for stl_path in stl_files:
        stats = read_stl(stl_path)
        estimates.append(estimate(stats, settings, stl_path.name))

    if out_path is None:
        out_path = folder / f"estimates.{fmt}"
    write_output(estimates, fmt, out_path)
    print(f"Wrote {len(estimates)} records to {out_path}")


WEB_TEMPLATE = """
<!doctype html>
<html>
<head>
  <title>STL Print Estimator</title>
  <style>
    body { font-family: Arial, sans-serif; margin: 2rem; }
    form { margin-bottom: 1.5rem; }
    table { border-collapse: collapse; width: 100%; }
    th, td { border: 1px solid #ccc; padding: 0.4rem; text-align: left; }
    th { background: #f2f2f2; }
    .download { margin: 1rem 0; }
  </style>
</head>
<body>
  <h1>STL Print Estimator</h1>
  <form method="post" enctype="multipart/form-data">
    <label>Upload STL or ZIP of STLs: <input type="file" name="file" required></label><br><br>
    <label><input type="radio" name="format" value="csv" checked> CSV</label>
    <label><input type="radio" name="format" value="json"> JSON</label>
    <label><input type="radio" name="format" value="xlsx"> XLSX</label>
    <label><input type="checkbox" name="download"> Provide download</label><br><br>
    <button type="submit">Analyze</button>
  </form>
  {% if estimates %}
  <div class="download">
    {% if download_link %}
      <a href="{{ download_link }}" download="{{ download_name }}">Download {{ download_name }}</a>
    {% endif %}
  </div>
  <table>
    <thead>
      <tr>
        {% for key in headers %}
        <th>{{ key }}</th>
        {% endfor %}
      </tr>
    </thead>
    <tbody>
      {% for row in estimates %}
        <tr>
        {% for key in headers %}
          <td>{{ row[key] }}</td>
        {% endfor %}
        </tr>
      {% endfor %}
    </tbody>
  </table>
  {% endif %}
</body>
</html>
"""


def to_csv_string(rows: List[Dict[str, Any]]) -> str:
    buf = io.StringIO()
    writer = csv.DictWriter(buf, fieldnames=list(rows[0].keys()))
    writer.writeheader()
    for r in rows:
        writer.writerow(r)
    return buf.getvalue()


def load_flask():
    block_colorama()
    block_click_winconsole()
    try:
        from flask import Flask, request, render_template_string
    except Exception:
        print("Flask is not installed. Install it with: pip install Flask", file=sys.stderr)
        sys.exit(1)
    return Flask, request, render_template_string


def handle_uploaded_file(file_storage, tmpdir: Path) -> List[Path]:
    filename = file_storage.filename or "upload"
    dest = tmpdir / filename
    file_storage.save(dest)
    paths: List[Path] = []
    if dest.suffix.lower() == ".zip":
        with zipfile.ZipFile(dest, "r") as zf:
            zf.extractall(tmpdir)
        for p in tmpdir.rglob("*.stl"):
            paths.append(p)
    elif dest.suffix.lower() == ".stl":
        paths.append(dest)
    return paths


def run_server(settings: PrinterSettings, fmt: str, port: int) -> None:
    Flask, request, render_template_string = load_flask()
    app = Flask(__name__)

    @app.route("/", methods=["GET", "POST"])
    def index():
        estimates_rows: List[Dict[str, Any]] = []
        download_link = None
        download_name = None
        chosen_format = fmt
        if request.method == "POST":
            chosen_format = request.form.get("format", fmt)
            provide_download = request.form.get("download") is not None
            upload = request.files.get("file")
            if not upload:
                return render_template_string(WEB_TEMPLATE, estimates=None, headers=[], download_link=None)
            with tempfile.TemporaryDirectory() as td:
                tmpdir = Path(td)
                stl_paths = handle_uploaded_file(upload, tmpdir)
                for p in stl_paths:
                    stats = read_stl(p)
                    est = estimate(stats, settings, p.name)
                    estimates_rows.append(asdict(est))
            if estimates_rows:
                if chosen_format == "json":
                    payload = json.dumps(estimates_rows, indent=2)
                    mime = "application/json"
                    download_name = "estimates.json"
                elif chosen_format == "xlsx":
                    estimate_objs = [Estimate(**row) for row in estimates_rows]
                    payload_bytes = xlsx_bytes(estimate_objs)
                    payload = base64.b64encode(payload_bytes).decode("ascii")
                    mime = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                    download_name = "estimates.xlsx"
                else:
                    payload = to_csv_string(estimates_rows)
                    mime = "text/csv"
                    download_name = "estimates.csv"
                if provide_download:
                    if chosen_format == "xlsx":
                        download_link = f"data:{mime};base64,{payload}"
                    else:
                        download_link = f"data:{mime};charset=utf-8,{payload.replace('%','%25').replace('\\n','%0A').replace(' ','%20')}"
        headers = list(estimates_rows[0].keys()) if estimates_rows else []
        return render_template_string(
            WEB_TEMPLATE,
            estimates=estimates_rows or None,
            headers=headers,
            download_link=download_link,
            download_name=download_name,
        )

    app.run(host="0.0.0.0", port=port, debug=False)


def main() -> None:
    args = parse_args()
    settings = load_config(Path(args.config))
    if args.serve:
        run_server(settings, args.output_format, args.port)
    else:
        batch_cli(settings, args.output_format, Path(args.out) if args.out else None)


if __name__ == "__main__":
    main()