import re
import json
import os
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

INDIAN_STATES = {
    "Andhra Pradesh", "Arunachal Pradesh", "Assam", "Bihar", "Chhattisgarh",
    "Goa", "Gujarat", "Haryana", "Himachal Pradesh", "Jharkhand", "Karnataka",
    "Kerala", "Madhya Pradesh", "Maharashtra", "Manipur", "Meghalaya", "Mizoram",
    "Nagaland", "Odisha", "Punjab", "Rajasthan", "Sikkim", "Tamil Nadu",
    "Telangana", "Tripura", "Uttar Pradesh", "Uttarakhand", "West Bengal",
    "Delhi", "New Delhi", "Puducherry"
}

def preprocess_ocr_address(text):
    """
    Clean OCR text:
    - Keep everything up to the first 6-digit pincode
    - Strip trailing garbage like emails, URLs, underscores
    """
    if not text:
        return ""
    
    match = re.search(r"(.*\b\d{6}\b)", text)
    if match:
        text = match.group(1)
    
    text = re.sub(r"[\s_\[\]@]+$", "", text)
    return text.strip()

def clean_component(text):
    """Remove trailing hyphens, underscores, dots, spaces from address components."""
    if text:
        return re.sub(r'[^A-Za-z\s]', '', text).strip()
    return text

def parse_indian_address(text):
    """
    Convert cleaned OCR text into structured JSON fields:
    address_line_1, city, state, pincode
    """
    text = preprocess_ocr_address(text)
    
    if not text:
        return {
            "address_line_1": None, "city": None, "state": None, "pincode": None
        }
    
    pincode_match = re.search(r"\b\d{6}\b", text)
    pincode = pincode_match.group(0) if pincode_match else None
    
    text_without_pincode = text.replace(pincode, "").rstrip(", ") if pincode else text
    parts = [part.strip() for part in text_without_pincode.split(",") if part.strip()]
    
    address_line_1 = city = state = None
    
    for i in range(len(parts)-1, -1, -1):
        part = parts[i]
        if part in INDIAN_STATES:
            state = part
            if i > 0:
                city = parts[i-1]
                address_line_1 = ", ".join(parts[:i-1]) if i-1 > 0 else parts[0]
            else:
                address_line_1 = ", ".join(parts[:i])
            break
    else:
        if len(parts) >= 2:
            state = parts[-1]
            city = parts[-2]
            address_line_1 = ", ".join(parts[:-2])
        elif len(parts) == 1:
            address_line_1 = parts[0]

    return {
        "address_line_1": address_line_1 if address_line_1 else None,
        "city": clean_component(city) if clean_component(city) else None,
        "state": clean_component(state) if clean_component(state) else None,
        "pincode": pincode if pincode else None
    }

def get_structured_address(ocr_text: str, app_id: int, provider: str = "gemini", logger=None) -> dict:
    """Use AI providers (Gemini/OpenAI) to structure OCR address text."""
    provider = provider.lower()
    if logger:
        logger.send_log_event(f"Application_id: {app_id}. In get_structured_address. OCR Text: {ocr_text}.")
    
    config = {
        "gemini": {
            "api_key": os.getenv("GEMINI_API_KEY"),
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "model": "gemini-2.5-flash",
        },
        "openai": {
            "api_key": os.getenv("OPENAI_API_KEY"),
            "base_url": None,
            "model": "gpt-4o-mini",
        },
    }

    if provider not in config:
        raise ValueError("Provider must be 'gemini' or 'openai'")

    cfg = config[provider]
    client = OpenAI(api_key=cfg["api_key"], base_url=cfg.get("base_url"))

    system_prompt = (
        "You are an expert postal address parser.\n"
        "Return CLEAN, STANDARDIZED JSON only.\n"
        "- Fix OCR typos (e.g., Banglore → Bengaluru, R0AD → Road).\n"
        "- Remove symbol or special characters.\n"
        "- Remove aadhaar card words like 'unique identification authority of india' which are not part of address.\n"
        "- Use official city/state names.\n"
        "- Missing fields must be null.\n"
        "Required keys: address_line_1, city, state, pincode."
    )

    try:
        response = client.chat.completions.create(
            model=cfg["model"],
            temperature=0,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"OCR Text to parse:\n{ocr_text}"}
            ],
        )

        content = response.choices[0].message.content
        # Remove markdown code fences if present
        if content.startswith("```"):
            content = content.split("```")[1]
            content = content.replace("json", "", 1).strip()
        if logger:
            logger.send_log_event(f"Application_id: {app_id}. AI Response Content: {content}")
        return json.loads(content)

    except Exception as e:
        if logger:
            logger.send_log_event(f"Application_id: {app_id}. Error in get_structured_address: {str(e)}")
        return {
            "address_line_1": None, "city": None, "state": None, "pincode": None,
        }