#!/usr/bin/env python3
"""
Poboljšana skripta za analizu tehničkih nacrta
Otkriva i kategorizira dimenzije, navoje, hrapavost i GD&T oznake

Autor: Claude
Verzija: 1.0
"""

import pdfplumber
import re
import json
import os
import sys
import argparse
from openpyxl import Workbook
from datetime import datetime
import fitz  # PyMuPDF
import tempfile
import numpy as np
from PIL import Image
import pytesseract

# Definiranje kategorija i potkategorija
CATEGORIES = {
    'metadata': {
        'name': 'Meta data',
        'subtypes': ['identification', 'measurements', 'material', 'bom']
    },
    'features': {
        'name': 'Features',
        'subtypes': ['dimensions', 'threads', 'roughnesses', 'gdts']
    },
    'type': {
        'name': 'Type',
        'subtypes': ['company_data', 'logo']
    }
}

# Regularni izrazi za prepoznavanje različitih formata dimenzija
PATTERNS = {
    # Promjeri
    'diameter': r'Ø\s*\d+(\.\d+)?(\s*[a-zA-Z]\d+)?(\s*\(\+?[\d\.-]+\/?[+-]?[\d\.-]+\))?',
    
    # Navoji
    'thread': r'M\d+(\.\d+)?\s*[×x]\s*\d+(\.\d+)?(?:---\d+[a-zA-Z]\/\d+[a-zA-Z])?',
    
    # Hrapavosti
    'roughness': r'▽√?U?Ra\s*\d+(\.\d+)?',
    
    # GD&T oznake
    'gdt': r'\[\s*[↗⟂◎⌭]\s*\|\s*[Ø]?\d+(\.\d+)?\s*\|\s*[A-Z]-[A-Z]\s*\]|\[\s*[↗⟂◎⌭]\s*\|\s*[Ø]?\d+(\.\d+)?\s*\]',
    
    # Linearne dimenzije
    'linear': r'\d+(\.\d+)?(\s*[±]\s*\d+(\.\d+)?)?',
    
    # Radijusi
    'radius': r'R\d+(\.\d+)?',
    
    # Kutevi
    'angle': r'\d+(\.\d+)?°'
}

def analyze_pdf_dimensions(pdf_path, output_pdf, output_excel, manual_mode=False, dimensions_list=None):
    """
    Analizira PDF dokument i detektira dimenzije
    
    Args:
        pdf_path: Putanja do PDF datoteke
        output_pdf: Putanja za označeni PDF
        output_excel: Putanja za Excel datoteku
        manual_mode: Uključuje ručni način označavanja
        dimensions_list: JSON string s dimenzijama za ručni način
        
    Returns:
        Lista detektiranih dimenzija
    """
    dimensions = []
    dimension_count = 0

    try:
        # Ako je ručni režim, koristi prosleđene dimenzije
        if manual_mode and dimensions_list:
            dimensions = json.loads(dimensions_list)
            for i, dim in enumerate(dimensions, 1):
                dimensions[i-1]['number'] = i
        else:
            # Automatsko prepoznavanje dimenzija
            with pdfplumber.open(pdf_path) as pdf:
                for page_number, page in enumerate(pdf.pages, start=1):
                    text = page.extract_text()

                    if text:
                        # Pronalaženje svih dimenzija na stranici prema uzorcima
                        for pattern_name, pattern in PATTERNS.items():
                            matches = re.finditer(pattern, text, re.IGNORECASE)
                            for match in matches:
                                dimension_text = match.group(0).strip()

                                # Provjera je li već dodana ova dimenzija
                                if not any(d["text"] == dimension_text for d in dimensions):
                                    # Pronalaženje pozicije teksta na stranici
                                    words = page.extract_words()
                                    for word in words:
                                        if dimension_text in word['text']:
                                            dimension_count += 1
                                            
                                            # Kategoriziranje dimenzije
                                            dimension_type, subtype = categorize_dimension(dimension_text, pattern_name)
                                            
                                            dimensions.append({
                                                "text": dimension_text,
                                                "page": page_number,
                                                "x0": word['x0'],
                                                "y0": word['top'],
                                                "x1": word['x1'],
                                                "y1": word['bottom'],
                                                "number": dimension_count,
                                                "type": dimension_type,
                                                "subtype": subtype
                                            })
                                            break

        # Dodavanje numerisanih crvenih krugova na PDF
        doc = fitz.open(pdf_path)

        for dim in dimensions:
            page_idx = dim['page'] - 1
            if page_idx < len(doc):  # Provjeri da li stranica postoji
                page = doc[page_idx]

                # Izračunaj poziciju za crveni krug
                if manual_mode and 'x' in dim and 'y' in dim:
                    circle_x = dim['x']
                    circle_y = dim['y']
                else:
                    circle_x = dim['x0'] - 20
                    circle_y = (dim['y0'] + dim['y1']) / 2

                # Dodaj crveni krug
                page.draw_circle(fitz.Point(circle_x, circle_y), 10, color=(1, 0, 0), fill=(1, 0, 0))

                # Dodaj broj u krug
                text_point = fitz.Point(circle_x - 3 if dim['number'] < 10 else circle_x - 6, circle_y + 3)
                page.insert_text(text_point, str(dim['number']), fontsize=12, color=(1, 1, 1))

        # Sačuvaj označeni PDF
        doc.save(output_pdf)
        doc.close()

        # Kreiranje Excel datoteke s dimenzijama
        create_excel_report(dimensions, output_excel)

        return dimensions

    except Exception as e:
        raise Exception(f"Error processing PDF: {str(e)}")

def categorize_dimension(text, pattern_name):
    """
    Kategorizira dimenziju prema njenom tekstu i uzorku
    
    Args:
        text: Tekst dimenzije
        pattern_name: Ime uzorka koji je prepoznao dimenziju
        
    Returns:
        Tuple (tip dimenzije, podtip)
    """
    # Zadane vrijednosti
    dimension_type = "features"
    subtype = "dimensions"
    
    # Određivanje kategorije i podkategorije prema uzorku
    if pattern_name == 'thread':
        subtype = "threads"
    elif pattern_name == 'roughness':
        subtype = "roughnesses"
    elif pattern_name == 'gdt':
        subtype = "gdts"
    elif pattern_name == 'diameter':
        subtype = "dimensions"  # Ostaje dimenzija, ali možemo dodati specifičniju informaciju
    
    return dimension_type, subtype

def create_excel_report(dimensions, output_excel):
    """
    Kreira Excel izvještaj s dimenzijama kategoriziranim po listovima
    
    Args:
        dimensions: Lista dimenzija
        output_excel: Putanja za Excel datoteku
    """
    wb = Workbook()
    
    # Ukloni zadani sheet
    default_sheet = wb.active
    wb.remove(default_sheet)
    
    # Grupiraj dimenzije po tipovima
    dimension_types = {}
    for dim in dimensions:
        dim_type = dim.get('type', 'features')
        subtype = dim.get('subtype', 'dimensions')
        
        key = f"{dim_type}_{subtype}"
        if key not in dimension_types:
            dimension_types[key] = []
        
        dimension_types[key].append(dim)
    
    # Kreiraj sheet za svaki tip
    for key, dims in dimension_types.items():
        if not dims:
            continue
            
        sheet_name = key.replace('_', ' ').title()
        if len(sheet_name) > 31:  # Excel ograničenje za ime sheet-a
            sheet_name = sheet_name[:31]
            
        ws = wb.create_sheet(sheet_name)
        
        # Dodaj zaglavlja
        ws.append(["Number", "Dimension", "Page"])
        
        # Dodaj dimenzije
        for dim in dims:
            ws.append([dim["number"], dim["text"], dim["page"]])
        
        # Podesi širinu kolona
        for column in ws.columns:
            max_length = 0
            column = list(column)
            for cell in column:
                try:
                    if len(str(cell.value)) > max_length:
                        max_length = len(str(cell.value))
                except:
                    pass
            adjusted_width = (max_length + 2)
            ws.column_dimensions[column[0].column_letter].width = adjusted_width
    
    # Kreiraj sheet sa svim dimenzijama
    all_dims_sheet = wb.create_sheet("All Dimensions")
    all_dims_sheet.append(["Number", "Dimension", "Type", "Subtype", "Page"])
    
    for dim in dimensions:
        all_dims_sheet.append([
            dim["number"], 
            dim["text"], 
            dim.get("type", "features"),
            dim.get("subtype", "dimensions"),
            dim["page"]
        ])
    
    # Podesi širinu kolona
    for column in all_dims_sheet.columns:
        max_length = 0
        column = list(column)
        for cell in column:
            try:
                if len(str(cell.value)) > max_length:
                    max_length = len(str(cell.value))
            except:
                pass
        adjusted_width = (max_length + 2)
        all_dims_sheet.column_dimensions[column[0].column_letter].width = adjusted_width
    
    # Sačuvaj Excel
    wb.save(output_excel)

def process_image(image_path):
    """
    Obrađuje sliku za prepoznavanje teksta
    
    Args:
        image_path: Putanja do slike
        
    Returns:
        Tekst prepoznat na slici
    """
    # Učitaj sliku
    image = Image.open(image_path)
    
    # Pretprocesiranje slike za bolji OCR
    # Konverzija u sivu skalu
    gray_image = image.convert('L')
    
    # Binarizacija (threshold) za bolji kontrast
    threshold = 200
    binary_image = gray_image.point(lambda p: p > threshold and 255)
    
    # OCR - prepoznavanje teksta
    text = pytesseract.image_to_string(binary_image, config='--psm 6')
    
    return text

def analyze_image_dimensions(image_path, output_image, output_excel):
    """
    Analizira sliku i detektira dimenzije koristeći OCR
    
    Args:
        image_path: Putanja do slike
        output_image: Putanja za označenu sliku
        output_excel: Putanja za Excel datoteku
        
    Returns:
        Lista detektiranih dimenzija
    """
    dimensions = []
    dimension_count = 0
    
    try:
        # Prepoznavanje teksta na slici
        text = process_image(image_path)
        
        if text:
            # Pronalaženje svih dimenzija prema uzorcima
            for pattern_name, pattern in PATTERNS.items():
                matches = re.finditer(pattern, text, re.IGNORECASE)
                for match in matches:
                    dimension_text = match.group(0).strip()
                    
                    # Provjera je li već dodana ova dimenzija
                    if not any(d["text"] == dimension_text for d in dimensions):
                        dimension_count += 1
                        
                        # Kategoriziranje dimenzije
                        dimension_type, subtype = categorize_dimension(dimension_text, pattern_name)
                        
                        # Postavljamo dummy poziciju jer ne možemo lako odrediti poziciju na slici
                        dimensions.append({
                            "text": dimension_text,
                            "page": 1,
                            "x0": 100,  # dummy vrijednost
                            "y0": 100 + (dimension_count * 30),  # dummy vrijednost
                            "x1": 200,  # dummy vrijednost
                            "y1": 120 + (dimension_count * 30),  # dummy vrijednost
                            "number": dimension_count,
                            "type": dimension_type,
                            "subtype": subtype
                        })
        
        # Stvaranje označene slike
        # Učitaj originalnu sliku
        original_image = Image.open(image_path)
        
        # Konverzija u RGB ako nije
        if original_image.mode != 'RGB':
            original_image = original_image.convert('RGB')
        
        # Za renderiranje krugova koristimo numpy
        img_array = np.array(original_image)
        
        # Za svaku dimenziju, dodaj crveni krug
        for dim in dimensions:
            x, y = dim['x0'], dim['y0']
            
            # Jednostavan crveni krug (u praksi bi trebalo koristiti OpenCV za ovo)
            # Ovo je pojednostavljeno za demonstraciju
            for i in range(-10, 11):
                for j in range(-10, 11):
                    if i*i + j*j <= 100:  # Krug radijusa 10
                        try:
                            img_array[int(y + j), int(x + i)] = [255, 0, 0]  # RGB za crvenu
                        except IndexError:
                            pass
        
        # Konverzija natrag u PIL sliku
        marked_image = Image.fromarray(img_array)
        
        # Spremi označenu sliku
        marked_image.save(output_image)
        
        # Kreiranje Excel datoteke s dimenzijama
        create_excel_report(dimensions, output_excel)
        
        return dimensions
        
    except Exception as e:
        raise Exception(f"Error processing image: {str(e)}")

def main():
    """
    Glavna funkcija za analizu tehničkih nacrta
    """
    parser = argparse.ArgumentParser(description='Analyze technical drawing, mark dimensions and export to Excel')
    parser.add_argument('--file_path', required=True, help='Path to the input file (PDF, JPEG, PNG, etc.)')
    parser.add_argument('--output_pdf', required=False, help='Path for the output marked PDF file')
    parser.add_argument('--output_image', required=False, help='Path for the output marked image file')
    parser.add_argument('--output_excel', required=True, help='Path for the output Excel file')
    parser.add_argument('--manual_mode', action='store_true', help='Use manual mode with provided dimensions')
    parser.add_argument('--dimensions', help='JSON string with dimensions (required for manual mode)')

    args = parser.parse_args()

    try:
        if not os.path.exists(args.file_path):
            print(json.dumps({"success": False, "error": "File not found"}))
            sys.exit(1)

        if args.manual_mode and not args.dimensions:
            print(json.dumps({"success": False, "error": "Dimensions are required for manual mode"}))
            sys.exit(1)
            
        # Određivanje tipa datoteke
        file_extension = os.path.splitext(args.file_path)[1].lower()
        
        if file_extension == '.pdf':
            # Za PDF koristimo analizu PDF-a
            output_pdf = args.output_pdf or args.file_path.replace('.pdf', '_marked.pdf')
            output_excel = args.output_excel
            
            results = analyze_pdf_dimensions(
                args.file_path,
                output_pdf,
                output_excel,
                args.manual_mode,
                args.dimensions
            )
            
            print(json.dumps({
                "success": True,
                "data": results,
                "marked_file": output_pdf,
                "excel_file": output_excel,
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            }))
            
        elif file_extension in ['.jpg', '.jpeg', '.png', '.tiff', '.bmp']:
            # Za slike koristimo OCR analizu
            output_image = args.output_image or args.file_path.replace(file_extension, '_marked' + file_extension)
            output_excel = args.output_excel
            
            results = analyze_image_dimensions(
                args.file_path,
                output_image,
                output_excel
            )
            
            print(json.dumps({
                "success": True,
                "data": results,
                "marked_file": output_image,
                "excel_file": output_excel,
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            }))
            
        else:
            print(json.dumps({
                "success": False,
                "error": f"Unsupported file format: {file_extension}",
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            }))

    except Exception as e:
        print(json.dumps({
            "success": False,
            "error": str(e),
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }))

if __name__ == "__main__":
    main()