import logging
import json
from typing import Dict, Any, List, Optional, Tuple
from sqlalchemy.orm import Session
from sqlalchemy import text
from openai import AzureOpenAI

try:
    from src.utils.settings import settings
except ImportError:
    from utils.settings import settings

from .services import chat_session_service, chat_history_service

logger = logging.getLogger(__name__)


# Database schema context for the LLM - only relevant tables and columns
DATABASE_SCHEMA = """
## Database Schema for Smart Inventory

### companies
- id: INTEGER (Primary Key)
- company_id: INTEGER (External company ID)
- company_name: VARCHAR(255)

### products
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- fk_product_category_id: INTEGER (Foreign Key -> categories.id)
- product_id: INTEGER (External product ID)
- product_name: VARCHAR(500)
- short_name: VARCHAR(255)
- description: TEXT
- brand_name: VARCHAR(255)
- sku: VARCHAR(255)
- is_perishable: BOOLEAN
- image_path: VARCHAR(1000)

### categories
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- category_id: INTEGER (External category ID)
- category_name: VARCHAR(255)

### locations
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- location_id: INTEGER (External location ID)
- location_name: VARCHAR(255)

### vendors
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- vendor_id: INTEGER (External vendor ID)
- vendor_name: VARCHAR(255)
- vendor_code: VARCHAR(100)

### sales_orders
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- location_id: INTEGER (Foreign Key -> locations.id)
- sold_at: TIMESTAMP WITH TIME ZONE
- channel: VARCHAR(50) -- e.g., 'store', 'online'
- order_date: TIMESTAMP WITH TIME ZONE
- ref_number: VARCHAR(100)

### sales_order_lines
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- sales_order_id: INTEGER (Foreign Key -> sales_orders.id)
- product_id: INTEGER (Foreign Key -> products.id)
- quantity: INTEGER
- unit_price: FLOAT
- promotion_id: INTEGER (Foreign Key -> discounts.id, nullable)

### inventory_batches
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- product_id: INTEGER (Foreign Key -> products.id)
- location_id: INTEGER (Foreign Key -> locations.id)
- batch_ref: VARCHAR(100) (Unique)
- quantity_on_hand: INTEGER
- expiry_date: TIMESTAMP WITH TIME ZONE (nullable)
- received_date: TIMESTAMP WITH TIME ZONE
- status: ENUM ('active', 'sold_out', 'expired', 'disposed', 'donated')

### inventory_movements
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- product_id: INTEGER (Foreign Key -> products.id)
- location_id: INTEGER (Foreign Key -> locations.id)
- batch_id: INTEGER (Foreign Key -> inventory_batches.id, nullable)
- movement_type: ENUM ('sale', 'receipt', 'sale_return', 'purchase_return', 'adjustment', 'transfer_in', 'transfer_out')
- quantity_delta: INTEGER (positive for increase, negative for decrease)
- reference: VARCHAR(100)
- created_at: TIMESTAMP WITH TIME ZONE

### purchase_orders
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- supplier_id: INTEGER
- location_id: INTEGER (Foreign Key -> locations.id)
- status: ENUM ('draft', 'pending_approval', 'sent', 'received', 'returned', 'closed')
- expected_delivery_date: TIMESTAMP WITH TIME ZONE
- order_date: TIMESTAMP WITH TIME ZONE
- ref_number: VARCHAR(100)

### purchase_order_lines
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- purchase_order_id: INTEGER (Foreign Key -> purchase_orders.id)
- product_id: INTEGER (Foreign Key -> products.id)
- ordered_qty: INTEGER
- received_qty: INTEGER
- unit_cost: FLOAT

### daily_sales
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- product_id: INTEGER (Foreign Key -> products.id)
- location_id: INTEGER (Foreign Key -> locations.id)
- sale_date: DATE
- quantity_sold: INTEGER
- total_amount: FLOAT

### service_level_daily
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- product_id: INTEGER (Foreign Key -> products.id)
- location_id: INTEGER (Foreign Key -> locations.id)
- date: DATE
- demand_qty: INTEGER
- fulfilled_qty: INTEGER
- lost_sales_qty: INTEGER
- service_level: FLOAT (0.0 - 1.0)

### inventory_snapshot_daily
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- product_id: INTEGER (Foreign Key -> products.id)
- location_id: INTEGER (Foreign Key -> locations.id)
- snapshot_date: DATE
- on_hand_qty: FLOAT
- inbound_qty: FLOAT
- outbound_qty: FLOAT

### slow_mover_snapshot (for slow/fast moving products analysis)
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- product_id: INTEGER (Foreign Key -> products.id)
- location_id: INTEGER (Foreign Key -> locations.id)
- snapshot_date: DATE
- on_hand_qty: FLOAT
- total_sold_7d: FLOAT (units sold in last 7 days)
- total_sold_30d: FLOAT (units sold in last 30 days)
- total_sold_90d: FLOAT (units sold in last 90 days)
- ads_7d: FLOAT (Average Daily Sales - 7 days)
- ads_30d: FLOAT (Average Daily Sales - 30 days)
- ads_90d: FLOAT (Average Daily Sales - 90 days)
- doh_90d: FLOAT (Days of Inventory on Hand based on ADS_90d)
- days_since_last_sale: INTEGER
- is_slow_mover: BOOLEAN
- slow_mover_severity: VARCHAR(20) -- 'watchlist', 'slow', 'dead'
- slow_mover_reason: VARCHAR(255)

### inventory_planning_snapshot (for reorder and stock planning)
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- product_id: INTEGER (Foreign Key -> products.id)
- location_id: INTEGER (Foreign Key -> locations.id)
- snapshot_date: DATE
- avg_daily_demand: FLOAT
- sigma_daily_demand: FLOAT
- lead_time_days: INTEGER
- review_period_days: INTEGER
- service_level_target: FLOAT
- current_safety_stock: FLOAT
- current_reorder_point: FLOAT
- forecast_avg_daily_demand_90d: FLOAT
- forecast_safety_stock_90d: FLOAT
- forecasted_reorder_point_90d: FLOAT
- on_hand_qty: FLOAT
- inbound_qty: FLOAT
- available_stock: FLOAT
- min_target: FLOAT
- max_target: FLOAT
- stock_status: VARCHAR(20) -- 'out_of_stock', 'low_stock', 'in_stock', 'overstock'
- recommended_order_qty: FLOAT
- should_reorder: BOOLEAN
- days_of_cover: FLOAT
- days_until_stockout: FLOAT
- is_urgent: BOOLEAN
- urgency_score: FLOAT

### monthly_forecasts
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- product_id: INTEGER (Foreign Key -> products.id)
- location_id: INTEGER (Foreign Key -> locations.id)
- target_month: DATE (First day of the forecasted month, e.g., 2026-01-01 for January 2026)
- forecast_qty: FLOAT (Predicted quantity for the entire month)
- model_version: VARCHAR(50)
- forecast_date: DATE (When the forecast was generated)
- created_at: TIMESTAMP WITH TIME ZONE

### product_prices
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- product_id: INTEGER (Foreign Key -> products.id)
- location_id: INTEGER (Foreign Key -> locations.id)
- cost_price_per_unit: FLOAT
- retail_price_excl_tax: FLOAT
- compare_at_price: FLOAT

### discounts
- id: INTEGER (Primary Key)
- company_id: INTEGER (Foreign Key -> companies.id)
- discount_id: INTEGER (External discount ID)
- discount_name: VARCHAR(255)
- discount_code: VARCHAR(100)
- discount_type: VARCHAR(50) -- e.g., 'percentage', 'fixed'
- value: FLOAT

## Important Relationships:
- All tables have company_id for multi-tenant filtering
- products -> categories via fk_product_category_id
- sales_order_lines -> sales_orders via sales_order_id
- inventory_movements -> inventory_batches via batch_id
- purchase_order_lines -> purchase_orders via purchase_order_id
"""


def get_openai_client() -> AzureOpenAI:
    """Initialize Azure OpenAI client with settings from environment"""
    return AzureOpenAI(
        api_key=settings.OPENAI_API_KEY,
        api_version=settings.OPENAI_API_VERSION,
        azure_endpoint=settings.OPENAI_API_ENDPOINT
    )


def generate_sql_query(client: AzureOpenAI, question: str, company_id: int, 
                       conversation_history: Optional[List[Dict[str, str]]] = None) -> str:
    """
    Agent 1: Convert natural language question to SQL query
    Uses conversation history for context in follow-up questions
    """
    system_prompt = f"""You are a SQL query generator for a PostgreSQL inventory management database.
Your task is to convert natural language questions into valid SQL queries.

{DATABASE_SCHEMA}

IMPORTANT RULES:
1. ALWAYS filter by company_id = {company_id} in your queries to ensure data isolation
2. Use proper JOINs when relating tables
3. Return ONLY the SQL query, no explanations or markdown
4. Use PostgreSQL syntax
5. Limit results to 100 rows maximum unless specifically asked for more
6. For date comparisons, use CURRENT_DATE or specific dates in 'YYYY-MM-DD' format
7. Always include relevant column names in SELECT (avoid SELECT *)
8. Use table aliases for readability
9. Handle NULL values appropriately
10. For aggregations, include appropriate GROUP BY clauses

CONVERSATION CONTEXT:
- You may receive conversation history for follow-up questions
- If user refers to "that product", "the first one", "those items", etc., use context from previous messages
- If user says "tell me more about X" or "what about location Y", build on previous context

CRITICAL - Product-Location Data:
- Tables like slow_mover_snapshot and inventory_planning_snapshot store data PER PRODUCT-LOCATION COMBINATION
- When asked for "top products" or similar, you should:
  a) Get the LATEST snapshot_date first (use subquery or MAX)
  b) Include location_id and location_name in results
  c) Order by the appropriate metric (ads_30d for fast, ads_90d for slow/stagnant, urgency_score for urgent)

QUERY PATTERNS FOR COMMON REQUESTS:

For "top N fastest/fast-moving products":
- Use slow_mover_snapshot table
- Filter by latest snapshot_date
- Order by ads_30d DESC (highest daily sales = fastest)
- Include product_name, location_name, ads_30d, total_sold_30d

For "top N slowest/slow-moving products":
- Use slow_mover_snapshot table  
- Filter by is_slow_mover = TRUE and latest snapshot_date
- Order by ads_90d ASC (lowest daily sales = slowest)
- Include product_name, location_name, ads_90d, days_since_last_sale, slow_mover_severity

For "most stagnant products" or "dead stock":
- Use slow_mover_snapshot table
- Filter by latest snapshot_date, optionally slow_mover_severity = 'dead'
- Order by days_since_last_sale DESC, ads_90d ASC
- Include on_hand_qty to show excess stock

For "most urgent orders" or "products needing reorder":
- Use inventory_planning_snapshot table
- Filter by latest snapshot_date AND is_urgent = TRUE or should_reorder = TRUE
- Order by urgency_score DESC
- Include product_name, location_name, on_hand_qty, days_of_cover, recommended_order_qty

For "out of stock" or "stockouts":
- Use inventory_planning_snapshot table
- Filter by stock_status = 'out_of_stock' and latest snapshot_date

For "low stock" products:
- Use inventory_planning_snapshot table
- Filter by stock_status = 'low_stock' and latest snapshot_date

For "overstock" products:
- Use inventory_planning_snapshot table
- Filter by stock_status = 'overstock' and latest snapshot_date

LATEST SNAPSHOT PATTERN:
Use this subquery pattern to get latest date:
WHERE snapshot_date = (SELECT MAX(snapshot_date) FROM table_name WHERE company_id = {company_id})

Return ONLY the raw SQL query, nothing else."""

    # Build messages with conversation history
    messages = [{"role": "system", "content": system_prompt}]
    
    # Add conversation history for context (if available)
    if conversation_history:
        for msg in conversation_history:
            if msg["role"] == "user":
                messages.append({"role": "user", "content": f"Previous question: {msg['content']}"})
            elif msg["role"] == "assistant":
                # Only include a brief summary of previous answers for context
                content = msg["content"]
                if len(content) > 500:
                    content = content[:500] + "..."
                messages.append({"role": "assistant", "content": f"Previous answer summary: {content}"})
    
    # Add current question
    messages.append({"role": "user", "content": f"Convert this question to a SQL query: {question}"})

    try:
        response = client.chat.completions.create(
            model=settings.OPENAI_DEPLOYMENT,
            messages=messages,
            temperature=0.1,
            max_completion_tokens=1000
        )
        
        sql_query = response.choices[0].message.content.strip()
        
        # Clean up the query - remove markdown code blocks if present
        if sql_query.startswith("```sql"):
            sql_query = sql_query[6:]
        if sql_query.startswith("```"):
            sql_query = sql_query[3:]
        if sql_query.endswith("```"):
            sql_query = sql_query[:-3]
        
        return sql_query.strip()
        
    except Exception as e:
        logger.error(f"Error generating SQL query: {str(e)}")
        raise


def execute_sql_query(db: Session, sql_query: str) -> Tuple[List[Dict[str, Any]], int]:
    """
    Execute the generated SQL query and return results
    """
    try:
        # Basic SQL injection prevention - only allow SELECT queries
        sql_upper = sql_query.upper().strip()
        if not sql_upper.startswith("SELECT"):
            raise ValueError("Only SELECT queries are allowed")
        
        # Check for dangerous keywords
        dangerous_keywords = ["DROP", "DELETE", "UPDATE", "INSERT", "ALTER", "TRUNCATE", "EXEC", "EXECUTE"]
        for keyword in dangerous_keywords:
            if keyword in sql_upper:
                raise ValueError(f"Query contains forbidden keyword: {keyword}")
        
        result = db.execute(text(sql_query))
        rows = result.fetchall()
        columns = result.keys()
        
        # Convert to list of dictionaries
        data = []
        for row in rows:
            row_dict = {}
            for i, col in enumerate(columns):
                value = row[i]
                # Handle non-serializable types
                if hasattr(value, 'isoformat'):
                    value = value.isoformat()
                row_dict[col] = value
            data.append(row_dict)
        
        return data, len(data)
        
    except Exception as e:
        logger.error(f"Error executing SQL query: {str(e)}")
        raise


def summarize_results(client: AzureOpenAI, question: str, sql_query: str, 
                      data: List[Dict[str, Any]], data_count: int,
                      conversation_history: Optional[List[Dict[str, str]]] = None) -> str:
    """
    Agent 2: Answer the user's question directly based on the data
    Uses conversation history for contextual follow-up answers
    """
    system_prompt = """You are a helpful inventory assistant for a business owner.
Your job is to DIRECTLY ANSWER their question using the data provided - not to summarize or explain what the data shows.

IMPORTANT GUIDELINES:
1. **Answer the question directly** - If they ask "what are the top 10 products?", just list them clearly
2. **Be concise** - Give the answer first, keep explanations brief
3. **Use simple language** - No technical jargon like "query", "database", "records"
4. **Format for easy reading** - Use numbered lists, bullet points when listing items
5. **Include relevant numbers** - Show quantities, percentages rounded nicely
6. **Add a brief insight** - One or two sentences of helpful context after the answer
7. **Be direct** - Start with the answer, not "Based on the data..." or "Here's what I found..."
8. **Use conversation context** - If this is a follow-up question, connect your answer to the previous conversation

Approach:
"Your top 10 fast-moving products are:
1. Product A - 1,500 units sold
2. Product B - 1,200 units sold
..."

If duplicates appear in the data, consolidate them and mention it briefly.
If there's no data, simply say "I don't have any data matching that request."

Keep your response SHORT and FOCUSED - just answer what they asked."""

    # Limit data in prompt to avoid token limits
    data_sample = data[:50] if len(data) > 50 else data
    data_str = json.dumps(data_sample, indent=2, default=str)
    
    # Build messages with conversation history
    messages = [{"role": "system", "content": system_prompt}]
    
    # Add relevant conversation history
    if conversation_history:
        # Include last few exchanges for context
        recent_history = conversation_history[-6:]  # Last 3 Q&A pairs
        for msg in recent_history:
            if msg["role"] == "user":
                messages.append({"role": "user", "content": msg["content"]})
            elif msg["role"] == "assistant":
                # Truncate long previous answers
                content = msg["content"]
                if len(content) > 300:
                    content = content[:300] + "..."
                messages.append({"role": "assistant", "content": content})
    
    # Add current question and data
    user_prompt = f"""Question: "{question}"

Data ({data_count} items):
{data_str}

Answer the question directly and concisely."""
    
    messages.append({"role": "user", "content": user_prompt})

    try:
        response = client.chat.completions.create(
            model=settings.OPENAI_DEPLOYMENT,
            messages=messages,
            temperature=0.3,
            max_completion_tokens=1500
        )
        
        return response.choices[0].message.content.strip()
        
    except Exception as e:
        logger.error(f"Error summarizing results: {str(e)}")
        raise


def summarize_conversation_history(client: AzureOpenAI, messages: List[Dict[str, str]]) -> str:
    """
    Agent 3: Summarize older conversation history when context gets too long
    """
    system_prompt = """You are a conversation summarizer. 
Your task is to create a concise summary of the conversation between a user and an inventory assistant.

GUIDELINES:
1. Capture the key questions asked and main insights provided
2. Include any specific products, locations, or metrics mentioned
3. Note any decisions or conclusions reached
4. Keep the summary under 500 words
5. Focus on information that would be useful for follow-up questions

Format: Write in third person, past tense. Example:
"The user asked about top-selling products. The assistant reported that Product X at Location Y had the highest sales..."
"""

    # Format conversation for summarization
    conversation_text = ""
    for msg in messages:
        role = "User" if msg["role"] == "user" else "Assistant"
        content = msg["content"]
        if len(content) > 500:
            content = content[:500] + "..."
        conversation_text += f"{role}: {content}\n\n"

    try:
        response = client.chat.completions.create(
            model=settings.OPENAI_DEPLOYMENT,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Summarize this conversation:\n\n{conversation_text}"}
            ],
            temperature=0.3,
            max_completion_tokens=600
        )
        
        return response.choices[0].message.content.strip()
        
    except Exception as e:
        logger.error(f"Error summarizing conversation: {str(e)}")
        # Return a basic fallback summary
        return "Previous conversation about inventory data analysis."


def process_chat_query(db: Session, company_id: int, question: str, 
                       user_id: int, store_id: int, branch_id: int,
                       chat_session_id: Optional[str] = None) -> Dict[str, Any]:
    """
    Main function to process a chat query with conversation continuity:
    1. Get or create session (Redis for in-memory, DB for persistence)
    2. Check if history needs summarization
    3. Generate SQL from natural language (with context from last 5 Q&A)
    4. Execute the SQL query
    5. Summarize the results (with context)
    6. Save to session history (both Redis and DB)
    """
    try:
        # Initialize OpenAI client
        client = get_openai_client()
        
        # Handle session - create new or get existing
        is_new_session = False
        if chat_session_id is None:
            # New conversation - create session in Redis
            chat_session_id = chat_session_service.create_session(company_id)
            is_new_session = True
            conversation_history = []
            logger.info(f"Created new chat session: {chat_session_id}")
            
            # Create session in database
            chat_history_service.create_session(
                db=db,
                chat_session_id=chat_session_id,
                company_id=company_id,
                user_id=user_id,
                store_id=store_id,
                branch_id=branch_id,
                first_question=question
            )
        else:
            # Existing conversation - check Redis first
            session_data = chat_session_service.get_session(chat_session_id)
            if session_data is None:
                # Redis session expired - check if exists in DB
                db_session = chat_history_service.get_session_by_chat_id(db, chat_session_id)
                if db_session is None:
                    # Session doesn't exist anywhere - create new one
                    chat_session_id = chat_session_service.create_session(company_id)
                    is_new_session = True
                    logger.info(f"Session not found, created new: {chat_session_id}")
                    
                    # Create in database
                    chat_history_service.create_session(
                        db=db,
                        chat_session_id=chat_session_id,
                        company_id=company_id,
                        user_id=user_id,
                        store_id=store_id,
                        branch_id=branch_id,
                        first_question=question
                    )
                    conversation_history = []
                else:
                    # Session exists in DB but Redis expired - recreate Redis session
                    chat_session_service.create_session(company_id)
                    # Get last 10 messages from DB for context (5 Q&A pairs)
                    conversation_history = chat_history_service.get_last_n_messages(db, chat_session_id, n=10)
                    logger.info(f"Redis expired, restored from DB: {chat_session_id}")
            else:
                # Redis session exists - get conversation history from DB (last 5 Q&A = 10 messages)
                conversation_history = chat_history_service.get_last_n_messages(db, chat_session_id, n=10)
                logger.info(f"Using existing session: {chat_session_id} with {len(conversation_history)} messages from DB")
        
        # Check if we need to summarize old context (Redis only - for LLM efficiency)
        if not is_new_session and chat_session_service.needs_summarization(chat_session_id):
            logger.info("Conversation history too long, summarizing...")
            session_data = chat_session_service.get_session(chat_session_id)
            if session_data:
                old_messages = session_data.get("messages", [])
                summary = summarize_conversation_history(client, old_messages)
                chat_session_service.set_context_summary(chat_session_id, summary, keep_last_n=4)
                # Refresh conversation history after summarization
                conversation_history = chat_session_service.get_conversation_history(chat_session_id)
                logger.info("Context summarized and history trimmed")
        
        # Agent 1: Generate SQL query (with conversation context - last 5 Q&A)
        logger.info(f"Generating SQL for question: {question}")
        sql_query = generate_sql_query(client, question, company_id, conversation_history)
        logger.info(f"Generated SQL: {sql_query}")
        
        # Execute the query
        logger.info("Executing SQL query...")
        data, data_count = execute_sql_query(db, sql_query)
        logger.info(f"Query returned {data_count} records")
        
        # Agent 2: Summarize results (with conversation context)
        logger.info("Summarizing results...")
        answer = summarize_results(client, question, sql_query, data, data_count, conversation_history)
        
        # Save messages to Redis session
        chat_session_service.add_message(chat_session_id, "user", question)
        chat_session_service.add_message(chat_session_id, "assistant", answer, raw_data=data[:10])
        
        # Save messages to database
        chat_history_service.add_message(db, chat_session_id, "user", question)
        chat_history_service.add_message(db, chat_session_id, "assistant", answer)
        
        return {
            "chat_session_id": chat_session_id,
            "raw_data": data[:100],  # Limit raw data in response
            "answer": answer
        }
        
    except ValueError as e:
        # Handle validation errors (dangerous queries)
        logger.warning(f"Query validation error: {str(e)}")
        
        # Still need to return a session ID
        if chat_session_id is None:
            chat_session_id = chat_session_service.create_session(company_id)
            # Also create in DB
            chat_history_service.create_session(
                db=db,
                chat_session_id=chat_session_id,
                company_id=company_id,
                user_id=user_id,
                store_id=store_id,
                branch_id=branch_id,
                first_question=question
            )
        
        return {
            "chat_session_id": chat_session_id,
            "raw_data": None,
            "answer": f"I couldn't process your request: {str(e)}"
        }
        
    except Exception as e:
        logger.error(f"Error processing chat query: {str(e)}")
        raise