package ua.be.wee.controller;

import java.io.FileNotFoundException;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Stack;
import java.util.stream.Collectors;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.Part;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.env.Environment;
import org.springframework.stereotype.Controller;
import org.springframework.ui.Model;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.ModelAttribute;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestMethod;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.ResponseBody;

import ua.be.wee.dto.EventView;
import ua.be.wee.model.EnactmentController;
import ua.be.wee.model.Token;
import ua.be.wee.model.nodes.Activity;
import ua.be.wee.model.nodes.Artifact;
import ua.be.wee.model.nodes.AutomatedActivity;
import ua.be.wee.model.nodes.AutomatedStatus;
import ua.be.wee.model.nodes.FinalNode;
import ua.be.wee.model.nodes.ForkJoinNode;
import ua.be.wee.model.nodes.InvokingActivity;
import ua.be.wee.model.nodes.Node;
import ua.be.wee.model.nodes.ports.ControlInputPort;
import ua.be.wee.model.nodes.ports.ControlOutputPort;
import ua.be.wee.model.pm.PM;
import ua.be.wee.model.pt.EndActivityEvent;
import ua.be.wee.model.pt.EndTraceEvent;
import ua.be.wee.model.pt.Event;
import ua.be.wee.model.pt.PT;
import ua.be.wee.model.pt.StartActivityEvent;
import ua.be.wee.model.pt.StartTraceEvent;
import ua.be.wee.model.pt.TraceArtifact;
import ua.be.wee.model.repository.FusekiWrapper;
import ua.be.wee.model.util.Pair;

@Controller
@RequestMapping("/")
public class EnactmentControllerMVC {

    @Autowired
    private EnactmentController controller;

    @Autowired
    private Environment env;

    @Value("${base_url}")
    private String baseUrl;

    @ModelAttribute("baseUrl")
    public String getUrl() {
        return this.baseUrl;
    }

    private Map<String, Stack<String>> branchStackMap = new HashMap<>();
    private Map<String, String> invokingActivityMap = new HashMap<>();
    private Map<String, Pair<List<PMTrigger>, List<Map<String, Object>>>> currentActs = new HashMap<>();

    @RequestMapping("/")
    public String index(Model model) throws FileNotFoundException, IOException {
        if (!env.getProperty("endpoint").equals("")) {
            model.addAttribute("savedEndpoint", env.getProperty("endpoint"));
            model.addAttribute("hasEndpoint", true);
        } else {
            model.addAttribute("hasEndpoint", false);
        }
        return "index";
    }

    @PostMapping("/pm")
    public String getPMs(Model model, @RequestParam String endpoint) throws FileNotFoundException, IOException {
        FusekiWrapper.getInstance().setServiceURI(endpoint);
        if (FusekiWrapper.getInstance().testEndpoint()) {
            model.addAttribute("error", false);
            model.addAttribute("pms", controller.getAllPMs());
            model.addAttribute("traces", null);
            model.addAttribute("current", "1");
            return "pms";
        } else {
            model.addAttribute("error", true);
            return "index";
        }
    }
   

    @RequestMapping("/gettraces")
    public String getTraces(Model model, @RequestParam String iri, HttpServletRequest request) throws Exception {
        List<StartTraceEvent> activeTraces = controller.getActiveTraces(iri);
        if (activeTraces != null && activeTraces.size() > 0) {
            model.addAttribute("traces", activeTraces);
        } else {
            model.addAttribute("traces", null);
        }
        model.addAttribute("error", false);
        model.addAttribute("pms", controller.getAllPMs());
        model.addAttribute("current", iri);
        return "pms";
    }

    @PostMapping(value = "/enactpm", params = "action=new")
    public String getPMtoEnact(Model model, @RequestParam String pmiri, HttpServletRequest request) throws Exception {
        if (pmiri.equals("select PM")) {
            model.addAttribute("error", true);
            model.addAttribute("pms", controller.getAllPMs());
            model.addAttribute("traces", null);
            model.addAttribute("current", "1");
            return "pms";
        }

        PM pm = controller.getPM(pmiri);
        PT trace = controller.createTrace(pm);

        String initialBranchId = controller.generateBranchId(trace.getIri(), null);
        List<Pair<String, String>> iris = controller.findNextNodes(pm.getInitial().getIri(), trace.getIri(), controller.getExecutionFlowManager());
        List<PMTrigger> acts = findElements(pm, iris, initialBranchId, null);

        Map<String, PM> pmMap = new HashMap<>();
        branchStackMap.put(trace.getIri(), new Stack<>());
        branchStackMap.get(trace.getIri()).push(initialBranchId);
        pmMap.put(initialBranchId, pm);

        invokingActivityMap = new HashMap<>();

        request.getSession().setAttribute("pmMap", pmMap);
        request.getSession().setAttribute("invokingActivityMap", invokingActivityMap);
        request.getSession().setAttribute("pm", pm);
        request.getSession().setAttribute("trace", trace);
        request.getSession().setAttribute("acts", acts);
        request.getSession().setAttribute("endacts", new ArrayList<>());
        request.getSession().setAttribute("automated", controller.getAutomatedActivities(trace));

        List<Map<String, String>> activeBranches = new ArrayList<>();
        Map<String, String> mainBranch = new HashMap<>();
        mainBranch.put("branchId", initialBranchId);
        mainBranch.put("pmName", pm.getName());
        mainBranch.put("invokingActivityName", null);
        activeBranches.add(mainBranch);
        model.addAttribute("activeBranches", activeBranches);

        request.getSession().removeAttribute("currentBranchId");
        model.addAttribute("error", false);
        model.addAttribute("arts", null);
        model.addAttribute("current", "1");
        model.addAttribute("invokingActivityMap", invokingActivityMap);
        model.addAttribute("eventViews", buildEventViewList(trace));
        updateActiveBranches(trace, pmMap, invokingActivityMap, request, model);
        
        return "enact";
    }


    @RequestMapping(value = "/enactpm", method = RequestMethod.POST, params = "action=continue")
    public String continueEnact(Model model, @RequestParam String pmiri, @RequestParam String contTrace,
            HttpServletRequest request) throws Exception {
        if (pmiri.equals("select PM") || contTrace.equals("Select Trace")) {
            model.addAttribute("error", true);
            model.addAttribute("pms", controller.getAllPMs());
            model.addAttribute("traces", null);
            model.addAttribute("current", "1");
            return "pms";
        }

        PM pm = controller.getPM(pmiri);
        Stack<String> branchStack = branchStackMap.getOrDefault(contTrace, new Stack<>());

        String branchId = branchStack.isEmpty()
                ? controller.generateBranchId(contTrace, null)
                : branchStack.peek();

        List<Event> events = controller.getEvents(contTrace, branchId);

        PT trace = new PT();
        trace.setEvents(events);
        trace.setIri(contTrace);
        trace.setPmEnacted(pm);

        Map<String, PM> pmMap = (Map<String, PM>) request.getSession().getAttribute("pmMap");
        if (pmMap == null) {
            pmMap = new HashMap<>();
            pmMap.put(branchId, pm);
            request.getSession().setAttribute("pmMap", pmMap);
        }

        invokingActivityMap = (Map<String, String>) request.getSession().getAttribute("invokingActivityMap");
        if (invokingActivityMap == null) {
            invokingActivityMap = new HashMap<>();
            request.getSession().setAttribute("invokingActivityMap", invokingActivityMap);
        }


        List<PMTrigger> acts = controller.getAvailableActivities(trace, pm, branchId, invokingActivityMap);
        List<Map<String, Object>> endActs = findEndActs(events);

        request.getSession().setAttribute("acts", acts);
        request.getSession().setAttribute("pm", pm);
        request.getSession().setAttribute("trace", trace);
        request.getSession().setAttribute("endacts", endActs);
        request.getSession().setAttribute("automated", controller.getAutomatedActivities(trace));

        updateActiveBranches(trace, pmMap, invokingActivityMap, request, model);

        model.addAttribute("error", false);
        model.addAttribute("arts", null);
        model.addAttribute("current", "1");
        model.addAttribute("invokingActivityMap", invokingActivityMap);
        model.addAttribute("eventViews", buildEventViewList(trace));
        
        return "enact";
    }
    
    private List<Map<String, String>> updateActiveBranches(PT pt, Map<String, PM> pmMap, Map<String, String> invokingActivityMap, HttpServletRequest request, Model model) {
        List<Map<String, String>> activeBranches = new ArrayList<>();
        List<String> availableActsBranches = new ArrayList<>();
        List<String> runningActsBranches = new ArrayList<>();

        for (String branchId : pmMap.keySet()) {
            try {
                PM pm = pmMap.get(branchId);
                if (pm == null) continue;

                List<PMTrigger> availableActs = controller.getAvailableActivities(pt, pm, branchId, invokingActivityMap);
                List<PMTrigger> runningActs = controller.getRunningActivities(pt, pm, branchId, invokingActivityMap);

                if (!availableActs.isEmpty()) {
                    availableActsBranches.add(branchId);
                }
                if (!runningActs.isEmpty()) {
                    runningActsBranches.add(branchId);
                }

                if (!availableActs.isEmpty() || !runningActs.isEmpty()) {
                    Map<String, String> branchInfo = new HashMap<>();
                    branchInfo.put("branchId", branchId);
                    branchInfo.put("pmName", pm.getName());
                    branchInfo.put("invokingActivityName", invokingActivityMap.get(branchId));
                    activeBranches.add(branchInfo);
                }

            } catch (Exception e) {
                System.err.println("Error while checking branch status " + branchId + ": " + e.getMessage());
            }
        }

        request.getSession().setAttribute("activeBranches", activeBranches);
        request.getSession().setAttribute("availableActsBranches", availableActsBranches);
        request.getSession().setAttribute("runningActsBranches", runningActsBranches);
        model.addAttribute("activeBranches", activeBranches);
        model.addAttribute("availableActsBranches", availableActsBranches);
        model.addAttribute("runningActsBranches", runningActsBranches);

        return activeBranches;
    }
    
    @RequestMapping("/checkDropdownVisibility")
    @ResponseBody
    public Map<String, Boolean> checkDropdownVisibility(HttpServletRequest request) {
        Map<String, Boolean> visibility = new HashMap<>();
        List<String> availableActsBranches = (List<String>) request.getSession().getAttribute("availableActsBranches");
        List<String> runningActsBranches = (List<String>) request.getSession().getAttribute("runningActsBranches");

        visibility.put("showActivities", availableActsBranches != null && !availableActsBranches.isEmpty());
        visibility.put("showIri", runningActsBranches != null && !runningActsBranches.isEmpty());

        return visibility;
    }
    
    
    @PostMapping("/selectBranch")
    public String selectBranch(Model model, @RequestParam String branchId, HttpServletRequest request) throws Exception {
        Map<String, PM> pmMap = (Map<String, PM>) request.getSession().getAttribute("pmMap");
        PT pt = (PT) request.getSession().getAttribute("trace");
        invokingActivityMap = (Map<String, String>) request.getSession().getAttribute("invokingActivityMap");
        if (invokingActivityMap == null) {
            invokingActivityMap = new HashMap<>();
            request.getSession().setAttribute("invokingActivityMap", invokingActivityMap);
        }

        PM pm = pmMap.get(branchId);
        request.getSession().setAttribute("pm", pm);
        request.getSession().setAttribute("currentBranchId", branchId);

        List<PMTrigger> availableActs = controller.getAvailableActivities(pt, pm, branchId, invokingActivityMap);
        List<PMTrigger> runningActs = controller.getRunningActivities(pt, pm, branchId, invokingActivityMap);

        List<Map<String, Object>> endActs = new ArrayList<>();
        for (PMTrigger trigger : runningActs) {
            Map<String, Object> endAct = new HashMap<>();
            endAct.put("activity", trigger.getPort().getActivity());
            endAct.put("branchId", trigger.getBranchId());
            endActs.add(endAct);
        }

        Map<String, List<PMTrigger>> branchActsMap = (Map<String, List<PMTrigger>>) request.getSession().getAttribute("branchActsMap");
        if (branchActsMap == null) {
            branchActsMap = new HashMap<>();
        }
        branchActsMap.put(branchId, availableActs);

        request.getSession().setAttribute("acts", availableActs);
        request.getSession().setAttribute("endacts", endActs);
        request.getSession().setAttribute("branchActsMap", branchActsMap);

        model.addAttribute("acts", availableActs);
        model.addAttribute("endacts", endActs);
        model.addAttribute("invokingActivityMap", invokingActivityMap);
        updateActiveBranches(pt, pmMap, invokingActivityMap, request, model);

        return "enact";
    }

    private List<PMTrigger> findStartActs(PT trace, String branchId) throws Exception {
        System.out.println("[findStartActs] Starting for branch: " + branchId);
        PM pmEnacted = trace.getPmEnacted();
        List<Node> nodes = pmEnacted.getNodes();
        List<PMTrigger> result = new ArrayList<>();
        List<Event> events = trace.getEvents();

        List<String> startedActivities = new ArrayList<>();
        for (Event event : events) {
            if (event instanceof StartActivityEvent) {
                StartActivityEvent startEvent = (StartActivityEvent) event;
                String activityIri = startEvent.getRelatesTo().getActivity().getIri();
                boolean ended = false;
                for (Event endEvent : events) {
                    if (endEvent instanceof EndActivityEvent) {
                        EndActivityEvent endEventCast = (EndActivityEvent) endEvent;
                        if (endEventCast.getRelatesTo().getActivity().getIri().equals(activityIri)) {
                            ended = true;
                            break;
                        }
                    }
                }
                if (!ended) {
                    startedActivities.add(activityIri);
                    System.out.println("[findStartActs] Activity started but not ended: " + activityIri + " in branch: " + startEvent.getBranchId());
                }
            }
        }

        System.out.println("[findStartActs] Started activities: " + startedActivities);
        for (Node node : nodes) {
            if (node instanceof ForkJoinNode) {
                ForkJoinNode fj = (ForkJoinNode) node;
                if (fj.getNextNodes().size() > 1) {
                	List<Pair<String, String>> next = controller.findNextNodes(fj.getIri(), trace.getIri(), controller.getExecutionFlowManager());
                    System.out.println("[findStartActs] findNextNodes for fork " + fj.getIri() + " returned: " + next);

                    List<String> actIris = next.stream()
                            .map(Pair::getSnd)
                            .filter(Objects::nonNull)
                            .collect(Collectors.toList());

                    List<String> availableActs = actIris.stream()
                            .filter(iri -> !startedActivities.contains(iri))
                            .collect(Collectors.toList());

                    System.out.println("[findStartActs] Available activities: " + availableActs);
                    next.removeIf(pair -> pair.getSnd() == null || !availableActs.contains(pair.getSnd()));

                    List<PMTrigger> triggers = findElements(pmEnacted, next, branchId, null);
                    triggers.removeIf(trigger -> {
                        if (trigger.getPort() != null) {
                            String iri = trigger.getPort().getActivity().getIri();
                            return startedActivities.contains(iri);
                        }
                        return false;
                    });

                    System.out.println("[findStartActs] Final activities for branch " + branchId + ": " + triggers);
                    result.addAll(triggers);
                }
            }
        }

        System.out.println("[findStartActs] Returning result: " + result);
        return result;
    }
    
    private List<Map<String, Object>> findEndActs(List<Event> events) {
        List<Map<String, Object>> result = new ArrayList<>();
        for (Event event : events) {
            if (event instanceof StartActivityEvent) {
                StartActivityEvent startEv = (StartActivityEvent) event;

                if (startEv.getRelatesTo().getActivity() instanceof InvokingActivity) {
                    continue;
                }

                String index = startEv.getIri().split("#")[1].split("start_activity")[1];
                boolean found = false;
                for (Event event2 : events) {
                    if (event2 instanceof EndActivityEvent && event2.getIri().split("#")[1].equals("end_activity" + index)) {
                        found = true;
                        break;
                    }
                }
                if (!found) {
                    Map<String, Object> endAct = new HashMap<>();
                    endAct.put("activity", startEv.getRelatesTo().getActivity());
                    endAct.put("branchId", startEv.getBranchId());
                    result.add(endAct);
                }
            }
        }
        return result;
    }

    @PostMapping("/reload")
    public String reload(Model model, HttpServletRequest request) throws Exception {
        String currentTraceIri = (String) request.getSession().getAttribute("currentTraceIri");
        PT pt = (PT) request.getSession().getAttribute("trace");
        PM pm = (PM) request.getSession().getAttribute("pm");
        Map<String, PM> pmMap = (Map<String, PM>) request.getSession().getAttribute("pmMap");

        if (pmMap == null) {
            pmMap = new HashMap<>();
            request.getSession().setAttribute("pmMap", pmMap);
        }

        Stack<String> branchStack = branchStackMap.getOrDefault(currentTraceIri, new Stack<>());
        String branchId = branchStack.isEmpty() 
                ? controller.generateBranchId(currentTraceIri, null)
                : branchStack.peek();

        List<Event> events = controller.getEvents(pt.getIri(), branchId);
        pt.setEvents(events);

        List<Map<String, Object>> endActs = findEndActs(pt.getEvents());
        List<PMTrigger> acts = findStartActs(pt, branchId);

        if (pt.getEvents().get(pt.getEvents().size() - 1) instanceof EndActivityEvent) {
            EndActivityEvent last = (EndActivityEvent) pt.getEvents().get(pt.getEvents().size() - 1);
            List<Pair<String, String>> iris = controller.findNextNodes(last.getRelatesTo().getIri(), pt.getIri(), controller.getExecutionFlowManager());
            acts.addAll(findElements(pm, iris, branchId, null));
        }

        Map<String, List<PMTrigger>> branchActsMap = (Map<String, List<PMTrigger>>) request.getSession().getAttribute("branchActsMap");
        if (branchActsMap == null) branchActsMap = new HashMap<>();
        branchActsMap.put(branchId, acts);

        invokingActivityMap = (Map<String, String>) request.getSession().getAttribute("invokingActivityMap");
        if (invokingActivityMap == null) {
            invokingActivityMap = new HashMap<>();
            request.getSession().setAttribute("invokingActivityMap", invokingActivityMap);
        }

        request.getSession().setAttribute("acts", acts);
        request.getSession().setAttribute("branchActsMap", branchActsMap);
        request.getSession().setAttribute("endacts", endActs);
        request.getSession().setAttribute("trace", pt);
        request.getSession().setAttribute("automated", controller.getAutomatedActivities(pt));

        model.addAttribute("error", false);
        model.addAttribute("arts", null);
        model.addAttribute("current", "1");
        model.addAttribute("invokingActivityMap", invokingActivityMap);
        model.addAttribute("eventViews", buildEventViewList(pt));
        updateActiveBranches(pt, pmMap, invokingActivityMap, request, model);

        return "enact";
    }
    
    
    @PostMapping("/loadActivitiesFragment")
    public String loadActivitiesFragment(@RequestParam String pmId, Model model, HttpServletRequest request) throws Exception {
        PT pt = (PT) request.getSession().getAttribute("trace");
        Map<String, PM> pmMap = (Map<String, PM>) request.getSession().getAttribute("pmMap");

        if (pt == null || pmMap == null) {
            model.addAttribute("activities", new ArrayList<>()); 
            return "fragments/activities :: activitiesFragment";
        }

        PM selectedPM = pmMap.get(pmId);
        if (selectedPM == null) {
            model.addAttribute("activities", new ArrayList<>());
            return "fragments/activities :: activitiesFragment";
        }

        List<PMTrigger> acts = controller.getAvailableActivities(pt, selectedPM, pmId, invokingActivityMap);

        model.addAttribute("activities", acts);
        model.addAttribute("branchId", pmId);
        return "fragments/activities :: activitiesFragment";
    }
    
    private List<EventView> buildEventViewList(PT trace) {
        List<EventView> eventViews = new ArrayList<>();

        for (Event ev : trace.getEvents()) {
            if (ev instanceof StartActivityEvent) {
                String target = ((StartActivityEvent) ev).getRelatesTo().getActivity().getName();
                eventViews.add(new EventView("StartActivity", "Begin", target));

            } else if (ev instanceof EndActivityEvent) {
                String target = ((EndActivityEvent) ev).getRelatesTo().getActivity().getName();
                eventViews.add(new EventView("EndActivity", "End", target));

            } else if (ev instanceof StartTraceEvent) {
                String target = trace.getPmEnacted().getName();
                eventViews.add(new EventView("StartTrace", "Start Trace", target));

            } else if (ev instanceof ua.be.wee.model.pt.EndTraceEvent) {
                String target = trace.getPmEnacted().getName();
                eventViews.add(new EventView("EndTrace", "End Trace", target));
            }
        }

        return eventViews;
    }

    @PostMapping("/startAct")
    public String startActivity(Model model, @RequestParam String iri, HttpServletRequest request) throws Exception {
        System.out.println("[startAct] Starting activity: " + iri);
        PT pt = (PT) request.getSession().getAttribute("trace");
        String traceIri = pt.getIri();
        String currentBranchId = (String) request.getSession().getAttribute("currentBranchId");

        Map<String, PM> pmMap = (Map<String, PM>) request.getSession().getAttribute("pmMap");
        invokingActivityMap = (Map<String, String>) request.getSession().getAttribute("invokingActivityMap");
        if (invokingActivityMap == null) {
            invokingActivityMap = new HashMap<>();
            request.getSession().setAttribute("invokingActivityMap", invokingActivityMap);
        }

        Map<String, List<PMTrigger>> branchActsMap = (Map<String, List<PMTrigger>>) request.getSession().getAttribute("branchActsMap");
        if (branchActsMap == null) {
            branchActsMap = new HashMap<>();
            request.getSession().setAttribute("branchActsMap", branchActsMap);
        }

        List<PMTrigger> acts = branchActsMap.getOrDefault(currentBranchId, new ArrayList<>());
        List<Map<String, Object>> endacts = (List<Map<String, Object>>) request.getSession().getAttribute("endacts");
        List<TraceArtifact> arts = (List<TraceArtifact>) request.getSession().getAttribute("arts");

        PMTrigger element = findPMTrigger(acts, iri);
        if (element == null) {
            System.out.println("[startAct] PMTrigger not found for IRI: " + iri + ". Checking if it is a FinalNode.");
            if (iri.toLowerCase().contains("final")) {
                 PM currentPm = (PM) request.getSession().getAttribute("pm");
                 FinalNode finalNode = (FinalNode) currentPm.getNodes().stream()
                    .filter(n -> n instanceof FinalNode && n.getIri().equals(iri))
                    .findFirst().orElse(null);
                 if(finalNode != null) {
                    element = new PMTrigger();
                    element.setNode(finalNode);
                 }
            }
            if(element == null){
                System.out.println("[startAct] Trigger element is still null. Cannot proceed.");
                updateActiveBranches(pt, pmMap, invokingActivityMap, request, model);
                return "enact";
            }
        }

        if (element.getPort() != null) {
            ControlInputPort port = (ControlInputPort) element.getPort();
            Activity activity = port.getActivity();

            if (activity == null) {
                System.out.println("[startAct] Activity is null for port: " + port.getIri());
                return "error";
            }

            System.out.println("[startAct] Processing activity: " + activity.getIri());

            if (activity instanceof InvokingActivity) {
                InvokingActivity invokingActivity = (InvokingActivity) activity;
                String eventIri = controller.addStartEvent(pt, port, invokingActivity, arts, currentBranchId);
                if (eventIri == null) {
                    updateActiveBranches(pt, pmMap, invokingActivityMap, request, model);
                    return "enact";
                }
                String newBranchId = controller.generateBranchId(pt.getIri(), currentBranchId);
                controller.getExecutionFlowManager().createExecutionFlow(pt, activity, newBranchId, invokingActivity.getName());
                PM invokedPM = controller.getPM(invokingActivity.getInvokedPM());
                String pmTraceIri = controller.createStartTraceEvent(pt, invokedPM, eventIri, pt.getIri());
                pmMap.put(newBranchId, invokedPM);
                invokingActivityMap.put(newBranchId, invokingActivity.getName());
                
                List<PMTrigger> newActs = controller.getAvailableActivities(pt, invokedPM, newBranchId, invokingActivityMap);

                branchActsMap.put(newBranchId, newActs);
                request.getSession().setAttribute("pm", invokedPM);
                request.getSession().setAttribute("currentBranchId", newBranchId);
                request.getSession().setAttribute("currentTraceIri", pmTraceIri);
            } else {
                String eventIri = controller.addStartEvent(pt, port, activity, arts, currentBranchId);
                if (eventIri == null) {
                    updateActiveBranches(pt, pmMap, invokingActivityMap, request, model);
                    return "enact";
                }
                controller.getExecutionFlowManager().createExecutionFlow(pt, activity, currentBranchId, invokingActivityMap.get(currentBranchId));
                if (activity instanceof AutomatedActivity) {
                    controller.callAutomatedActivity(pt, (AutomatedActivity) activity, port, arts, null);
                    ((AutomatedActivity) activity).setAutomatedStatus(AutomatedStatus.IN_PROGRESS);
                } else {
                    Map<String, Object> endAct = new HashMap<>();
                    endAct.put("activity", activity);
                    endAct.put("branchId", currentBranchId);
                    endacts.add(endAct);
                }
                removeElement(acts, iri);
                PM currentPm = pmMap.get(currentBranchId);
                List<PMTrigger> updatedActs = controller.getAvailableActivities(pt, currentPm, currentBranchId, invokingActivityMap);
                
                branchActsMap.put(currentBranchId, updatedActs);
            }
        } else if (element.getPort() == null && element.getNode() != null) {
            System.out.println("[startAct] PM finished. Checking if it was invoked by another activity.");
            if (invokingActivityMap.containsKey(currentBranchId)) {
                System.out.println("[startAct] Sub-process finished. Automating return to parent process.");
                String invokingActivityName = invokingActivityMap.get(currentBranchId);
                Event endTraceEvent = controller.addEndTraceEvent(pt.getIri(), pt.getLastEvent().getIri(), pt.getPmEnacted().getIri());
                pt.addEvent(endTraceEvent);
                controller.getExecutionFlowManager().clearTokensForTrace(pt.getIri());
                String parentBranchId = null;
                int lastSlashIndex = currentBranchId.lastIndexOf('/');
                if (lastSlashIndex > 0) {
                    parentBranchId = currentBranchId.substring(0, lastSlashIndex);
                }
                if (parentBranchId == null) {
                    throw new IllegalStateException("Could not derive parent branch for sub-process branch: " + currentBranchId);
                }
                PM parentPM = pmMap.get(parentBranchId);
                InvokingActivity invokingActivity = (InvokingActivity) parentPM.getNodes().stream()
                        .filter(node -> node instanceof InvokingActivity && ((InvokingActivity) node).getName().equals(invokingActivityName))
                        .findFirst()
                        .orElseThrow(() -> new IllegalStateException("Could not find InvokingActivity with name '" + invokingActivityName + "' in parent PM."));
                ControlOutputPort outputPort = invokingActivity.getCtrlOutPorts().stream().findFirst()
                        .orElseThrow(() -> new IllegalStateException("No ControlOutputPort found on InvokingActivity " + invokingActivity.getName()));
                
                controller.addEndEvent(pt, new ArrayList<>(), outputPort, parentBranchId);
                
                List<PMTrigger> parentActs = controller.getAvailableActivities(pt, parentPM, parentBranchId, invokingActivityMap);
                
                if (parentActs.size() == 1 && parentActs.get(0).getNode() instanceof FinalNode) {
                    if (!invokingActivityMap.containsKey(parentBranchId)) {
                        System.out.println("[startAct] getAvailableActivities returned only FinalNode. Auto-finalizing.");
                        return finalizeProcess(model, request);
                    }
                }

                branchActsMap.put(parentBranchId, parentActs);
                pmMap.remove(currentBranchId);
                invokingActivityMap.remove(currentBranchId);
                branchActsMap.remove(currentBranchId);
                request.getSession().setAttribute("pm", parentPM);
                request.getSession().setAttribute("currentBranchId", parentBranchId);
                request.getSession().setAttribute("acts", parentActs);
                endacts.removeIf(endAct -> ((Activity)endAct.get("activity")).getIri().equals(invokingActivity.getIri()));
            } else {
                System.out.println("[startAct] Main process finished. Redirecting to end page.");
                return finalizeProcess(model, request);
            }
        }

        List<Map<String, String>> activeBranches = updateActiveBranches(pt, pmMap, invokingActivityMap, request, model);
        request.getSession().setAttribute("pmMap", pmMap);
        request.getSession().setAttribute("invokingActivityMap", invokingActivityMap);
        request.getSession().setAttribute("automated", controller.getAutomatedActivities(pt));
        request.getSession().setAttribute("endacts", endacts);
        request.getSession().setAttribute("branchActsMap", branchActsMap);
        request.getSession().removeAttribute("arts");
        
        currentBranchId = (String) request.getSession().getAttribute("currentBranchId");
        request.getSession().setAttribute("acts", branchActsMap.getOrDefault(currentBranchId, new ArrayList<>()));
        
        model.addAttribute("invokingActivityMap", invokingActivityMap);
        model.addAttribute("eventViews", buildEventViewList(pt));
        model.addAttribute("activeBranches", activeBranches);
        model.addAttribute("acts", branchActsMap.getOrDefault(currentBranchId, new ArrayList<>()));
        
        System.out.println("[startAct] Returning model with activeBranches: " + activeBranches);
        return "enact";
    }

    private String finalizeProcess(Model model, HttpServletRequest request) throws Exception {
        System.out.println("[finalizeProcess] Main process is finalizing.");
        PT pt = (PT) request.getSession().getAttribute("trace");
        PM finishedPM = pt.getPmEnacted();

        List<Event> rawEvents = controller.getAllEventsForTrace(pt.getIri());
        List<Event> timelineEvents = new ArrayList<>();

        for (Event ev : rawEvents) {
            if (ev instanceof StartTraceEvent) {
                if (ev.getIri().equals(pt.getIri())) {
                    timelineEvents.add(ev);
                }
            } else if (!(ev instanceof EndTraceEvent)) {
                timelineEvents.add(ev);
            }
        }

        String finalEventIri = pt.getIri() + "_end";
        rawEvents.stream()
            .filter(e -> e instanceof EndTraceEvent && e.getIri().equals(finalEventIri))
            .findFirst()
            .ifPresent(timelineEvents::add); 

        pt.setEvents(timelineEvents);
        
        boolean alreadyEndedInDb = rawEvents.stream().anyMatch(e -> e instanceof EndTraceEvent && e.getIri().equals(finalEventIri));
        if (!alreadyEndedInDb) {
            Event endTraceEvent = controller.addEndTraceEvent(pt.getIri(), pt.getLastEvent().getIri(), finishedPM.getIri());
            pt.addEvent(endTraceEvent); 
        }

        controller.getExecutionFlowManager().clearTokensForTrace(pt.getIri());
        branchStackMap.remove(pt.getIri());

        // Clean up session
        request.getSession().removeAttribute("acts");
        request.getSession().removeAttribute("endacts");
        request.getSession().removeAttribute("branchActsMap");
        request.getSession().removeAttribute("currentBranchId");
        request.getSession().removeAttribute("pmMap");
        request.getSession().removeAttribute("invokingActivityMap");
        request.getSession().removeAttribute("activeBranches");
        request.getSession().removeAttribute("pm");
        
        model.addAttribute("pm", finishedPM);
        model.addAttribute("trace", pt); 
        
        request.getSession().removeAttribute("trace");

        return "endEnactment";
    }


    @PostMapping("/endAct")
    public String endActivityWithArtifacts(Model model, @RequestParam String port, @RequestParam String activity, @RequestParam String branchId,
            HttpServletRequest request) throws Exception {

        PT pt = (PT) request.getSession().getAttribute("trace");
        String traceIri = pt.getIri();
        List<Map<String, Object>> endacts = (List<Map<String, Object>>) request.getSession().getAttribute("endacts");
        invokingActivityMap = (Map<String, String>) request.getSession().getAttribute("invokingActivityMap");
        Map<String, PM> pmMap = (Map<String, PM>) request.getSession().getAttribute("pmMap");

        if (pmMap == null) throw new IllegalStateException("pmMap not initialized.");
        if (invokingActivityMap == null) {
            invokingActivityMap = new HashMap<>();
            request.getSession().setAttribute("invokingActivityMap", invokingActivityMap);
        }

        PM pm = pmMap.get(branchId);
        if (pm == null) return "error";

        Activity act = (Activity) pm.getNode(activity);
        if (act == null) return "error";

        Map<String, List<PMTrigger>> branchActsMap = (Map<String, List<PMTrigger>>) request.getSession().getAttribute("branchActsMap");
        if (branchActsMap == null) branchActsMap = new HashMap<>();

        if (act instanceof InvokingActivity) {
            InvokingActivity invokingAct = (InvokingActivity) act;

            ControlOutputPort selectedPort = invokingAct.getCtrlOutPorts().stream()
                .filter(p -> p.getIri().equals(port))
                .findFirst()
                .orElseThrow(() -> new IllegalStateException("Selected output port with IRI " + port + " not found for InvokingActivity."));

            List<TraceArtifact> traceArts = processArtifacts(invokingAct.getOutputs(), request);
            controller.addEndEvent(pt, traceArts, selectedPort, branchId);

            List<Pair<String, String>> iris = controller.findNextNodes(selectedPort.getIri(), traceIri, controller.getExecutionFlowManager());
            List<PMTrigger> acts = findElements(pm, iris, branchId, null);
            branchActsMap.put(branchId, acts);
            request.getSession().setAttribute("acts", acts);

            endacts.removeIf(endAct -> endAct.get("activity").equals(invokingAct));
            request.getSession().setAttribute("endacts", endacts);
        } else {
            List<ControlOutputPort> ctrlOutPorts = act.getCtrlOutPorts();
            ControlOutputPort p = ctrlOutPorts.stream().filter(out -> out.getIri().equals(port)).findFirst().orElse(null);

            List<TraceArtifact> traceArts = processArtifacts(act.getOutputs(), request);
            controller.addEndEvent(pt, traceArts, p, branchId);
            
            endacts.removeIf(endAct -> ((Activity) endAct.get("activity")).getIri().equals(act.getIri()));

            if (p != null) {
                List<PMTrigger> acts = controller.getAvailableActivities(pt, pm, branchId, invokingActivityMap);

                if (acts.size() == 1 && acts.get(0).getNode() instanceof FinalNode) {
                    if (invokingActivityMap == null || !invokingActivityMap.containsKey(branchId)) {
                        System.out.println("[endActivityWithArtifacts] Join completed, automatically finalizing main process.");
                        return finalizeProcess(model, request);
                    }
                }
                
                branchActsMap.put(branchId, acts);
                request.getSession().setAttribute("acts", acts);
            }
        }

        List<Map<String, String>> activeBranches = new ArrayList<>();
        for (String bid : pmMap.keySet()) {
            Map<String, String> branchInfo = new HashMap<>();
            branchInfo.put("branchId", bid);
            branchInfo.put("pmName", pmMap.get(bid).getName());
            branchInfo.put("invokingActivityName", invokingActivityMap.get(bid));
            activeBranches.add(branchInfo);
        }
        model.addAttribute("activeBranches", activeBranches);

        request.getSession().setAttribute("trace", pt);
        request.getSession().setAttribute("branchActsMap", branchActsMap);
        request.getSession().setAttribute("automated", controller.getAutomatedActivities(pt));
        model.addAttribute("arts", null);
        model.addAttribute("current", "1");
        model.addAttribute("invokingActivityMap", invokingActivityMap);
        model.addAttribute("eventViews", buildEventViewList(pt));
        updateActiveBranches(pt, pmMap, invokingActivityMap, request, model);

        return "enact";
    }


    private List<TraceArtifact> processArtifacts(List<Artifact> arts, HttpServletRequest request) throws Exception {
        List<TraceArtifact> traceArts = new ArrayList<>();
        ArrayList<Part> parts = (ArrayList<Part>) request.getParts();

        for (Artifact artifact : arts) {
            for (Part part : parts) {
                if (part.getName().equals(artifact.getName())) {
                    TraceArtifact tArt = new TraceArtifact();
                    tArt.setLocation(part.getSubmittedFileName());
                    tArt.setRelatesTo(artifact);
                    traceArts.add(tArt);
                    controller.uploadArtifact(part.getInputStream(), part.getSubmittedFileName());
                }
            }
        }

        return traceArts;
    }

    @RequestMapping("/endselect")
    public String selectActivityToEnd(Model model, @RequestParam String iri, HttpServletRequest request)
            throws Exception {
        PM pm = (PM) request.getSession().getAttribute("pm");
        Map<String, PM> pmMap = (Map<String, PM>) request.getSession().getAttribute("pmMap");
        invokingActivityMap = (Map<String, String>) request.getSession().getAttribute("invokingActivityMap");
        if (invokingActivityMap == null) {
            invokingActivityMap = new HashMap<>();
            request.getSession().setAttribute("invokingActivityMap", invokingActivityMap);
        }

        if (!iri.equals("1")) {
            if (pm.getNode(iri) instanceof Activity) {
                Activity act = (Activity) pm.getNode(iri);
                model.addAttribute("act", pm.getNode(iri));
                model.addAttribute("arts", act.getOutputs());
                String invokingActivityName = invokingActivityMap.get(act.getIri());
                model.addAttribute("invokingActivityName", invokingActivityName);
            } else {
                model.addAttribute("endBool", true);
            }
        } else {
            model.addAttribute("arts", null);
        }

        List<Map<String, String>> activeBranches = new ArrayList<>();
        for (String bid : pmMap.keySet()) {
            Map<String, String> branchInfo = new HashMap<>();
            branchInfo.put("branchId", bid);
            branchInfo.put("pmName", pmMap.get(bid).getName());
            branchInfo.put("invokingActivityName", invokingActivityMap.get(bid));
            activeBranches.add(branchInfo);
        }
        model.addAttribute("activeBranches", activeBranches);

        model.addAttribute("invokingActivityMap", invokingActivityMap);
        return "enactEnd";
    }

    @PostMapping("/inarts")
    public String inarts(@RequestParam String iri, HttpServletRequest request, Model model) throws Exception {
        PT pt = (PT) request.getSession().getAttribute("trace");
        Map<String, PM> pmMap = (Map<String, PM>) request.getSession().getAttribute("pmMap");
        invokingActivityMap = (Map<String, String>) request.getSession().getAttribute("invokingActivityMap");

        if (invokingActivityMap == null) {
            invokingActivityMap = new HashMap<>();
            request.getSession().setAttribute("invokingActivityMap", invokingActivityMap);
        }

        String currentBranchId = (String) request.getSession().getAttribute("currentBranchId");
        PM pm = pmMap.get(currentBranchId);

        List<PMTrigger> acts = (List<PMTrigger>) request.getSession().getAttribute("acts");
        Map<String, List<PMTrigger>> branchActsMap = (Map<String, List<PMTrigger>>) request.getSession().getAttribute("branchActsMap");
        if (branchActsMap == null) {
            branchActsMap = new HashMap<>();
        }

        PMTrigger element = findPMTrigger(acts, iri);
        if (element == null) {
            System.out.println("[inarts] PMTrigger not found for IRI: " + iri);
            return "error";
        }

        if (element.getPort() == null) {
            model.addAttribute("errorMessage", "Selected element has no port.");
            return "error";
        }

        ControlInputPort port = (ControlInputPort) element.getPort();
        Activity activity = port.getActivity();

        List<TraceArtifact> arts = new ArrayList<>();
        if (activity != null && activity.getInputs() != null) {
            for (Artifact art : activity.getInputs()) {
                TraceArtifact t = new TraceArtifact();
                t.setRelatesTo(art);
                arts.add(t);
            }
        }

        request.getSession().setAttribute("arts", arts);
        request.getSession().setAttribute("act", activity);
        request.getSession().setAttribute("acts", acts);

        branchActsMap.put(currentBranchId, acts);
        request.getSession().setAttribute("branchActsMap", branchActsMap);

        model.addAttribute("arts", arts);
        model.addAttribute("current", "1");
        model.addAttribute("invokingActivityMap", invokingActivityMap);
        model.addAttribute("eventViews", buildEventViewList(pt));
        updateActiveBranches(pt, pmMap, invokingActivityMap, request, model);

        return "inarts";
    }


    private List<PMTrigger> findElements(PM pm, List<Pair<String, String>> iris, String branchId, String invokingActivityName) {
        List<PMTrigger> acts = new ArrayList<>();
        for (Pair<String, String> pair : iris) {
            PMTrigger tr = new PMTrigger();
            String nodeIri = (pair.getSnd() != null) ? pair.getSnd() : pair.getFst();
            Node node = pm.getNode(nodeIri);
            if (node == null) continue;
            if (pair.getSnd() != null && node instanceof Activity) {
                Activity activity = (Activity) node;
                ControlInputPort port = activity.getCtrlInPorts().stream()
                        .filter(ctr -> ctr.getIri().equals(pair.getFst())).findFirst().orElse(null);
                if (port != null) {
                    tr.setPort(port);
                    tr.setBranchId(branchId);
                    tr.setInvokingActivityName(invokingActivityName);
                    acts.add(tr);
                }
            } else if (node instanceof FinalNode) {
                tr.setNode((FinalNode) node);
                tr.setBranchId(branchId);
                acts.add(tr);
            }
        }
        return acts;
    }

    private PMTrigger findPMTrigger(List<PMTrigger> acts, String iri) {
        for (PMTrigger pmTrigger : acts) {
            if (pmTrigger.getIri().equals(iri)) {
                return pmTrigger;
            }
        }
        return null;
    }

    private void removeElement(List<PMTrigger> acts, String iri) {
        for (PMTrigger pmTrigger : acts) {
            if (pmTrigger.getIri().equals(iri)) {
                acts.remove(pmTrigger);
                break;
            }
        }
    }
}