"""
Database utilities for autolabel workflow.
"""

import json
import psycopg2
from psycopg2.extras import execute_values
from typing import Set, Dict, Any, Optional
import sys
from pathlib import Path

from shared.config import get_db_connection_string


class DocumentDB:
    """Database interface for document processing status."""

    def __init__(self, connection_string: str = None):
        self.connection_string = connection_string or get_db_connection_string()
        self.conn = None

    def connect(self):
        """Connect to database."""
        if self.conn is None:
            self.conn = psycopg2.connect(self.connection_string)
        return self.conn

    def create_tables(self):
        """Create database tables if they don't exist."""
        conn = self.connect()
        with conn.cursor() as cursor:
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS documents (
                    document_id TEXT PRIMARY KEY,
                    pdf_path TEXT,
                    pdf_type TEXT,
                    success BOOLEAN,
                    total_pages INTEGER,
                    fields_matched INTEGER,
                    fields_total INTEGER,
                    annotations_generated INTEGER,
                    processing_time_ms REAL,
                    timestamp TIMESTAMPTZ,
                    errors JSONB DEFAULT '[]',
                    -- Extended CSV format fields
                    split TEXT,
                    customer_number TEXT,
                    supplier_name TEXT,
                    supplier_organisation_number TEXT,
                    supplier_accounts TEXT
                );

                CREATE TABLE IF NOT EXISTS field_results (
                    id SERIAL PRIMARY KEY,
                    document_id TEXT NOT NULL REFERENCES documents(document_id) ON DELETE CASCADE,
                    field_name TEXT,
                    csv_value TEXT,
                    matched BOOLEAN,
                    score REAL,
                    matched_text TEXT,
                    candidate_used TEXT,
                    bbox JSONB,
                    page_no INTEGER,
                    context_keywords JSONB DEFAULT '[]',
                    error TEXT
                );

                CREATE INDEX IF NOT EXISTS idx_documents_success ON documents(success);
                CREATE INDEX IF NOT EXISTS idx_field_results_document_id ON field_results(document_id);
                CREATE INDEX IF NOT EXISTS idx_field_results_field_name ON field_results(field_name);
                CREATE INDEX IF NOT EXISTS idx_field_results_matched ON field_results(matched);

                -- Add new columns to existing tables if they don't exist (for migration)
                DO $$
                BEGIN
                    IF NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name='documents' AND column_name='split') THEN
                        ALTER TABLE documents ADD COLUMN split TEXT;
                    END IF;
                    IF NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name='documents' AND column_name='customer_number') THEN
                        ALTER TABLE documents ADD COLUMN customer_number TEXT;
                    END IF;
                    IF NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name='documents' AND column_name='supplier_name') THEN
                        ALTER TABLE documents ADD COLUMN supplier_name TEXT;
                    END IF;
                    IF NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name='documents' AND column_name='supplier_organisation_number') THEN
                        ALTER TABLE documents ADD COLUMN supplier_organisation_number TEXT;
                    END IF;
                    IF NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name='documents' AND column_name='supplier_accounts') THEN
                        ALTER TABLE documents ADD COLUMN supplier_accounts TEXT;
                    END IF;
                END $$;
            """)
        conn.commit()

    def close(self):
        """Close database connection."""
        if self.conn:
            self.conn.close()
            self.conn = None

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def get_successful_doc_ids(self) -> Set[str]:
        """Get all document IDs that have been successfully processed."""
        conn = self.connect()
        with conn.cursor() as cursor:
            cursor.execute("SELECT document_id FROM documents WHERE success = true")
            return {row[0] for row in cursor.fetchall()}

    def get_failed_doc_ids(self) -> Set[str]:
        """Get all document IDs that failed processing."""
        conn = self.connect()
        with conn.cursor() as cursor:
            cursor.execute("SELECT document_id FROM documents WHERE success = false")
            return {row[0] for row in cursor.fetchall()}

    def check_document_status(self, doc_id: str) -> Optional[bool]:
        """
        Check if a document exists and its success status.

        Returns:
            True if exists and success=true
            False if exists and success=false
            None if not exists
        """
        conn = self.connect()
        with conn.cursor() as cursor:
            cursor.execute(
                "SELECT success FROM documents WHERE document_id = %s",
                (doc_id,)
            )
            row = cursor.fetchone()
            if row is None:
                return None
            return row[0]

    def check_documents_status_batch(self, doc_ids: list[str]) -> Dict[str, Optional[bool]]:
        """
        Batch check document status for multiple IDs.

        Returns:
            Dict mapping doc_id to status:
                True if exists and success=true
                False if exists and success=false
                (missing from dict if not exists)
        """
        if not doc_ids:
            return {}

        conn = self.connect()
        with conn.cursor() as cursor:
            cursor.execute(
                "SELECT document_id, success FROM documents WHERE document_id = ANY(%s)",
                (doc_ids,)
            )
            return {row[0]: row[1] for row in cursor.fetchall()}

    def delete_document(self, doc_id: str):
        """Delete a document and its field results (for re-processing)."""
        conn = self.connect()
        with conn.cursor() as cursor:
            # field_results will be cascade deleted
            cursor.execute("DELETE FROM documents WHERE document_id = %s", (doc_id,))
        conn.commit()

    def get_document(self, doc_id: str) -> Optional[Dict[str, Any]]:
        """Get a single document with its field results."""
        conn = self.connect()
        with conn.cursor() as cursor:
            # Get document
            cursor.execute("""
                SELECT document_id, pdf_path, pdf_type, success, total_pages,
                       fields_matched, fields_total, annotations_generated,
                       processing_time_ms, timestamp, errors,
                       split, customer_number, supplier_name,
                       supplier_organisation_number, supplier_accounts
                FROM documents WHERE document_id = %s
            """, (doc_id,))
            row = cursor.fetchone()
            if not row:
                return None

            doc = {
                'document_id': row[0],
                'pdf_path': row[1],
                'pdf_type': row[2],
                'success': row[3],
                'total_pages': row[4],
                'fields_matched': row[5],
                'fields_total': row[6],
                'annotations_generated': row[7],
                'processing_time_ms': row[8],
                'timestamp': str(row[9]) if row[9] else None,
                'errors': row[10] if isinstance(row[10], list) else json.loads(row[10] or '[]'),
                # New fields
                'split': row[11],
                'customer_number': row[12],
                'supplier_name': row[13],
                'supplier_organisation_number': row[14],
                'supplier_accounts': row[15],
                'field_results': []
            }

            # Get field results
            cursor.execute("""
                SELECT field_name, csv_value, matched, score, matched_text,
                       candidate_used, bbox, page_no, context_keywords, error
                FROM field_results WHERE document_id = %s
            """, (doc_id,))

            for fr in cursor.fetchall():
                doc['field_results'].append({
                    'field_name': fr[0],
                    'csv_value': fr[1],
                    'matched': fr[2],
                    'score': fr[3],
                    'matched_text': fr[4],
                    'candidate_used': fr[5],
                    'bbox': fr[6] if isinstance(fr[6], list) else json.loads(fr[6]) if fr[6] else None,
                    'page_no': fr[7],
                    'context_keywords': fr[8] if isinstance(fr[8], list) else json.loads(fr[8] or '[]'),
                    'error': fr[9]
                })

            return doc

    def get_all_documents_summary(self, success_only: bool = False, limit: int = None) -> list[Dict[str, Any]]:
        """Get summary of all documents (without field_results for efficiency)."""
        conn = self.connect()
        with conn.cursor() as cursor:
            query = """
                SELECT document_id, pdf_path, pdf_type, success, total_pages,
                       fields_matched, fields_total
                FROM documents
            """
            params = []
            if success_only:
                query += " WHERE success = true"
            query += " ORDER BY timestamp DESC"
            if limit:
                # Use parameterized query instead of f-string
                query += " LIMIT %s"
                params.append(limit)

            cursor.execute(query, params if params else None)
            return [
                {
                    'document_id': row[0],
                    'pdf_path': row[1],
                    'pdf_type': row[2],
                    'success': row[3],
                    'total_pages': row[4],
                    'fields_matched': row[5],
                    'fields_total': row[6]
                }
                for row in cursor.fetchall()
            ]

    def get_field_stats(self) -> Dict[str, Dict[str, int]]:
        """Get match statistics per field."""
        conn = self.connect()
        with conn.cursor() as cursor:
            cursor.execute("""
                SELECT field_name,
                       COUNT(*) as total,
                       SUM(CASE WHEN matched THEN 1 ELSE 0 END) as matched
                FROM field_results
                GROUP BY field_name
                ORDER BY field_name
            """)
            return {
                row[0]: {'total': row[1], 'matched': row[2]}
                for row in cursor.fetchall()
            }

    def get_failed_matches(self, field_name: str = None, limit: int = 100) -> list[Dict[str, Any]]:
        """Get field results that failed to match."""
        conn = self.connect()
        with conn.cursor() as cursor:
            query = """
                SELECT fr.document_id, fr.field_name, fr.csv_value, fr.error,
                       d.pdf_type
                FROM field_results fr
                JOIN documents d ON fr.document_id = d.document_id
                WHERE fr.matched = false
            """
            params = []
            if field_name:
                query += " AND fr.field_name = %s"
                params.append(field_name)
            # Use parameterized query instead of f-string
            query += " LIMIT %s"
            params.append(limit)

            cursor.execute(query, params)
            return [
                {
                    'document_id': row[0],
                    'field_name': row[1],
                    'csv_value': row[2],
                    'error': row[3],
                    'pdf_type': row[4]
                }
                for row in cursor.fetchall()
            ]

    def get_documents_batch(self, doc_ids: list[str]) -> Dict[str, Dict[str, Any]]:
        """
        Get multiple documents with their field results in a single batch query.

        This is much more efficient than calling get_document() in a loop.

        Args:
            doc_ids: List of document IDs to fetch

        Returns:
            Dict mapping doc_id to document data (with field_results)
        """
        if not doc_ids:
            return {}

        conn = self.connect()
        result: Dict[str, Dict[str, Any]] = {}

        with conn.cursor() as cursor:
            # Batch fetch all documents
            cursor.execute("""
                SELECT document_id, pdf_path, pdf_type, success, total_pages,
                       fields_matched, fields_total, annotations_generated,
                       processing_time_ms, timestamp, errors,
                       split, customer_number, supplier_name,
                       supplier_organisation_number, supplier_accounts
                FROM documents WHERE document_id = ANY(%s)
            """, (doc_ids,))

            for row in cursor.fetchall():
                result[row[0]] = {
                    'document_id': row[0],
                    'pdf_path': row[1],
                    'pdf_type': row[2],
                    'success': row[3],
                    'total_pages': row[4],
                    'fields_matched': row[5],
                    'fields_total': row[6],
                    'annotations_generated': row[7],
                    'processing_time_ms': row[8],
                    'timestamp': str(row[9]) if row[9] else None,
                    'errors': row[10] if isinstance(row[10], list) else json.loads(row[10] or '[]'),
                    # New fields
                    'split': row[11],
                    'customer_number': row[12],
                    'supplier_name': row[13],
                    'supplier_organisation_number': row[14],
                    'supplier_accounts': row[15],
                    'field_results': []
                }

            if not result:
                return {}

            # Batch fetch all field results for these documents
            cursor.execute("""
                SELECT document_id, field_name, csv_value, matched, score,
                       matched_text, candidate_used, bbox, page_no, context_keywords, error
                FROM field_results WHERE document_id = ANY(%s)
            """, (list(result.keys()),))

            for fr in cursor.fetchall():
                doc_id = fr[0]
                if doc_id in result:
                    result[doc_id]['field_results'].append({
                        'field_name': fr[1],
                        'csv_value': fr[2],
                        'matched': fr[3],
                        'score': fr[4],
                        'matched_text': fr[5],
                        'candidate_used': fr[6],
                        'bbox': fr[7] if isinstance(fr[7], list) else json.loads(fr[7]) if fr[7] else None,
                        'page_no': fr[8],
                        'context_keywords': fr[9] if isinstance(fr[9], list) else json.loads(fr[9] or '[]'),
                        'error': fr[10]
                    })

        return result

    def save_document(self, report: Dict[str, Any]):
        """Save or update a document and its field results using batch operations."""
        conn = self.connect()
        doc_id = report.get('document_id')

        with conn.cursor() as cursor:
            # Delete existing record if any (for update)
            cursor.execute("DELETE FROM documents WHERE document_id = %s", (doc_id,))

            # Insert document
            cursor.execute("""
                INSERT INTO documents
                (document_id, pdf_path, pdf_type, success, total_pages,
                 fields_matched, fields_total, annotations_generated,
                 processing_time_ms, timestamp, errors,
                 split, customer_number, supplier_name, supplier_organisation_number, supplier_accounts)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """, (
                doc_id,
                report.get('pdf_path'),
                report.get('pdf_type'),
                report.get('success'),
                report.get('total_pages'),
                report.get('fields_matched'),
                report.get('fields_total'),
                report.get('annotations_generated'),
                report.get('processing_time_ms'),
                report.get('timestamp'),
                json.dumps(report.get('errors', [])),
                # New fields
                report.get('split'),
                report.get('customer_number'),
                report.get('supplier_name'),
                report.get('supplier_organisation_number'),
                report.get('supplier_accounts'),
            ))

            # Batch insert field results using execute_values
            field_results = report.get('field_results', [])
            if field_results:
                field_values = [
                    (
                        doc_id,
                        field.get('field_name'),
                        field.get('csv_value'),
                        field.get('matched'),
                        field.get('score'),
                        field.get('matched_text'),
                        field.get('candidate_used'),
                        json.dumps(field.get('bbox')) if field.get('bbox') else None,
                        field.get('page_no'),
                        json.dumps(field.get('context_keywords', [])),
                        field.get('error')
                    )
                    for field in field_results
                ]
                execute_values(cursor, """
                    INSERT INTO field_results
                    (document_id, field_name, csv_value, matched, score,
                     matched_text, candidate_used, bbox, page_no, context_keywords, error)
                    VALUES %s
                """, field_values)

        conn.commit()

    def save_documents_batch(self, reports: list[Dict[str, Any]]):
        """Save multiple documents in a batch."""
        if not reports:
            return

        conn = self.connect()
        doc_ids = [r['document_id'] for r in reports]

        with conn.cursor() as cursor:
            # Delete existing records
            cursor.execute(
                "DELETE FROM documents WHERE document_id = ANY(%s)",
                (doc_ids,)
            )

            # Batch insert documents
            doc_values = [
                (
                    r.get('document_id'),
                    r.get('pdf_path'),
                    r.get('pdf_type'),
                    r.get('success'),
                    r.get('total_pages'),
                    r.get('fields_matched'),
                    r.get('fields_total'),
                    r.get('annotations_generated'),
                    r.get('processing_time_ms'),
                    r.get('timestamp'),
                    json.dumps(r.get('errors', [])),
                    # New fields
                    r.get('split'),
                    r.get('customer_number'),
                    r.get('supplier_name'),
                    r.get('supplier_organisation_number'),
                    r.get('supplier_accounts'),
                )
                for r in reports
            ]
            execute_values(cursor, """
                INSERT INTO documents
                (document_id, pdf_path, pdf_type, success, total_pages,
                 fields_matched, fields_total, annotations_generated,
                 processing_time_ms, timestamp, errors,
                 split, customer_number, supplier_name, supplier_organisation_number, supplier_accounts)
                VALUES %s
            """, doc_values)

            # Batch insert field results
            field_values = []
            for r in reports:
                doc_id = r.get('document_id')
                for field in r.get('field_results', []):
                    field_values.append((
                        doc_id,
                        field.get('field_name'),
                        field.get('csv_value'),
                        field.get('matched'),
                        field.get('score'),
                        field.get('matched_text'),
                        field.get('candidate_used'),
                        json.dumps(field.get('bbox')) if field.get('bbox') else None,
                        field.get('page_no'),
                        json.dumps(field.get('context_keywords', [])),
                        field.get('error')
                    ))

            if field_values:
                execute_values(cursor, """
                    INSERT INTO field_results
                    (document_id, field_name, csv_value, matched, score,
                     matched_text, candidate_used, bbox, page_no, context_keywords, error)
                    VALUES %s
                """, field_values)

        conn.commit()
