import streamlit as st
import json
import os

# ---- Optional: define parser directly here or import from module ----
def parse_dcm_assembly(uploaded_file):
    lines = uploaded_file.read().decode("utf-8").splitlines()
    geometry = {}
    constraints = []
    parts = []
    set_geom_map = {}
    current = {}
    label = None

    for line in lines:
        line = line.strip()
        if line.startswith("CLASS "):
            if current:
                if current["type"] == "SET":
                    parts.append({
                        "name": current.get("label", ""),
                        "member_tags": current.get("member_tags", [])
                    })
                elif current["type"] == "PLANE":
                    geometry[current["id"]] = {
                        "base_point": current.get("base_point"),
                        "normal": current.get("normal"),
                        "label": current.get("label", "")
                    }
                elif current["type"] == "WORK_PLANE":
                    tag = current.get("tag")
                    if tag in geometry:
                        geometry[tag]["label"] = current.get("label", "")
                elif current["type"] == "CONSTRAINT":
                    constraints.append({
                        "from": current.get("end_1"),
                        "to": current.get("end_2"),
                        "label": current.get("label", "")
                    })
            parts_of_line = line.split()
            current = {
                "type": parts_of_line[1],
                "id": int(parts_of_line[2])
            }
        elif line.startswith("FBASE_POINT"):
            current["base_point"] = tuple(map(float, line.split()[1:]))
        elif line.startswith("FNORMAL"):
            current["normal"] = tuple(map(float, line.split()[1:]))
        elif line.startswith("FMY_PLANE TAG"):
            current["tag"] = int(line.split()[-1])
        elif line.startswith("DCM3_TEXT_LABEL"):
            current["label"] = line.split("DCM3_TEXT_LABEL", 1)[-1].strip()
        elif line.startswith("FMEMBER TAGS"):
            current["member_tags"] = list(map(int, line.split()[2:]))
        elif line.startswith("FEND_1 TAG"):
            current["end_1"] = int(line.split()[-1])
        elif line.startswith("FEND_2 TAG"):
            current["end_2"] = int(line.split()[-1])
    if current:
        if current["type"] == "SET":
            parts.append({
                "name": current.get("label", ""),
                "member_tags": current.get("member_tags", [])
            })
        elif current["type"] == "PLANE":
            geometry[current["id"]] = {
                "base_point": current.get("base_point"),
                "normal": current.get("normal"),
                "label": current.get("label", "")
            }
        elif current["type"] == "WORK_PLANE":
            tag = current.get("tag")
            if tag in geometry:
                geometry[tag]["label"] = current.get("label", "")
        elif current["type"] == "CONSTRAINT":
            constraints.append({
                "from": current.get("end_1"),
                "to": current.get("end_2"),
                "label": current.get("label", "")
            })

    return {
        "geometry": geometry,
        "constraints": constraints,
        "parts": parts
    }

# ---- Streamlit UI and file handling ----
st.set_page_config(layout="wide")
st.title("DCM Assembly Visualizer")

file1 = st.file_uploader("Upload File A", type=["txt"])
file2 = st.file_uploader("Upload File B", type=["txt"])

parsed_files = []

if file1:
    parsed_files.append(("red", "blue", parse_dcm_assembly(file1)))
if file2:
    parsed_files.append(("purple", "orange", parse_dcm_assembly(file2)))

if parsed_files:
    os.makedirs("static", exist_ok=True)
    with open("static/parsed_data.json", "w") as f:
        json.dump({"files": parsed_files}, f, indent=2)

    st.components.v1.html(
        '<iframe src="static/index.html" width="100%" height="800" style="border:none;"></iframe>',
        height=800,
    )