import sys
import json
from hugchat import hugchat
from hugchat.login import Login
import os
import re
import torch
from transformers import pipeline
import librosa

# HugChat login credentials from environment variables (secrets)
EMAIL = os.environ.get("EMAIL")
PASSWD = os.environ.get("PASSWORD")

# Directory to store cookies
cookie_path_dir = "./cookies/"
os.makedirs(cookie_path_dir, exist_ok=True)

# Login to HugChat
sign = Login(EMAIL, PASSWD)
cookies = sign.login(cookie_dir_path=cookie_path_dir, save_cookies=True)
chatbot = hugchat.ChatBot(cookies=cookies.get_dict())

# Model and device configuration for Whisper transcription
MODEL_NAME = "openai/whisper-large-v3-turbo"
device = 0 if torch.cuda.is_available() else "cpu"

# Initialize Whisper pipeline
pipe = pipeline(
    task="automatic-speech-recognition",
    model=MODEL_NAME,
    chunk_length_s=30,
    device=device,
)

def transcribe_audio(audio_path):
    """
    Transcribe a local audio file using the Whisper pipeline.
    """
    try:
        # Ensure audio is mono and resampled to 16kHz
        audio, sr = librosa.load(audio_path, sr=16000, mono=True)

        # Perform transcription
        transcription = pipe(audio, batch_size=8, generate_kwargs={"language": "urdu"})["text"]
        return transcription

    except Exception as e:
        return f"Error processing audio: {e}"

# Get command-line arguments: audio file path and file name
audio_path = sys.argv[1]  # Path to the audio file
file_name = sys.argv[2]   # File name for metadata

# Transcribe the audio to get Urdu text
urdu_text = transcribe_audio(audio_path)
if "Error" in urdu_text:
    print(json.dumps({"error": urdu_text}))
    sys.exit(1)

def extract_metadata(file_name):
    """
    Extract metadata from the file name.
    Assumes the second-last chunk is the city, e.g., 
    'agent2_5_Multan_Pakistan.mp3' -> location = 'Multan'.
    
    Args:
        file_name (str): The name of the audio file.
    
    Returns:
        dict: Contains agent_username and location.
    """
    base = file_name.split(".")[0]  # Remove extension
    parts = base.split("_")
    if len(parts) >= 3:
        return {
            "agent_username": parts[0],
            "location": parts[-2]  # Second-last chunk
        }
    return {"agent_username": "Unknown", "location": "Unknown"}

# Extract metadata from file name
metadata = extract_metadata(file_name)
location = metadata["location"]

# Step 1: Translate Urdu to English with context-aware correction
english_text = chatbot.chat(
    f"The following Urdu text is about crops and their diseases, but it may contain errors or misheard words due to audio transcription issues. Please use context to infer the most likely correct crop names and disease terms, and then translate the text to English:\n\n{urdu_text}"
).wait_until_done()

# Step 2: Extract specific crops and diseases from the English text
extraction_prompt = f"""
Below is an English text about specific crops and possible diseases/pests:

{english_text}

Identify each specific Crop (like wheat, rice, cotton, etc.) mentioned and list any Diseases or Pests affecting that crop.

- If a disease or pest is mentioned without specifying a particular crop, list it under "No crop:".
- If a crop is mentioned but no diseases or pests are specified for it, include it with an empty diseases list.
- Do not include general terms like "crops" as a specific crop name.

Format your answer in this style (one entry at a time):

For specific crops with diseases:
1. CropName:
Diseases:
- DiseaseName
- AnotherDisease

For specific crops with no diseases:
2. NextCrop:
Diseases:

For standalone diseases:
3. No crop:
Diseases:
- StandaloneDisease

No extra text, just the structured bullet list.
"""
extraction_response = chatbot.chat(extraction_prompt).wait_until_done()

# Step 3: Parse the extraction response
lines = extraction_response.splitlines()
crops_and_diseases = []
current_crop = None
current_diseases = []

for line in lines:
    line = line.strip()
    if not line:
        continue

    # Match lines like "1. Wheat:" or "3. No crop:"
    match_crop = re.match(r'^(\d+)\.\s*(.+?):$', line)
    if match_crop:
        # Save previous crop and diseases if any
        if current_crop is not None or current_diseases:
            crops_and_diseases.append({
                "crop": current_crop,
                "diseases": current_diseases
            })
        # Process new crop
        crop_name = match_crop.group(2).strip()
        # Handle general terms as "No crop"
        if crop_name.lower() in ["no crop", "crops", "general crops"]:
            current_crop = None  # Standalone diseases
        else:
            current_crop = crop_name
        current_diseases = []
        continue

    # Skip "Diseases:" line
    if line.lower().startswith("diseases:"):
        continue

    # Add disease if line starts with '-'
    if line.startswith('-'):
        disease_name = line.lstrip('-').strip()
        if disease_name:
            current_diseases.append(disease_name)

# Append the last crop/diseases if present
if current_crop is not None or current_diseases:
    crops_and_diseases.append({
        "crop": current_crop,
        "diseases": current_diseases
    })

# Step 4: Get temperature for the location
temp_prompt = f"Give me weather of {location} in Celsius numeric only."
temperature_response = chatbot.chat(temp_prompt).wait_until_done()

# Parse temperature (extract first number found)
temperature = None
temp_match = re.search(r'(\d+)', temperature_response)
if temp_match:
    temperature = int(temp_match.group(1))

# Step 5: Build and output the final JSON
output = {
    "urdu_text": urdu_text,
    "english_text": english_text,
    "crops_and_diseases": crops_and_diseases,
    "temperature": temperature,
    "location": location
}

print(json.dumps(output))