#!/usr/bin/env python3
"""
KRS PDF Parser - wyciąga dane zarządu i wspólników z odpisu KRS

Parsuje odpisy pełne pobrane z ekrs.ms.gov.pl i wyciąga:
- Członków zarządu (funkcja, imię, nazwisko, PESEL)
- Wspólników (imię, nazwisko, PESEL, udziały)
- Prokurentów

Usage:
    python scripts/parse_krs_pdf.py --file /path/to/odpis.pdf
    python scripts/parse_krs_pdf.py --dir /path/to/pdfs/
"""

import re
import json
import argparse
from pathlib import Path
from dataclasses import dataclass, asdict
from typing import List, Optional, Dict, Any

try:
    import pdfplumber
except ImportError:
    print("Wymagana biblioteka pdfplumber. Zainstaluj: pip install pdfplumber")
    exit(1)


@dataclass
class Person:
    """Osoba powiązana z firmą"""
    nazwisko: str
    imiona: str
    pesel: Optional[str] = None
    rola: str = ""  # PREZES ZARZĄDU, CZŁONEK ZARZĄDU, WSPÓLNIK, PROKURENT
    udzialy: Optional[str] = None  # dla wspólników

    def full_name(self) -> str:
        return f"{self.imiona} {self.nazwisko}"


@dataclass
class KRSData:
    """Dane wyciągnięte z odpisu KRS"""
    krs: str
    nazwa: str
    nip: Optional[str] = None
    regon: Optional[str] = None
    zarzad: List[Person] = None
    wspolnicy: List[Person] = None
    prokurenci: List[Person] = None
    zrodlo: str = "ekrs.ms.gov.pl"

    def __post_init__(self):
        if self.zarzad is None:
            self.zarzad = []
        if self.wspolnicy is None:
            self.wspolnicy = []
        if self.prokurenci is None:
            self.prokurenci = []

    def to_dict(self) -> Dict[str, Any]:
        return {
            'krs': self.krs,
            'nazwa': self.nazwa,
            'nip': self.nip,
            'regon': self.regon,
            'zarzad': [asdict(p) for p in self.zarzad],
            'wspolnicy': [asdict(p) for p in self.wspolnicy],
            'prokurenci': [asdict(p) for p in self.prokurenci],
            'zrodlo': self.zrodlo
        }


def extract_text_from_pdf(pdf_path: str) -> str:
    """Wyciąga tekst z PDF"""
    with pdfplumber.open(pdf_path) as pdf:
        text = ""
        for page in pdf.pages:
            page_text = page.extract_text()
            if page_text:
                text += page_text + "\n"
    return text


def parse_person_block(lines: List[str], start_idx: int) -> Optional[Person]:
    """
    Parsuje blok danych osoby z linii PDF

    Format w PDF:
    1.Nazwisko / Nazwa lub firma 1 - NAZWISKO
    2.Imiona 1 - IMIĘ DRUGIE_IMIĘ
    3.Numer PESEL/REGON lub data 1 - 12345678901, ------
    """
    person = Person(nazwisko="", imiona="")
    found_nazwisko = False
    found_imiona = False
    found_pesel = False

    for i in range(start_idx, min(start_idx + 10, len(lines))):
        line = lines[i].strip()

        # Wykryj początek następnej osoby - przestań parsować
        if i > start_idx and ('1.Nazwisko' in line or 'Nazwisko / Nazwa' in line):
            # Początek nowej osoby - koniec bloku
            break

        # Nazwisko (tylko pierwsze znalezione)
        if not found_nazwisko and 'Nazwisko' in line and ' - ' in line:
            match = re.search(r' - ([A-ZĄĆĘŁŃÓŚŹŻ\-]+)$', line)
            if match:
                person.nazwisko = match.group(1)
                found_nazwisko = True

        # Imiona (tylko pierwsze znalezione)
        if not found_imiona and 'Imiona' in line and ' - ' in line:
            match = re.search(r' - ([A-ZĄĆĘŁŃÓŚŹŻ ]+)$', line)
            if match:
                person.imiona = match.group(1).strip()
                found_imiona = True

        # PESEL (tylko pierwsze znalezione)
        if not found_pesel and 'PESEL' in line and ' - ' in line:
            match = re.search(r' - (\d{11})', line)
            if match:
                person.pesel = match.group(1)
                found_pesel = True

        # Funkcja (dla zarządu)
        if 'Funkcja' in line and ' - ' in line:
            match = re.search(r' - ([A-ZĄĆĘŁŃÓŚŹŻ ]+)$', line)
            if match:
                person.rola = match.group(1).strip()

    if person.nazwisko and person.imiona:
        return person
    return None


def parse_krs_pdf(pdf_path: str) -> KRSData:
    """
    Parsuje odpis KRS i wyciąga dane
    """
    text = extract_text_from_pdf(pdf_path)
    lines = text.split('\n')

    # Extract basic info
    krs_match = re.search(r'Numer KRS:\s*(\d{10})', text)
    krs = krs_match.group(1) if krs_match else ""

    # Find company name - format: "3.Firma, pod którą spółka działa 1 - NAZWA FIRMY"
    nazwa = ""
    nazwa_match = re.search(r'3\.Firma,?\s+pod którą spółka działa\s+\d+\s+-\s+([^\n]+)', text)
    if nazwa_match:
        nazwa = nazwa_match.group(1).strip()

    # NIP and REGON - format: "REGON: 369796786, NIP: 5862329746"
    nip_match = re.search(r'NIP:\s*(\d{10})', text)
    nip = nip_match.group(1) if nip_match else None

    regon_match = re.search(r'REGON:\s*(\d{9,14})', text)
    regon = regon_match.group(1) if regon_match else None

    data = KRSData(krs=krs, nazwa=nazwa, nip=nip, regon=regon)

    # Parse sections
    in_zarzad = False
    in_wspolnicy = False
    in_prokurenci = False

    for i, line in enumerate(lines):
        line_stripped = line.strip()

        # Detect sections
        if 'ZARZĄD' in line_stripped.upper() and 'Nazwa organu' in line_stripped:
            in_zarzad = True
            in_wspolnicy = False
            in_prokurenci = False
            continue

        # Wspólnicy - szukaj "Dane wspólników" żeby nie łapać "Wspólnik może mieć:"
        if 'Dane wspólników' in line_stripped or 'WSPÓLNICY' in line_stripped.upper():
            in_wspolnicy = True
            in_zarzad = False
            in_prokurenci = False

        if 'PROKURENCI' in line_stripped.upper() or 'Prokurent' in line_stripped:
            in_prokurenci = True
            in_zarzad = False
            in_wspolnicy = False

        # Parse person data when we find "Nazwisko"
        if '1.Nazwisko' in line_stripped or 'Nazwisko / Nazwa' in line_stripped:
            person = parse_person_block(lines, i)
            if person:
                if in_zarzad:
                    # Look for function in nearby lines
                    for j in range(i, min(i + 8, len(lines))):
                        if 'Funkcja' in lines[j] and ' - ' in lines[j]:
                            func_match = re.search(r' - ([A-ZĄĆĘŁŃÓŚŹŻ ]+)$', lines[j])
                            if func_match:
                                person.rola = func_match.group(1).strip()
                            break
                    if not person.rola:
                        person.rola = "CZŁONEK ZARZĄDU"
                    data.zarzad.append(person)
                elif in_wspolnicy:
                    person.rola = "WSPÓLNIK"
                    # Look for share info
                    for j in range(i, min(i + 10, len(lines))):
                        if 'udziałów' in lines[j].lower() or 'udział' in lines[j].lower():
                            data.wspolnicy.append(person)
                            break
                    else:
                        data.wspolnicy.append(person)
                elif in_prokurenci:
                    person.rola = "PROKURENT"
                    data.prokurenci.append(person)

    return data


def main():
    parser = argparse.ArgumentParser(description="Parse KRS PDF files")
    parser.add_argument("--file", type=str, help="Single PDF file to parse")
    parser.add_argument("--dir", type=str, help="Directory with PDF files")
    parser.add_argument("--output", type=str, help="Output JSON file")
    args = parser.parse_args()

    results = []

    if args.file:
        print(f"Parsing: {args.file}")
        data = parse_krs_pdf(args.file)
        results.append(data.to_dict())

        # Print summary
        print(f"\n=== {data.nazwa} (KRS: {data.krs}) ===")
        print(f"NIP: {data.nip}, REGON: {data.regon}")

        print(f"\nZarząd ({len(data.zarzad)} osób):")
        for p in data.zarzad:
            print(f"  - {p.full_name()} - {p.rola}")

        print(f"\nWspólnicy ({len(data.wspolnicy)} osób):")
        for p in data.wspolnicy:
            print(f"  - {p.full_name()}")

        if data.prokurenci:
            print(f"\nProkurenci ({len(data.prokurenci)} osób):")
            for p in data.prokurenci:
                print(f"  - {p.full_name()}")

    elif args.dir:
        pdf_dir = Path(args.dir)
        pdf_files = list(pdf_dir.glob("*.pdf"))
        print(f"Found {len(pdf_files)} PDF files")

        for pdf_file in pdf_files:
            print(f"Parsing: {pdf_file.name}...")
            try:
                data = parse_krs_pdf(str(pdf_file))
                results.append(data.to_dict())
                print(f"  OK: {data.nazwa}")
            except Exception as e:
                print(f"  ERROR: {e}")

    # Save results
    if args.output and results:
        with open(args.output, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        print(f"\nResults saved to: {args.output}")
    elif results:
        print("\n=== JSON OUTPUT ===")
        print(json.dumps(results, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()