import asyncio
import json
import os
import statistics
from datetime import datetime, date
from typing import Any, Dict, List, Optional
from sqlalchemy import select
from celery import Celery
from pydantic import BaseModel, Field, AliasChoices
from google import genai
from app.core.config import settings
from app.core.logging import logger
from app.db.session import AsyncSessionLocal
from app.repositories.job_repository import JobRepository
from app.db.models.transaction import Transaction
from app.db.models.job_summary import JobSummary
from app.clients.exchange_rate_client import ExchangeRateClient
from app.core.observability import time_it, benchmark

exchange_rate_client = ExchangeRateClient()


def clean_json_response(text: str) -> str:
    """Extracts the first JSON object or array from the text, handling markdown fences and conversational prefix/suffix."""
    text = text.strip()
    start_idx = -1
    for i, char in enumerate(text):
        if char in ("{", "["):
            start_idx = i
            break

    end_idx = -1
    for i in range(len(text) - 1, -1, -1):
        if text[i] in ("}", "]"):
            end_idx = i
            break

    if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
        return text[start_idx : end_idx + 1]

    return text


class TransactionClassification(BaseModel):
    txn_id: str = Field(
        description="The unique transaction ID (txn_id) from the input."
    )
    category: str = Field(
        description="Assigned category matching one of the allowed categories list."
    )


class BatchClassificationResponse(BaseModel):
    classifications: List[TransactionClassification]


@time_it("LLM Batch Classification Request")
async def classify_transactions_batch(
    client: genai.Client,
    batch_txns: List[Dict[str, Any]],
    allowed_categories: List[str],
) -> Dict[str, Dict[str, Any]]:
    """
    Calls the Gemini model to classify a batch of transactions.
    Implements up to 3 retries (4 attempts total) with exponential backoff.
    Returns a dictionary mapping txn_id to classification results.
    """
    # Format the input for the LLM
    prompt_data = []
    for t in batch_txns:
        prompt_data.append(
            {
                "txn_id": t["txn_id"],
                "merchant": t["merchant"],
                "amount": t["amount"],
                "currency": t.get("currency", "INR"),
            }
        )

    input_prompt = (
        "Classify the following financial transactions. For each transaction, match its txn_id "
        f"and assign one of the allowed categories: {', '.join(allowed_categories)} based on the merchant name and details:\n\n"
        f"{json.dumps(prompt_data, indent=2)}"
    )

    for attempt in range(1, 5):
        try:
            logger.info(
                f"Sending batch of {len(batch_txns)} transactions to LLM (Attempt {attempt})..."
            )
            loop = asyncio.get_running_loop()
            # Wrap synchronous genai SDK call in executor
            interaction = await loop.run_in_executor(
                None,
                lambda: client.interactions.create(
                    model="gemini-2.5-flash-lite",
                    input=input_prompt,
                    response_format={
                        "type": "text",
                        "mime_type": "application/json",
                        "schema": BatchClassificationResponse.model_json_schema(),
                    },
                ),
            )

            clean_text = clean_json_response(interaction.output_text)
            resp_obj = BatchClassificationResponse.model_validate_json(clean_text)

            results = {}
            for item in resp_obj.classifications:
                item_json = json.dumps(item.model_dump())
                results[item.txn_id] = {
                    "category": item.category,
                    "raw_response": item_json,
                    "failed": False,
                }
            return results
        except Exception as e:
            logger.warning(f"LLM Classification attempt {attempt} failed: {e}")
            if attempt == 4:
                break
            sleep_seconds = 2**attempt
            await asyncio.sleep(sleep_seconds)

    logger.error(
        "LLM Classification completely failed after all retries. Falling back to default."
    )
    results = {}
    for t in batch_txns:
        results[t["txn_id"]] = {
            "category": "Other",
            "raw_response": None,
            "failed": True,
        }
    return results


class JobNarrativeSummary(BaseModel):
    total_spend_by_currency: Dict[str, float] = Field(
        validation_alias=AliasChoices(
            "total_spend_by_currency", "spend_by_currency", "total_spend"
        ),
        description="Total spend grouped by original currency",
    )
    top_3_merchants: Dict[str, float] = Field(
        validation_alias=AliasChoices("top_3_merchants", "top_merchants", "top_3"),
        description="Top 3 merchants by spend volume (value in USD equivalent)",
    )
    anomaly_count: int = Field(
        validation_alias=AliasChoices(
            "anomaly_count", "total_anomalies", "anomalies", "anomaly_transactions"
        ),
        description="Total number of anomaly transactions flagged in the dataset",
    )
    spending_narrative: str = Field(
        validation_alias=AliasChoices("spending_narrative", "narrative", "summary"),
        description="A cohesive 2-3 sentence spending narrative",
    )
    risk_level: str = Field(
        validation_alias=AliasChoices("risk_level", "overall_risk_level", "risk"),
        description="Overall risk level: Low, Medium, or High",
    )


@time_it("LLM Job Narrative Summary Request")
async def generate_job_summary_llm(
    client: genai.Client,
    db_transactions: List[Transaction],
) -> Optional[JobNarrativeSummary]:
    """
    Calls the Gemini model to produce a structured JSON summary:
    total spend by currency, top 3 merchants, anomaly count, a 2-3 sentence spending narrative, and a risk level.
    """
    # Build a simplified list of transactions for context
    tx_data = [
        {
            "merchant": t.merchant,
            "amount": t.amount,
            "currency": t.currency,
            "category": t.category,
            "is_anomaly": t.is_anomaly,
            "anomaly_reason": t.anomaly_reason,
        }
        for t in db_transactions
    ]

    input_prompt = (
        "Analyze the following list of processed financial transactions and generate a structured JSON summary. "
        "The returned JSON MUST contain exactly the following keys:\n"
        '  - "total_spend_by_currency": dictionary mapping currency symbol/code to total amount spent (e.g. {"USD": 150.0, "INR": 25000.0})\n'
        '  - "top_3_merchants": dictionary mapping the top 3 merchants by spend to their total spend value in USD equivalent (assuming 83 INR per USD for conversion)\n'
        '  - "anomaly_count": integer showing the total count of flagged anomaly transactions (is_anomaly=True)\n'
        '  - "spending_narrative": string containing a cohesive 2-3 sentence narrative summarizing key spending insights and largest merchants\n'
        '  - "risk_level": string showing the overall risk level of the job: "Low", "Medium", or "High" (pick one based on the nature of anomalies)\n\n'
        f"{json.dumps(tx_data, indent=2)}"
    )

    for attempt in range(1, 4):
        try:
            logger.info(
                f"Calling LLM to generate job narrative summary (Attempt {attempt})..."
            )
            loop = asyncio.get_running_loop()
            interaction = await loop.run_in_executor(
                None,
                lambda: client.interactions.create(
                    model="gemini-2.5-flash-lite",
                    input=input_prompt,
                    response_format={
                        "type": "text",
                        "mime_type": "application/json",
                        "schema": JobNarrativeSummary.model_json_schema(),
                    },
                ),
            )
            clean_text = clean_json_response(interaction.output_text)
            resp_obj = JobNarrativeSummary.model_validate_json(clean_text)
            return resp_obj
        except Exception as e:
            logger.warning(f"Failed to generate LLM summary (Attempt {attempt}): {e}")
            if attempt == 3:
                break
            await asyncio.sleep(2)
    return None


# Define Celery app
celery_app = Celery(
    "alemno_worker",
    broker=settings.CELERY_BROKER_URL,
    backend=settings.CELERY_RESULT_BACKEND,
)

# Optional configurations
celery_app.conf.update(
    task_serializer="json",
    accept_content=["json"],
    result_serializer="json",
    timezone="UTC",
    enable_utc=True,
)


def run_async(coro):
    """Helper to run async functions synchronously in Celery worker context"""
    return asyncio.run(coro)


@celery_app.task(name="tasks.process_transaction_job")
def process_transaction_job(job_id: str, transactions: List[Dict[str, Any]]):
    """
    Background job to clean transactions, detect anomalies, calculate spend breakdown,
    generate narrative summary, and persist results to Transaction and JobSummary tables.
    """
    logger.info(f"Starting Celery task to process job_id: {job_id}")
    return run_async(_process_job_async(job_id, transactions))


@time_it("Process Job Background Task")
async def _process_job_async(job_id: str, transactions: List[Dict[str, Any]]):
    from app.db.session import engine

    with benchmark("Dispose Engine"):
        await engine.dispose()

    async with AsyncSessionLocal() as db:
        with benchmark("Initial Job Setup"):
            repo = JobRepository(db)
            job = await repo.get_by_id(job_id)
            if not job:
                logger.error(f"Job with ID {job_id} not found in database.")
                return

            # Update status to processing
            await repo.update(job, status="processing")
            await db.commit()

        try:
            db_transactions = []
            total_spend_inr = 0.0
            total_spend_usd = 0.0
            merchant_spends: Dict[str, float] = {}
            anomaly_count = 0

            USD_TO_INR = await exchange_rate_client.get_usd_to_inr_rate()

            with benchmark("Compute Account History Medians"):
                # 1. Fetch existing transaction amounts and currencies for all accounts in this job to calculate median
                account_ids = {
                    t.get("account_id") for t in transactions if t.get("account_id")
                }
                existing_amounts_inr = {}
                if account_ids:
                    stmt = select(
                        Transaction.account_id, Transaction.amount, Transaction.currency
                    ).where(Transaction.account_id.in_(account_ids))
                    res = await db.execute(stmt)
                    for acc_id, amt, curr in res.all():
                        if acc_id not in existing_amounts_inr:
                            existing_amounts_inr[acc_id] = []
                        # Standardize to INR for median calculation
                        amt_inr = amt * USD_TO_INR if curr.upper() == "USD" else amt
                        existing_amounts_inr[acc_id].append(amt_inr)

                # Group incoming transaction amounts standardized to INR
                incoming_amounts_inr = {}
                for t in transactions:
                    acc_id = t.get("account_id")
                    if acc_id:
                        if acc_id not in incoming_amounts_inr:
                            incoming_amounts_inr[acc_id] = []
                        amt = float(t["amount"])
                        curr = t.get("currency", "INR").strip().upper()
                        amt_inr = amt * USD_TO_INR if curr == "USD" else amt
                        incoming_amounts_inr[acc_id].append(amt_inr)

                # Precalculate median of (existing + incoming) transactions in INR
                account_medians_inr = {}
                for acc_id in account_ids:
                    all_amts = existing_amounts_inr.get(
                        acc_id, []
                    ) + incoming_amounts_inr.get(acc_id, [])
                    if all_amts:
                        account_medians_inr[acc_id] = statistics.median(all_amts)

            # 2. Call LLM to classify transactions without a category
            # Find all uncategorized incoming transactions
            uncategorized_txns = []
            for t in transactions:
                cat = t.get("category")
                if not cat or cat.strip().lower() in {"", "general", "other", "none"}:
                    uncategorized_txns.append(t)

            llm_results = {}
            api_key = settings.GEMINI_API_KEY or os.environ.get("GEMINI_API_KEY")
            if uncategorized_txns:
                # Load allowed categories from the database categories table
                from app.db.models.category import Category

                try:
                    res_cats = await db.execute(select(Category.name))
                    allowed_categories = [r[0] for r in res_cats.all()]
                except Exception as db_err:
                    logger.warning(f"Could not load categories from database: {db_err}")
                    allowed_categories = []

                if not allowed_categories:
                    allowed_categories = [
                        "Food",
                        "Shopping",
                        "Travel",
                        "Transport",
                        "Utilities",
                        "Cash Withdrawal",
                        "Entertainment",
                        "Other",
                    ]

                if api_key:
                    logger.info(
                        f"Found {len(uncategorized_txns)} uncategorized transactions. Initializing GenAI client..."
                    )
                    try:
                        client = genai.Client(api_key=api_key)
                        # Split into batches of 50 and fire all concurrently
                        batch_size = 50
                        batches = [
                            uncategorized_txns[i : i + batch_size]
                            for i in range(0, len(uncategorized_txns), batch_size)
                        ]
                        logger.info(
                            f"Dispatching {len(batches)} classification batch(es) concurrently..."
                        )
                        with benchmark(
                            "LLM Concurrent Batch Classification (all batches)"
                        ):
                            batch_results_list = await asyncio.gather(
                                *[
                                    classify_transactions_batch(
                                        client, batch, allowed_categories
                                    )
                                    for batch in batches
                                ],
                                return_exceptions=True,
                            )
                        for batch_result in batch_results_list:
                            if isinstance(batch_result, Exception):
                                logger.error(
                                    f"A classification batch failed: {batch_result}"
                                )
                            else:
                                llm_results.update(batch_result)
                    except Exception as e:
                        logger.error(
                            f"Failed to initialize GenAI Client or classify: {e}"
                        )
                        # Fallback for all
                        for t in uncategorized_txns:
                            llm_results[t["txn_id"]] = {
                                "category": "Other",
                                "raw_response": None,
                                "failed": True,
                            }
                else:
                    logger.warning(
                        "GEMINI_API_KEY is not set. Defaulting to local mock categorization."
                    )
                    # Mock categorization mapping based on merchant name keywords
                    for t in uncategorized_txns:
                        merchant_lower = t["merchant"].lower()
                        cat = (
                            "Other"
                            if "Other" in allowed_categories
                            else (
                                allowed_categories[0] if allowed_categories else "Other"
                            )
                        )
                        if (
                            any(
                                kw in merchant_lower
                                for kw in [
                                    "swiggy",
                                    "zomato",
                                    "restaurant",
                                    "cafe",
                                    "food",
                                ]
                            )
                            and "Food" in allowed_categories
                        ):
                            cat = "Food"
                        elif (
                            any(
                                kw in merchant_lower
                                for kw in ["ola", "uber", "taxi", "metro", "train"]
                            )
                            and "Transport" in allowed_categories
                        ):
                            cat = "Transport"
                        elif (
                            any(
                                kw in merchant_lower
                                for kw in [
                                    "amazon",
                                    "flipkart",
                                    "shopping",
                                    "store",
                                    "supermarket",
                                ]
                            )
                            and "Shopping" in allowed_categories
                        ):
                            cat = "Shopping"
                        elif (
                            any(
                                kw in merchant_lower
                                for kw in ["air", "flight", "hotel", "travel", "irctc"]
                            )
                            and "Travel" in allowed_categories
                        ):
                            cat = "Travel"
                        elif (
                            any(
                                kw in merchant_lower
                                for kw in [
                                    "electric",
                                    "power",
                                    "water",
                                    "bill",
                                    "utility",
                                    "utilities",
                                ]
                            )
                            and "Utilities" in allowed_categories
                        ):
                            cat = "Utilities"
                        elif (
                            any(
                                kw in merchant_lower
                                for kw in ["atm", "cash", "withdrawal"]
                            )
                            and "Cash Withdrawal" in allowed_categories
                        ):
                            cat = "Cash Withdrawal"
                        elif (
                            any(
                                kw in merchant_lower
                                for kw in [
                                    "netflix",
                                    "spotify",
                                    "movie",
                                    "show",
                                    "entertainment",
                                ]
                            )
                            and "Entertainment" in allowed_categories
                        ):
                            cat = "Entertainment"

                        llm_results[t["txn_id"]] = {
                            "category": cat,
                            "raw_response": f'{{"category": "{cat}", "confidence": 0.8, "status": "mocked_offline"}}',
                            "failed": False,
                        }

            for t in transactions:
                txn_id = t["txn_id"]
                raw_date = t["date"]
                # Parse date string to datetime.date object
                parsed_date = (
                    date.fromisoformat(raw_date)
                    if isinstance(raw_date, str)
                    else raw_date
                )
                merchant = t["merchant"].strip()
                amount = float(t["amount"])
                currency = t.get("currency", "INR").strip().upper()
                category = t.get("category") or ""
                account_id = t.get("account_id")

                # Spend calculation and currency conversion
                if currency == "USD":
                    amt_usd = amount
                    amt_inr = amount * USD_TO_INR
                else:
                    amt_inr = amount
                    amt_usd = amount / USD_TO_INR

                total_spend_inr += amt_inr
                total_spend_usd += amt_usd

                # Aggregate merchant spend (in USD for standardization)
                merchant_spends[merchant] = merchant_spends.get(merchant, 0.0) + amt_usd

                # Anomaly detection checks
                reasons = []

                # A. 3x Median Outlier Check
                if account_id and account_id in account_medians_inr:
                    median_inr = account_medians_inr[account_id]
                    if amt_inr > 3.0 * median_inr:
                        reasons.append(
                            "Statistical outlier (amount exceeds 3x account median)"
                        )

                # B. Domestic Brand transacted in USD check
                domestic_brands = {"swiggy", "ola", "irctc"}
                merchant_lower = merchant.lower()
                if currency == "USD" and any(
                    brand in merchant_lower for brand in domestic_brands
                ):
                    reasons.append("Domestic brand transacted in USD")

                # C. High value transaction check
                if amt_usd > 5000.0:
                    reasons.append("High value transaction (> $5,000 USD equivalent)")

                # D. Suspicious keyword check
                desc_lower = merchant.lower()
                if any(
                    k in desc_lower
                    for k in ["suspend", "hack", "error", "unknown", "fraud"]
                ):
                    reasons.append("Suspicious merchant keyword")

                is_anomaly = len(reasons) > 0
                anomaly_reason = "; ".join(reasons) if is_anomaly else None

                if is_anomaly:
                    anomaly_count += 1

                # Check if it was uncategorized initially and assign LLM category
                cat_lower = category.strip().lower()
                is_uncategorized = not cat_lower or cat_lower in {
                    "",
                    "general",
                    "other",
                    "none",
                }

                if is_uncategorized and txn_id in llm_results:
                    classification = llm_results[txn_id]
                    llm_category = classification["category"].strip()
                    llm_raw_response = classification["raw_response"]
                    llm_failed = classification["failed"]
                    category = llm_category
                else:
                    llm_category = category.strip() if category else "Other"
                    llm_raw_response = None
                    llm_failed = False

                db_txn = Transaction(
                    job_id=job_id,
                    txn_id=txn_id,
                    date=parsed_date,
                    merchant=merchant,
                    amount=amount,
                    currency=currency,
                    status="cleaned",
                    category=category,
                    account_id=account_id,
                    is_anomaly=is_anomaly,
                    anomaly_reason=anomaly_reason,
                    llm_category=llm_category,
                    llm_raw_response=llm_raw_response,
                    llm_failed=llm_failed,
                )
                db_transactions.append(db_txn)

            with benchmark("Save Cleaned Transactions to DB"):
                await repo.add_transactions(db_transactions)

            # Generate narrative and risk summary using LLM if api_key is available
            llm_summary = None
            if api_key:
                try:
                    client = genai.Client(api_key=api_key)
                    llm_summary = await generate_job_summary_llm(
                        client, db_transactions
                    )
                except Exception as e:
                    logger.error(f"Error in LLM summary generation: {e}")

            if llm_summary:
                # Normalize risk level value to match capitalized Low/Medium/High in our DB
                risk_level_str = llm_summary.risk_level.strip().lower()
                if risk_level_str == "high":
                    risk_level = "High"
                elif risk_level_str == "medium":
                    risk_level = "Medium"
                else:
                    risk_level = "Low"

                narrative = llm_summary.spending_narrative.strip()
                top_merchants_json = {
                    merchant: round(spend, 2)
                    for merchant, spend in llm_summary.top_3_merchants.items()
                }
                anomaly_count = llm_summary.anomaly_count
            else:
                top_merchant = (
                    max(merchant_spends, key=merchant_spends.get)
                    if merchant_spends
                    else "None"
                )
                risk_level = "Low"
                if anomaly_count > 2:
                    risk_level = "High"
                elif anomaly_count > 0:
                    risk_level = "Medium"

                narrative = (
                    f"Successfully parsed and clean-processed {len(transactions)} transaction records. "
                    f"Total spend is INR {total_spend_inr:,.2f} (${total_spend_usd:,.2f} USD). "
                    f"Top merchant by spend volume is '{top_merchant}' with total of ${merchant_spends.get(top_merchant, 0):,.2f} USD. "
                    f"Scan flagged {anomaly_count} potential anomalies, resulting in a overall risk level of '{risk_level}'."
                )

                sorted_merchants = sorted(
                    merchant_spends.items(), key=lambda x: x[1], reverse=True
                )[:3]
                top_merchants_json = {
                    merchant: round(spend, 2) for merchant, spend in sorted_merchants
                }

            # Create job summary
            with benchmark("Save Summary and Complete Job"):
                summary = JobSummary(
                    job_id=job_id,
                    total_spend_inr=round(total_spend_inr, 2),
                    total_spend_usd=round(total_spend_usd, 2),
                    top_merchants=top_merchants_json,
                    anomaly_count=anomaly_count,
                    narrative=narrative,
                    risk_level=risk_level,
                )
                await repo.add_summary(summary)

                # Update job state
                await repo.update(
                    job,
                    status="completed",
                    row_count_raw=len(transactions),
                    row_count_clean=len(db_transactions),
                    completed_at=datetime.utcnow(),
                )
                await db.commit()
            logger.info(f"Job {job_id} processing completed successfully.")

        except Exception as e:
            logger.exception(f"Error processing job {job_id}")
            await repo.update(
                job,
                status="failed",
                error_message=str(e),
                completed_at=datetime.utcnow(),
            )
            await db.commit()
            raise e
