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

try:
    from flask import Flask, request, render_template_string
except Exception:  # Flask is optional for CLI usage
    Flask = None  # type: ignore
    request = None  # type: ignore
    render_template_string = None  # type: ignore


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"], default="csv", help="Output format.")
    parser.add_argument("--out", "-o", default=None, help="Output file path. Defaults to estimates.csv/json.")
    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)
    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 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="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 ensure_flask():
    if Flask is None:
        print("Flask is not installed. Install it with: pip install Flask", file=sys.stderr)
        sys.exit(1)


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:
    ensure_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"
                else:
                    payload = to_csv_string(estimates_rows)
                    mime = "text/csv"
                    download_name = "estimates.csv"
                if provide_download:
                    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()