import os
import requests
import logging
import time
import pymysql
import json
from azure_upload import upload_to_azure_blob
from openai_helper import send_openai_analysis
from db_config import get_db_connection
from dotenv import load_dotenv
import urllib3

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
load_dotenv()

logging.basicConfig(level=logging.INFO)

SUBSCRIPTION_KEY = os.getenv('SARVAM_SUBSCRIPTION_KEY')
AUDIO_DIR = 'storage/audio/'


def send_to_sarvam(bid=None):
    conn = get_db_connection()
    cursor = conn.cursor(pymysql.cursors.DictCursor)
    
    # Use dynamic table name if bid is provided, otherwise use default
    if bid:
        calls_table = f"{bid}_calls"
        sarvam_response_table = f"{bid}_sarvamresponse"
    else:
        calls_table = "7417_calls"  # Default fallback
        sarvam_response_table = "7417_sarvamresponse"
    
    cursor.execute(f"SELECT * FROM {calls_table} WHERE status = 0")
    calls = cursor.fetchall()
    for call in calls:
        try:
            call_id = call['callid']
            file_url = call['fileUrl']
            local_file = os.path.join(AUDIO_DIR, f'translate_{call_id}.wav')
            response = requests.get(file_url)
            if response.status_code != 200:
                logging.error(f"❌ Failed to download file for callid {call_id}. Status code: {response.status_code}")
                continue

            with open(local_file, 'wb') as f:
                f.write(response.content)

            if not os.path.exists(local_file):
                logging.error(f"❌ File missing after download for callid {call_id}")
                continue

            job = init_sarvam_job()
            if not job:
                logging.error(f"❌ Failed to initialize Sarvam job for callid {call_id}")
                continue

            job_id = job['job_id']
            output_url = job['output_storage_path']
            azure_url = job['input_storage_path'].split('?')[0].rstrip('/') + '/' + os.path.basename(local_file) + '?' + job['input_storage_path'].split('?')[1]

            if not upload_to_azure_blob(azure_url, local_file):
                logging.error(f"❌ Failed to upload file to Azure for callid {call_id}")
                continue

            if not start_sarvam_job(job_id):
                logging.error(f"❌ Failed to start Sarvam job for jobid {job_id}")
                continue

            result = poll_sarvam_status(job_id, output_url)
            if not result:
                logging.error(f"❌ Failed to get result for jobid {job_id}")
                continue

            cursor.execute(f""" 
                INSERT INTO {sarvam_response_table} (callid, transcript, request_id, language, raw_response, status, created_at, updated_at)
                VALUES (%s, %s, %s, %s, %s, %s, NOW(), NOW()) 
            """, (call_id, result.get('transcript'), job_id, result.get('language_code', 'unknown'), str(result), 1))
            conn.commit()

            logging.info(f"✅ Translated transcript inserted for callid {call_id}")
            send_openai_analysis(call_id, bid)

            try:
                os.remove(local_file)
                logging.info(f"🧹 Removed local file for callid {call_id}")
            except Exception as e:
                logging.warning(f"⚠️ Could not delete file {local_file}: {e}")

        except Exception as e:
            logging.error(f"❌ Exception for callid {call['callid']}: {str(e)}")

    cursor.close()
    conn.close()
    return {"message": "Translation batch completed"}

def init_sarvam_job():
    headers = {
        'API-Subscription-Key': SUBSCRIPTION_KEY,
        'Content-Type': 'application/json'
    }
    try:
        response = requests.post('https://api.sarvam.ai/speech-to-text-translate/job/init', headers=headers, json={}, verify=True)
        if response.status_code == 202:
            return response.json()
        else:
            logging.error(f"❌ Job init failed: {response.status_code} - {response.text}")
            return None
    except requests.exceptions.RequestException as e:
        logging.error(f"❌ Request exception during job init: {str(e)}")
        return None

def start_sarvam_job(job_id):
    headers = {
        'API-Subscription-Key': SUBSCRIPTION_KEY,
        'Content-Type': 'application/json'
    }
    data = {
        "job_id": job_id,
        "job_parameters": {"with_diarization": False}
    }
    try:
        response = requests.post('https://api.sarvam.ai/speech-to-text-translate/job', headers=headers, json=data, verify=True)
        if response.status_code == 200:
            return True
        else:
            logging.error(f"❌ Failed to start Sarvam job: {response.status_code} - {response.text}")
            return False
    except requests.exceptions.RequestException as e:
        logging.error(f"❌ Request exception during job start: {str(e)}")
        return False

def poll_sarvam_status(job_id, output_url):
    headers = {'API-Subscription-Key': SUBSCRIPTION_KEY}
    for attempt in range(120):
        time.sleep(10)
        try:
            response = requests.get(f'https://api.sarvam.ai/speech-to-text-translate/job/{job_id}/status', headers=headers, verify=True)
            if response.status_code == 200:
                data = response.json()
                # Check if job_details exists and has elements before accessing
                if data.get('job_details') and len(data['job_details']) > 0:
                    logging.info(f"Job {job_id} state: {data['job_state']} - {data['job_details'][0]['state']}")
                    job_success = data['job_state'] == 'Completed' and data['job_details'][0]['state'] == 'Success'
                else:
                    logging.info(f"Job {job_id} state: {data['job_state']} - no job details available")
                    job_success = data['job_state'] == 'Completed'
                
                if job_success:
                    logging.info(f"✅ Job {job_id} completed successfully, fetching result...")
                    time.sleep(10)
                    if '?' in output_url:
                        base_url, query = output_url.split('?', 1)
                        final_url = f"{base_url.rstrip('/')}/0.json?{query}"
                    else:
                        final_url = output_url.rstrip('/') + '/0.json'
                    final_response = requests.get(final_url, verify=True)
                    logging.info(f"✅ Final response job {final_response}:{output_url}")
                    if final_response.status_code == 200:
                        if final_response.text:
                            try:
                                parsed = final_response.json()
                                logging.info(f"✅ Final JSON for job {job_id}: {parsed}")
                                return parsed
                            except json.JSONDecodeError:
                                logging.error(f"❌ Failed to decode JSON for job {job_id}. Raw response: {final_response.text}")
                                return None
                        else:
                            logging.error(f"❌ Empty response for job {job_id}. Response: {final_response.text}")
                            return None
                    else:
                        logging.error(f"❌ Failed to get content from output URL for job {job_id}. Response status: {final_response.status_code}")
                        return None
                else:
                    logging.info(f"❌ Job {job_id} not completed or failed, retrying... Current state: {data['job_state']}")
            else:
                logging.error(f"❌ Failed to check job status for job {job_id}. Status code: {response.status_code}")
        except requests.exceptions.RequestException as e:
            logging.error(f"❌ Exception while polling job status for job {job_id}: {e}")
            return None
    logging.warning(f"❌ Timeout polling job {job_id}")
    return None


def process_call_id(call_id, bid):
    conn = get_db_connection()
    cursor = conn.cursor(pymysql.cursors.DictCursor)

    # Use dynamic table name based on bid
    calls_table = f"{bid}_calls"
    sarvam_response_table = f"{bid}_sarvamresponse"
    
    cursor.execute(f"SELECT * FROM {calls_table} WHERE callid = %s AND status = 0", (call_id,))
    call = cursor.fetchone()
    if not call:
        logging.warning(f"⚠️ Call ID {call_id} not found or already processed in table {calls_table}.")
        return

    try:
        file_url = call['fileUrl']
        local_file = os.path.join(AUDIO_DIR, f'translate_{call_id}.wav')
        response = requests.get(file_url)
        if response.status_code != 200:
            logging.error(f"❌ Failed to download file for callid {call_id}. Status code: {response.status_code}")
            return

        with open(local_file, 'wb') as f:
            f.write(response.content)

        if not os.path.exists(local_file):
            logging.error(f"❌ File missing after download for callid {call_id}")
            return

        job = init_sarvam_job()
        if not job:
            logging.error(f"❌ Failed to initialize Sarvam job for callid {call_id}")
            return

        job_id = job['job_id']
        output_url = job['output_storage_path']
        azure_url = job['input_storage_path'].split('?')[0].rstrip('/') + '/' + os.path.basename(local_file) + '?' + job['input_storage_path'].split('?')[1]

        if not upload_to_azure_blob(azure_url, local_file):
            logging.error(f"❌ Failed to upload file to Azure for callid {call_id}")
            return

        if not start_sarvam_job(job_id):
            logging.error(f"❌ Failed to start Sarvam job for jobid {job_id}")
            return

        result = poll_sarvam_status(job_id, output_url)
        if not result:
            logging.error(f"❌ Failed to get result for jobid {job_id}")
            return

        cursor.execute(f""" 
            INSERT INTO {sarvam_response_table} (callid, transcript, request_id, language, raw_response, status, created_at, updated_at)
            VALUES (%s, %s, %s, %s, %s, %s, NOW(), NOW()) 
        """, (call_id, result.get('transcript'), job_id, result.get('language_code', 'unknown'), str(result), 1))

        cursor.execute(f"UPDATE {calls_table} SET status = 1 WHERE callid = %s", (call_id,))
        conn.commit()

        logging.info(f"✅ Translated transcript inserted and status updated for callid {call_id}")
        send_openai_analysis(call_id, bid)

        # Optional clean-up
        if os.path.exists(local_file):
            os.remove(local_file)
            logging.info(f"🧹 Deleted local audio file for callid {call_id}")

    except Exception as e:
        logging.error(f"❌ Exception for callid {call_id}: {str(e)}")

    finally:
        cursor.close()
        conn.close()

# Keep init_sarvam_job, start_sarvam_job, poll_sarvam_status here exactly as you had before