'''
Generated files comparison script.

Takes two directories as generated by anchorscad_runner.py and compares the file contents.

This can be used to verify any potential API changes that may have affected models.

Author: gianni
'''

import sys
import os
import filecmp
import argparse
import re
from dataclasses import dataclass, field
from dataclasses_json import dataclass_json
from typing import Optional, List, Tuple, Dict
from PIL import Image, ImageChops

import http.server
import socketserver
from urllib.parse import unquote
import mimetypes

    
@dataclass
class FileDifference:
    '''Represents a difference between two files.'''
    line_num: int
    char_position: int
    line1: str
    line2: str
    token1: str
    token2: str

    def get_snippet_with_diff(self, window_size: int = 100):
        '''Return a snippet of the lines with the differing tokens underlined.
        
        Returns a tuple of (snippet1, underline1, snippet2, underline2) where:
        - snippet1 is a window of the line containing the first differing token
        - underline1 is a string with the first differing token underlined
        - snippet2 is a window of the line containing the second differing token
        - underline2 is a string with the second differing token underlined
        '''
        def create_snippet(line, token, char_position):
            start = max(char_position - window_size // 2, 0)
            end = min(char_position + window_size // 2, len(line))
            snippet = line[start:end]
            if start > 0:
                snippet = "..." + snippet
            if end < len(line):
                snippet = snippet + "..."
            underline = ' ' * (char_position - start) + '^' * len(token)
            return snippet, underline

        snippet1, underline1 = create_snippet(self.line1, self.token1, self.char_position)
        snippet2, underline2 = create_snippet(self.line2, self.token2, self.char_position)
        
        return snippet1, underline1, snippet2, underline2

@dataclass_json
@dataclass
class ModelFiles:
    '''Model file paths.'''
    scad: str=None
    png: str=None
    stl: str=None
    f3mf: str=None
    graph: str=None
    path: str=None
    others: List[str]=field(default_factory=list)
    
    def add_file(self, fname: str):
        '''Add a file to the model files.'''
        if fname.endswith('.scad'):
            self.scad = fname
        elif fname.endswith('.png'):
            self.png = fname
        elif fname.endswith('.stl'):
            self.stl = fname
        elif fname.endswith('.3mf'):
            self.f3mf = fname
        elif fname.endswith('.dot.svg'):
            self.graph = fname
        elif fname.endswith('.html'):
            self.path = fname
        else:
            self.others.append(fname)


@dataclass_json
@dataclass
class ModelDiffResult:
    '''Model diff results.'''
    model_files: ModelFiles=field(default_factory=ModelFiles)
    png_diff: Optional[str]=None
    png_diff_result: Optional[bool]=None
    file_diff_text: List[str]=field(default_factory=list)
    files_only_in_one_dir: List[List[str]]=None
    has_file_diff: bool=False
    
    def add_file(self, name: str):
        '''Add a file to the diff result.'''
        # repalce '\' with '/' for windows paths
        uname = name.replace('\\', '/')
        self.model_files.add_file(uname)
        
    def add_file_to_only_in_one_dir(self, dir_number: int, name: str):
        '''Add a file to the list of files only in directory 1.'''
        uname = name.replace('\\', '/')
        if not self.files_only_in_one_dir:
            self.files_only_in_one_dir = [[], []]
        self.files_only_in_one_dir[dir_number].append(uname)
        
    def set_png_diff(self, diff_path: str):
        '''Set the path to the PNG diff image.'''
        self.png_diff = diff_path.replace('\\', '/')


@dataclass_json
@dataclass
class DirectoryDiffResults:
    '''Difference in all the directories.'''
    dirs: Tuple[str, str]
    model_diffs: Dict[str, ModelDiffResult]=field(default_factory=dict)
    
    def diff_for_model(self, model_file: str) -> ModelDiffResult:
        '''Return the diff result for the model.'''
        
        # Extract the model path including the common madel file prefix.
        model_path = model_file.split('.')[0]
        
        diff : ModelDiffResult = self.model_diffs.get(model_path, None)
        if not diff:
            diff = ModelDiffResult()
            self.model_diffs[model_path] = diff
        return diff

    def add_file_to_only_in_one_dir(self, dir_number: int, fname: str):
        diff: ModelDiffResult = self.diff_for_model(fname)
        diff.add_file_to_only_in_one_dir(dir_number, fname)
        return diff
        
    def add_no_diff_file(self, fname: str):
        diff: ModelDiffResult = self.diff_for_model(fname)
        diff.add_file(fname)
        return diff
    
def tokenize(text):
    """Tokenize the input text into a list of tokens."""
    token_pattern = re.compile(r'(\s+|[A-Za-z_]\w*|[-+]?\d*\.\d+|\d+|[-+]?\d*\.\d+[eE][-+]?\d+|[-+]?\d+[eE][-+]?\d+|\".*?\"|[^\s\w])')
    return token_pattern.findall(text)

def is_float(token):
    """Check if a token is a floating-point number."""
    try:
        float(token)
        return True
    except ValueError:
        return False
    
def is_string(token):
    """Check if a token is a string."""
    return token.startswith('"')

def compare_scad_tokens(tokens1, tokens2, line1, line2, line_num, epsilon) -> Optional[FileDifference]:
    """Compare lists of tokens with a given epsilon tolerance for floating-point numbers."""
    max_len = max(len(tokens1), len(tokens2))
    char_pos = 0

    last_token_size = 0
    for idx in range(max_len):
        char_pos += last_token_size
        if idx >= len(tokens1) or idx >= len(tokens2):
            t1 = tokens1[idx] if idx < len(tokens1) else "EOF"
            t2 = tokens2[idx] if idx < len(tokens2) else "EOF"
            return FileDifference(line_num, char_pos, line1, line2, t1, t2)
        
        t1, t2 = tokens1[idx], tokens2[idx]
        last_token_size = len(t1)
        
        if is_float(t1) and is_float(t2):
            if abs(float(t1) - float(t2)) > epsilon:
                return FileDifference(line_num, char_pos, line1, line2, t1, t2)
        elif is_string(t1) and is_string(t2):
            pass # Ignore strings
        elif t1 != t2:
            return FileDifference(line_num, char_pos, line1, line2, t1, t2)
    
    return None

def compare_scad_lines(lines1: list[str], lines2: list[str], epsilon=1e-6) -> Optional[FileDifference]:
    """Compare two lines while ignoring small differences in floating-point numbers."""

    max_len = max(len(lines1), len(lines2))
    
    for line_num in range(max_len):
        line1 = lines1[line_num] if line_num < len(lines1) else ""
        line2 = lines2[line_num] if line_num < len(lines2) else ""
        
        tokens1 = tokenize(line1)
        tokens2 = tokenize(line2)

        result = compare_scad_tokens(tokens1, tokens2, line1, line2, line_num + 1, epsilon)
        if result:
            return result
    
    return None

def compare_scad_files(file1, file2, epsilon=1e-6) -> Optional[FileDifference]:
    """Compare two files while ignoring small differences in floating-point numbers."""
    with open(file1, 'r') as f1, open(file2, 'r') as f2:
        lines1 = f1.readlines()
        lines2 = f2.readlines()
    return compare_scad_lines(lines1, lines2, epsilon)


def make_image_output_dir(dir1, dir2):
    '''Returns a name for the output directory for image diff files.'''
        # Extract the suffixes from dir1 and dir2
    suffix1 = dir1.split('.')[-1]
    suffix2 = dir2.split('.')[-1]

    # Determine the order based on the suffixes
    if suffix1 < suffix2:
        output_dir = f"{dir1}-{suffix2}"
    else:
        output_dir = f"{dir2}-{suffix1}"

    return output_dir
    

def compare_images(image1_path, image2_path, output_path) -> bool:
    '''Create a diff image from two input image files, save the result, and return whether a change was detected.

    Args:
        image1_path: Path to the first image file.
        image2_path: Path to the second image file.
        output_path: Path where the diff image will be saved.
        
    Returns True if a change was detected, False otherwise.'''

    image1 = Image.open(image1_path)
    image2 = Image.open(image2_path)

    # Create a diff image using ImageChops.difference
    diff_image = ImageChops.difference(image1, image2)

    # Check if the diff image is completely black (no differences)
    bbox = diff_image.getbbox()
    change_detected = bbox is not None

    if change_detected:
        # Make the diff image more visible by inverting the colors
        diff_image = ImageChops.invert(diff_image)
        
        # Save the diff image
        # Ensure the output directory exists
        output_dir = os.path.dirname(output_path)
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
            
        diff_image.save(output_path)

    return change_detected



def get_files_with_suffixes(directory, suffixes):
    """Return a set of filenames in the given directory and its subdirectories with the specified suffixes."""
    matching_files = set()
    for root, _, files in os.walk(directory):
        for f in files:
            if any(f.endswith(suffix) for suffix in suffixes):
                relative_path = os.path.relpath(os.path.join(root, f), directory)
                matching_files.add(relative_path)
    return matching_files


@dataclass(frozen=True)
class ComparisonResult:
    are_equal: bool=True
    difference: Optional[FileDifference] = field(default=None)

    def __post_init__(self):
        if self.difference is not None:
            object.__setattr__(self, 'are_equal', False)

def compare_file_contents(file1, file2, args) -> ComparisonResult:
    '''Compare the contents of two files.'''

    if file1.endswith('.scad') and file2.endswith('.scad'):
        difference = compare_scad_files(file1, file2, args.epsilon)
        return ComparisonResult(difference is None, difference)
    else:
        are_equal = filecmp.cmp(file1, file2, shallow=False)
        return ComparisonResult(are_equal)
        
def to_png_file_name(file_name):
    '''Convert a file name to a PNG file name.'''
    parts = file_name.split('.')
    parts[-1] = 'png'
    return '.'.join(parts)

def compare_directories(args) -> Tuple[bool, DirectoryDiffResults]:
    '''Compare files between two directories for given suffixes and print the results.
    
    Returns True if there are any differences between the directories for the given suffixes.
    '''
    diffResults = DirectoryDiffResults(dirs=(args.dir1, args.dir2))
    files1 = get_files_with_suffixes(args.dir1, args.suffixes)
    files2 = get_files_with_suffixes(args.dir2, args.suffixes)

    common_files = files1.intersection(files2)
    unique_to_dir1 = files1 - files2
    unique_to_dir2 = files2 - files1

    unequal_files = []
    equal_files = []

    for file in common_files:
        file1_path = os.path.join(args.dir1, file)
        file2_path = os.path.join(args.dir2, file)
        comparison_result = compare_file_contents(file1_path, file2_path, args)
        if comparison_result.are_equal:
            equal_files.append(file)
        else:
            unequal_files.append((file, comparison_result.difference))

    if unequal_files:
        print(f'Common files with different contents in both directories with suffixes {args.suffixes}:')
        for e in unequal_files:
            fname: str = e[0]
            uname = fname.replace('\\', '/')
            difference: FileDifference = e[1]
            diff: ModelDiffResult = diffResults.diff_for_model(uname)
            diff.has_file_diff = True  # Mark that there is a difference in the model.
            diff.add_file(uname)
            lines = []
            if difference:
                lines.append(f'  {uname} at line {difference.line_num}, '
                             f'character position {difference.char_position + 1}:')
                snippet = difference.get_snippet_with_diff()
                lines.append(f'    File 1: {snippet[0]}')
                lines.append(f'            {snippet[1]}')
                lines.append(f'    File 2: {snippet[2]}')
                lines.append(f'            {snippet[3]}')
                lines.append(f'    Token in File 1: {difference.token1}')
                lines.append(f'    Token in File 2: {difference.token2}')
            else:
                lines.append(f'  {uname} differs')
            
            print('\n'.join(lines))
            diff.file_diff_text.extend(lines)
            
            # If there is a diff, also create an image diff but only if we haven't already.
            if args.gen_diff_image and diff.png_diff_result is None:
                diff_image_file = to_png_file_name(uname)
                diff.add_file(diff_image_file)
                output_path = os.path.join(args.diff_output_dir, diff_image_file)
                file1_path = os.path.join(args.dir1, diff_image_file)
                file2_path = os.path.join(args.dir2, diff_image_file)
                change_detected = compare_images(file1_path, file2_path, output_path)
                diff.png_diff_result = change_detected
                if change_detected:
                    diff.set_png_diff(output_path)
                    print(f'    Diff image saved to {output_path}')
                else:
                    print('    No change detected in the images')

    if unique_to_dir1:
        print(f'\nFiles unique to {args.dir1} with suffixes {args.suffixes}:')
        for fname in unique_to_dir1:
            uname = fname.replace('\\', '/')
            diff: ModelDiffResult = diffResults.add_file_to_only_in_one_dir(0, uname)
            diff_image_file = to_png_file_name(uname)
            diff.add_file(diff_image_file)
            print(f'  {uname}')

    if unique_to_dir2:
        print(f'\nFiles unique to {args.dir2} with suffixes {args.suffixes}:')
        for fname in unique_to_dir2:
            uname = fname.replace('\\', '/')
            diff: ModelDiffResult = diffResults.add_file_to_only_in_one_dir(1, uname)
            diff_image_file = to_png_file_name(uname)
            diff.add_file(diff_image_file)
            print(f'  {uname}')

    if args.verbose:
        print(f'\nCommon files with the same contents in both directories with suffixes {args.suffixes}:')
        for uname in equal_files:
            diffResults.add_no_diff_file(uname)
            print(f'  {uname}')
            
    return unequal_files or len(unique_to_dir1) or len(unique_to_dir2), diffResults


@dataclass
class ServerConfig:
    dir1: str
    dir2: str
    diff_output_dir: str = '.'
    host: str = 'localhost'
    port: Optional[int] = None
    
    
def find_index_html():
    '''Find the index.html file in the same directory as this module.'''
    current_dir = os.path.dirname(__file__)
    index_path = os.path.join(current_dir, 'diff_viewer.html')
    if os.path.exists(index_path):
        return index_path
    return None

class MyHttpRequestHandler(http.server.SimpleHTTPRequestHandler):
    '''Custom HTTP request handler that serves files from the specified directories and 
    diff output directory.'''
    def __init__(self, *args, **kwargs):
        self.server_conf = kwargs.pop('server_conf', {})
        self.dir1 = self.server_conf.dir1
        self.dir2 = self.server_conf.dir2
        self.diff_output_dir = self.server_conf.diff_output_dir
        super().__init__(*args, **kwargs)

    def do_GET(self):
        url_path = unquote(self.path)
        
        if url_path == '/':
            self.send_response(200)
            self.send_header('Content-type', 'text/html')
            self.end_headers()
            index_html = find_index_html()
            if index_html is None:
                self.send_response(200)
                self.wfile.write(b"<!DOCTYPE html><html><body><h1>Index not found</h1></body></html>")
                return
            with open(index_html, 'rb') as file:
                self.send_response(200)
                self.wfile.write(file.read())
            return
        if url_path == '/diff.json':
            diff_json_path = os.path.join(self.diff_output_dir, "diff.json")
            if os.path.exists(diff_json_path):
                self.send_response(200)
                self.send_header('Content-type', 'application/json')
                self.end_headers()
                with open(diff_json_path, 'rb') as file:
                    self.wfile.write(file.read())
                return
            else:
                self.send_response(404)
                self.end_headers()
                return
        
        resource_path = url_path[1:] if url_path.startswith('/') else url_path
        if not os.path.exists(resource_path):
            resource_path = resource_file_url_match(self.server_conf, self.dir1, url_path)
            if resource_path is None:
                resource_path = resource_file_url_match(self.server_conf, self.dir2, url_path)
            if resource_path is None:
                resource_path = diff_file_url_match(self.server_conf, self.diff_output_dir, url_path)
            
        if resource_path and os.path.exists(resource_path):
            self.send_response(200)
            self.send_header('Content-type', mimetypes.guess_type(resource_path)[0] or 'application/octet-stream')
            self.end_headers()
            with open(resource_path, 'rb') as file:
                self.wfile.write(file.read())
        else:
            self.send_response(404)
            self.end_headers()
    
    def log_message(self, format, *args):
        # Suppress logging
        return

def resource_file_url_match(args, dir, url_path):
    potential_path = os.path.join(dir, url_path.lstrip('/'))
    return potential_path if os.path.exists(potential_path) else None

def diff_file_url_match(args, diff_output_dir, url_path):
    potential_path = os.path.join(diff_output_dir, url_path.lstrip('/'))
    return potential_path if os.path.exists(potential_path) else None

def start_server(server_conf: ServerConfig):
    handler = lambda *args, **kwargs: MyHttpRequestHandler(*args, server_conf=server_conf, **kwargs)
    
    host = server_conf.host
    port = server_conf.port

    if port is None:
        port = 4500
        while True:
            try:
                with socketserver.TCPServer((host, port), handler) as httpd:
                    print(f"Serving at: http://{host}:{port}")
                    httpd.serve_forever()
            except OSError:
                port += 1
                continue
            break
    else:
        with socketserver.TCPServer((host, port), handler) as httpd:
            print(f"Serving at: http://{host}:{port}")
            httpd.serve_forever()
    
def add_boolean_optional_argument(
    parser: argparse.ArgumentParser, name: str, dest: str=None, help: str=None, default: bool=False):
    if not dest:
        dest = name.replace('-', '_')
    parser.add_argument(
        f'--{name}', 
        dest=dest,
        action='store_true',
        help=help)
    parser.add_argument(
        f'--no-{name}', 
        dest=dest,
        action=argparse.BooleanOptionalAction,
        help=f"Do not {help[0].lower() + help[1:]}")
    
    parser.set_defaults(**{dest:default})

def main():
    parser = argparse.ArgumentParser(
        description='Compare contents of 2 directories for files with given suffix.')
    parser.add_argument('--dir1', type=str, default=None, help='First directory')
    parser.add_argument('--dir2', type=str, default=None, help='Second directory')
    parser.add_argument('--suffixes', type=str, nargs='+', help='List of file suffixes to compare (multiple suffixes can be provided)')
    parser.add_argument(
        '-v', '--verbose', action='store_true', 
        help='Verbose output including files with the same contents')
    parser.set_defaults(verbose=False)
    parser.add_argument(
        '--epsilon', type=float, default=1e-8, 
        help="Tolerance for floating-point differences.")
    
    add_boolean_optional_argument(
        parser,
        'gen-diff-image', 
        dest='gen_diff_image',
        help="Generate diff images for image files.")
    
    parser.add_argument(
        '--diff_output_dir', 
        type=str, 
        default=None,
        help="Whether to generate diff images for image files.")
    
    parser.add_argument(
        '--host', 
        type=str, 
        default="localhost",
        help="The local host address to start the server on.")
    
    parser.add_argument(
        '--port', 
        type=str, 
        default=None,
        help="The local host address to start the server on.")
        
    parser.add_argument(
        '--serve_diff_output_dir', 
        type=str, 
        default=None,
        help="If set, no comparison is performed, only the server is started.")
    
    add_boolean_optional_argument(
        parser,
        'start-server', 
        help="True if an http server is started once the diff is complete.",
        default=True)

    args = parser.parse_args()
    
    if args.serve_diff_output_dir:
        if args.diff_output_dir or args.dir1 or args.dir2 or args.suffixes:
            print('Error: --serve_diff_output_dir cannot be used with \n'
                  '--diff_output_dir, --dir1, --dir2, or --suffixes.', file=sys.stderr)
            return 1
        
        json_in_path = os.path.join(args.serve_diff_output_dir, "diff.json")
        with open(json_in_path, 'r') as f:
            json_data = f.read()
        diffResults: DirectoryDiffResults = DirectoryDiffResults.from_json(json_data)

        server_conf: ServerConfig = ServerConfig(
            diffResults.dirs[0], 
            diffResults.dirs[1], 
            args.serve_diff_output_dir, 
            args.host, 
            args.port)
        start_server(server_conf)
        return 0
    
    if not args.diff_output_dir:
        args.diff_output_dir = make_image_output_dir(args.dir1, args.dir2)

    print(f'Comparing files with suffixes {args.suffixes} between {args.dir1} and {args.dir2}...')
    has_diff, diffResults = compare_directories(args)
    
    if has_diff:
        diff_json = diffResults.to_json()
        output_path = os.path.join(args.diff_output_dir, "diff.json")
        with open(output_path, 'w') as f:
            f.write(diff_json)

        if args.start_server:
            start_server(args)
        return 0
    

    if args.start_server:
        print('Not starting server, no differences found. Nothing to serve.')
    return 0

if __name__ == '__main__':
    # args = ['--dir1', 'generated-dev-fobm-pmc', 
    #     '--dir2', 'generated-dev-fobm',
    #     '--suffixes', '.scad', '.png', 
    #     '--epsilon', '1e-8', 
    #     '--gen-diff-image']
    # args = ['--dir1', 'generated.tess1', 
    #         '--dir2', 'generated.master',
    #         '--suffixes', '.scad', 
    #         '--epsilon', '1e-8', 
    #         '--gen-diff-image']
    # args = ['--serve_diff_output_dir', 'generated.master-tess1',]
    #sys.argv = [sys.argv[0]] + args
    exit(main())