"""
Training Data Processing Worker Tasks
Handles AI training data processing and embedding generation
"""
import logging
from datetime import datetime
from typing import Optional, Dict, Any, List
import json
import hashlib

from celery import shared_task
from sqlalchemy import select, and_
from sqlalchemy.orm import selectinload

from app.workers.celery_app import celery
from app.database import get_sync_session
from app.models.knowledge_base import (
    TrainingData, TrainingJob, KnowledgeBaseArticle,
    FAQ, DataSourceType, ArticleStatus
)
from app.models.client import Client

logger = logging.getLogger(__name__)


@celery.task(
    bind=True,
    name="app.workers.tasks.training_tasks.process_training_data",
    max_retries=3,
    default_retry_delay=120,
    queue="low_priority"
)
def process_training_data(self, training_data_id: int) -> Dict[str, Any]:
    """
    Process uploaded training data

    Args:
        training_data_id: TrainingData ID
    """
    try:
        logger.info(f"Processing training data {training_data_id}")

        with get_sync_session() as session:
            training_data = session.execute(
                select(TrainingData)
                .where(TrainingData.id == training_data_id)
            ).scalar_one_or_none()

            if not training_data:
                return {"success": False, "error": "Training data not found"}

            # Process based on source type
            content = None
            error = None

            try:
                if training_data.source_type == DataSourceType.DOCUMENT:
                    content = _process_document(training_data.file_path)
                elif training_data.source_type == DataSourceType.CSV:
                    content = _process_csv(training_data.file_path)
                elif training_data.source_type == DataSourceType.JSON:
                    content = _process_json(training_data.file_path)
                elif training_data.source_type == DataSourceType.WEBSITE:
                    content = _process_website(training_data.source_url)
                elif training_data.source_type in [DataSourceType.FAQ, DataSourceType.ARTICLE]:
                    # Already processed content
                    content = training_data.content

            except Exception as e:
                error = str(e)
                logger.error(f"Error processing training data {training_data_id}: {error}")

            # Update training data
            if content:
                # Generate content hash for deduplication
                content_hash = hashlib.sha256(content.encode()).hexdigest()

                training_data.content = content
                training_data.content_hash = content_hash
                training_data.is_processed = True
                training_data.processed_at = datetime.utcnow()
                training_data.processing_error = None
            else:
                training_data.is_processed = False
                training_data.processing_error = error or "Unknown error"

            session.commit()

            if content:
                # Queue embedding generation
                generate_embeddings.delay(training_data_id)

                return {
                    "success": True,
                    "training_data_id": training_data_id,
                    "content_length": len(content)
                }
            else:
                return {
                    "success": False,
                    "error": error
                }

    except Exception as e:
        logger.error(f"Error processing training data {training_data_id}: {str(e)}")
        self.retry(exc=e)


@celery.task(
    bind=True,
    name="app.workers.tasks.training_tasks.generate_embeddings",
    max_retries=3,
    default_retry_delay=60,
    queue="low_priority"
)
def generate_embeddings(self, training_data_id: int) -> Dict[str, Any]:
    """
    Generate embeddings for processed training data

    Args:
        training_data_id: TrainingData ID
    """
    try:
        logger.info(f"Generating embeddings for training data {training_data_id}")

        with get_sync_session() as session:
            training_data = session.execute(
                select(TrainingData)
                .where(TrainingData.id == training_data_id)
            ).scalar_one_or_none()

            if not training_data:
                return {"success": False, "error": "Training data not found"}

            if not training_data.content:
                return {"success": False, "error": "No content to process"}

            # Chunk content for embedding
            chunks = _chunk_content(training_data.content)

            # Generate embeddings using OpenAI or local model
            embeddings = []
            for chunk in chunks:
                embedding = _generate_embedding(chunk)
                if embedding:
                    embeddings.append(embedding)

            # Store embeddings (in practice, use a vector database)
            training_data.is_trained = True
            training_data.trained_at = datetime.utcnow()
            training_data.embedding_count = len(embeddings)

            session.commit()

            logger.info(f"Generated {len(embeddings)} embeddings for training data {training_data_id}")

            return {
                "success": True,
                "training_data_id": training_data_id,
                "embedding_count": len(embeddings)
            }

    except Exception as e:
        logger.error(f"Error generating embeddings: {str(e)}")
        self.retry(exc=e)


@celery.task(
    bind=True,
    name="app.workers.tasks.training_tasks.run_training_job",
    max_retries=2,
    queue="low_priority"
)
def run_training_job(self, job_id: int) -> Dict[str, Any]:
    """
    Run a full training job for a client

    Args:
        job_id: TrainingJob ID
    """
    try:
        logger.info(f"Running training job {job_id}")

        with get_sync_session() as session:
            job = session.execute(
                select(TrainingJob)
                .where(TrainingJob.id == job_id)
            ).scalar_one_or_none()

            if not job:
                return {"success": False, "error": "Job not found"}

            # Update job status
            job.status = "running"
            job.started_at = datetime.utcnow()
            session.commit()

            results = {
                "processed": 0,
                "errors": 0,
                "error_details": []
            }

            try:
                # Process based on job type
                if job.job_type == "full":
                    # Process all training data
                    training_items = session.execute(
                        select(TrainingData)
                        .where(
                            and_(
                                TrainingData.client_id == job.client_id,
                                TrainingData.is_deleted == False
                            )
                        )
                    ).scalars().all()

                    job.total_items = len(training_items)
                    session.commit()

                    for item in training_items:
                        try:
                            process_training_data.delay(item.id)
                            results["processed"] += 1
                        except Exception as e:
                            results["errors"] += 1
                            results["error_details"].append({
                                "item_id": item.id,
                                "error": str(e)
                            })

                        job.processed_items = results["processed"]
                        job.error_count = results["errors"]
                        session.commit()

                elif job.job_type == "faq_only":
                    # Process FAQs only
                    faqs = session.execute(
                        select(FAQ)
                        .where(
                            and_(
                                FAQ.client_id == job.client_id,
                                FAQ.include_in_training == True,
                                FAQ.is_deleted == False
                            )
                        )
                    ).scalars().all()

                    job.total_items = len(faqs)
                    session.commit()

                    for faq in faqs:
                        try:
                            _process_faq(faq)
                            results["processed"] += 1
                        except Exception as e:
                            results["errors"] += 1
                            results["error_details"].append({
                                "faq_id": faq.id,
                                "error": str(e)
                            })

                        job.processed_items = results["processed"]
                        job.error_count = results["errors"]
                        session.commit()

                elif job.job_type == "incremental":
                    # Process only unprocessed items
                    unprocessed = session.execute(
                        select(TrainingData)
                        .where(
                            and_(
                                TrainingData.client_id == job.client_id,
                                TrainingData.is_processed == False,
                                TrainingData.is_deleted == False
                            )
                        )
                    ).scalars().all()

                    job.total_items = len(unprocessed)
                    session.commit()

                    for item in unprocessed:
                        try:
                            process_training_data.delay(item.id)
                            results["processed"] += 1
                        except Exception as e:
                            results["errors"] += 1
                            results["error_details"].append({
                                "item_id": item.id,
                                "error": str(e)
                            })

                        job.processed_items = results["processed"]
                        job.error_count = results["errors"]
                        session.commit()

                # Mark job as completed
                job.status = "completed"
                job.completed_at = datetime.utcnow()
                job.result_summary = json.dumps(results)
                if results["error_details"]:
                    job.error_log = json.dumps(results["error_details"])
                session.commit()

                logger.info(f"Training job {job_id} completed: {results}")

                return {
                    "success": True,
                    "job_id": job_id,
                    "results": results
                }

            except Exception as e:
                job.status = "failed"
                job.completed_at = datetime.utcnow()
                job.error_log = str(e)
                session.commit()
                raise

    except Exception as e:
        logger.error(f"Error running training job {job_id}: {str(e)}")
        self.retry(exc=e)


@celery.task(
    name="app.workers.tasks.training_tasks.update_article_embeddings",
    queue="low_priority"
)
def update_article_embeddings(article_id: int) -> Dict[str, Any]:
    """
    Update embeddings for a knowledge base article

    Args:
        article_id: Article ID
    """
    try:
        logger.info(f"Updating embeddings for article {article_id}")

        with get_sync_session() as session:
            article = session.execute(
                select(KnowledgeBaseArticle)
                .where(KnowledgeBaseArticle.id == article_id)
            ).scalar_one_or_none()

            if not article:
                return {"success": False, "error": "Article not found"}

            if not article.include_in_training:
                return {"success": False, "error": "Article not marked for training"}

            # Generate embedding for article content
            content = f"{article.title}\n\n{article.summary or ''}\n\n{article.content}"
            embedding = _generate_embedding(content)

            if embedding:
                article.embedding_vector = json.dumps(embedding)
                article.last_trained_at = datetime.utcnow()
                session.commit()

                return {
                    "success": True,
                    "article_id": article_id
                }
            else:
                return {
                    "success": False,
                    "error": "Failed to generate embedding"
                }

    except Exception as e:
        logger.error(f"Error updating article embeddings: {str(e)}")
        return {"success": False, "error": str(e)}


@celery.task(
    name="app.workers.tasks.training_tasks.update_faq_embeddings",
    queue="low_priority"
)
def update_faq_embeddings(faq_id: int) -> Dict[str, Any]:
    """
    Update embeddings for a FAQ

    Args:
        faq_id: FAQ ID
    """
    try:
        logger.info(f"Updating embeddings for FAQ {faq_id}")

        with get_sync_session() as session:
            faq = session.execute(
                select(FAQ)
                .where(FAQ.id == faq_id)
            ).scalar_one_or_none()

            if not faq:
                return {"success": False, "error": "FAQ not found"}

            _process_faq(faq)
            session.commit()

            return {
                "success": True,
                "faq_id": faq_id
            }

    except Exception as e:
        logger.error(f"Error updating FAQ embeddings: {str(e)}")
        return {"success": False, "error": str(e)}


def _process_document(file_path: str) -> Optional[str]:
    """Extract text from document files"""
    try:
        import os
        ext = os.path.splitext(file_path)[1].lower()

        if ext == ".txt":
            with open(file_path, "r", encoding="utf-8") as f:
                return f.read()

        elif ext == ".pdf":
            # Use PyPDF2 or pdfplumber
            try:
                import pdfplumber
                text = ""
                with pdfplumber.open(file_path) as pdf:
                    for page in pdf.pages:
                        text += page.extract_text() or ""
                return text
            except ImportError:
                logger.warning("pdfplumber not installed")
                return None

        elif ext in [".doc", ".docx"]:
            # Use python-docx
            try:
                from docx import Document
                doc = Document(file_path)
                return "\n".join([para.text for para in doc.paragraphs])
            except ImportError:
                logger.warning("python-docx not installed")
                return None

        else:
            logger.warning(f"Unsupported document type: {ext}")
            return None

    except Exception as e:
        logger.error(f"Error processing document: {str(e)}")
        return None


def _process_csv(file_path: str) -> Optional[str]:
    """Extract content from CSV file"""
    try:
        import csv
        content_parts = []

        with open(file_path, "r", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            for row in reader:
                content_parts.append(" | ".join(f"{k}: {v}" for k, v in row.items()))

        return "\n".join(content_parts)

    except Exception as e:
        logger.error(f"Error processing CSV: {str(e)}")
        return None


def _process_json(file_path: str) -> Optional[str]:
    """Extract content from JSON file"""
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            data = json.load(f)

        # Convert JSON to readable text
        def flatten_json(obj, prefix=""):
            parts = []
            if isinstance(obj, dict):
                for k, v in obj.items():
                    parts.extend(flatten_json(v, f"{prefix}{k}: "))
            elif isinstance(obj, list):
                for i, item in enumerate(obj):
                    parts.extend(flatten_json(item, f"{prefix}[{i}] "))
            else:
                parts.append(f"{prefix}{obj}")
            return parts

        return "\n".join(flatten_json(data))

    except Exception as e:
        logger.error(f"Error processing JSON: {str(e)}")
        return None


def _process_website(url: str) -> Optional[str]:
    """Extract content from website URL"""
    try:
        import requests
        from bs4 import BeautifulSoup

        response = requests.get(url, timeout=30)
        response.raise_for_status()

        soup = BeautifulSoup(response.text, "html.parser")

        # Remove script and style elements
        for script in soup(["script", "style"]):
            script.decompose()

        # Get text
        text = soup.get_text()

        # Clean up whitespace
        lines = (line.strip() for line in text.splitlines())
        chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
        text = "\n".join(chunk for chunk in chunks if chunk)

        return text

    except ImportError:
        logger.warning("beautifulsoup4 or requests not installed")
        return None
    except Exception as e:
        logger.error(f"Error processing website: {str(e)}")
        return None


def _chunk_content(content: str, chunk_size: int = 1000, overlap: int = 100) -> List[str]:
    """Split content into overlapping chunks for embedding"""
    chunks = []
    start = 0

    while start < len(content):
        end = start + chunk_size
        chunk = content[start:end]
        chunks.append(chunk)
        start = end - overlap

    return chunks


def _generate_embedding(text: str) -> Optional[List[float]]:
    """Generate embedding vector for text"""
    try:
        from app.config import settings

        if settings.OPENAI_API_KEY:
            import openai
            openai.api_key = settings.OPENAI_API_KEY

            response = openai.embeddings.create(
                model="text-embedding-3-small",
                input=text[:8000]  # Limit text length
            )

            return response.data[0].embedding

        else:
            # Return placeholder for development
            logger.warning("No embedding API configured, using placeholder")
            return [0.0] * 384  # Placeholder embedding

    except Exception as e:
        logger.error(f"Error generating embedding: {str(e)}")
        return None


def _process_faq(faq: FAQ) -> None:
    """Process FAQ for training"""
    # Combine question and variations for embedding
    content = faq.question

    if faq.question_variations:
        try:
            variations = json.loads(faq.question_variations)
            content += "\n" + "\n".join(variations)
        except json.JSONDecodeError:
            pass

    # Add answer
    content += f"\n\nAnswer: {faq.answer}"

    # Generate embedding
    embedding = _generate_embedding(content)
    if embedding:
        faq.embedding_vector = json.dumps(embedding)
        faq.last_trained_at = datetime.utcnow()