# -*- coding: utf-8 -*-
from pyrevit import revit, DB, HOST_APP
from pyrevit import script
from pyrevit.preflight import PreflightTestCase


output = script.get_output()
output.close_others()
doc = HOST_APP.doc

# get the linked documents
link_instances = DB.FilteredElementCollector(doc).OfClass(DB.RevitLinkInstance).ToElements()
link_docs = []
if link_instances:
    for link_instance in link_instances:
        link_doc = link_instance.GetLinkDocument()
        link_docs.append(link_doc)


def ef_tools_check(document = doc, op = output):
    # get the document info
    document_info(document, op)
    # get the walls info
    walls_info(document, op)
    # get the grids info
    grids_info(document, op)


def document_info(document = doc, op = output):
    # get the project name
    project_title = document.Title
    project_name = document.ProjectInformation.Name
    # get the project number
    project_number = document.ProjectInformation.Number
    # get the project address
    project_address = document.ProjectInformation.Address

    op.print_md("## Project Info")
    op.print_md("### {}".format(project_name))
    op.print_md("Project Title: {}".format(project_title))
    op.print_md("Project Number: {}".format(project_number))
    op.print_md("Project Address: {}".format(project_address))


def walls_info(document = doc, op = output):
    walls = DB.FilteredElementCollector(document).OfCategory(DB.BuiltInCategory.OST_Walls)
    walls_instances = walls.WhereElementIsNotElementType()
    if walls_instances.GetElementCount() > 0:
        # get the wall count
        walls_count = walls_instances.GetElementCount()
        op.print_md("**Walls: **{}".format(walls_count))
        # get the wall name
        walls_names = [wall.Name for wall in walls_instances]
        # get the wall id
        walls_ids = []
        for wall in walls_instances.ToElements():
            walls_ids.append(op.linkify(wall.Id))
        # get wall length
        wall_length = []
        for wall in walls_instances.ToElements():
            try:
                wall_length.append(wall.get_Parameter(DB.BuiltInParameter.CURVE_ELEM_LENGTH).AsValueString()+"mm")
            except:
                wall_length.append("N/A")
        # get the wall workset
        walls_workset = elements_worksets(document, walls_instances)

        op.print_table(title="Walls", columns=["Wall name", "Wall id", "Wall Length", "Wall Workset"], table_data=zip(walls_names, walls_ids, wall_length, walls_workset))
    else:
        op.print_md("**No walls found**")


def grids_info(document = doc, op = output):
    grids = DB.FilteredElementCollector(document).OfCategory(DB.BuiltInCategory.OST_Grids)
    grid_instances = grids.WhereElementIsNotElementType()
    if grid_instances.GetElementCount() > 0:
        # get the grid count
        grids_count = grid_instances.GetElementCount()
        op.print_md("**Grids: **{}".format(grids_count))
        # get the grid name
        grid_names = [grid.Name for grid in grid_instances]
        # get the grid id
        grids_ids = []
        for grid in grid_instances.ToElements():
            grids_ids.append(op.linkify(grid.Id))
        # get the grid workset
        grids_workset = elements_worksets(document, grid_instances)
        op.print_table(title="Grids", columns=["Grid name", "Grid id", "Grid workset"], table_data=zip(grid_names, grids_ids, grids_workset))
    else:
        op.print_md("**No grids found**")


def elements_worksets(document = doc, element_instances = None):
    element_worksets = []
    for element in element_instances.ToElements():
        if document.IsWorkshared:
            element_worksets.append(element.get_Parameter(DB.BuiltInParameter.ELEM_PARTITION_PARAM).AsValueString())
        else:
            element_worksets.append("File is not workshared")
    return element_worksets


class ModelChecker(PreflightTestCase):
    """
    List all walls and their info and get some of the model info
    This QC tool returns you with the following data:
        - Model info: project name, project number, project address
        - Wall count, wall name, wall id, wall length, wall workset (if workshared)
        - Grid count, grid name, grid id, grid workset (if workshared)
    """
    
    name = "EF_tools_checker"
    author = "Jean-Marc Couffin"

    def setUp(self, doc, output):
        pass

    def startTest(self, doc, output):
        output.print_md("# EF_Tools check")
        ef_tools_check(doc, output)
        for link_doc in link_docs:
            if link_doc:
                ef_tools_check(link_doc)

    def tearDown(self, doc, output):
        pass

    def doCleanups(self, doc, output):
        pass