# -*- coding: utf-8 -*-
import oci
import json
import tempfile
import os
import re
from datetime import datetime
from typing import List, Dict, Any

import oracledb
from extract_msg import Message as MsgReader
from oci.object_storage import ObjectStorageClient

from langchain_core.prompts import ChatPromptTemplate
from langchain_community.chat_models.oci_generative_ai import ChatOCIGenAI

# ============================================================
# DATABASE TABLE DEFINITION
# ============================================================
# CREATE TABLE PROCESSED_EMAILS (
#     ID NUMBER GENERATED ALWAYS AS IDENTITY,
#     DATE_SENT TIMESTAMP,
#     BOOKING VARCHAR2(100),
#     EMAIL_FROM VARCHAR2(500),
#     EMAIL_TO VARCHAR2(1000),
#     SUBJECT VARCHAR2(1000),
#     STATUS VARCHAR2(100),
#     INSERT_DATE TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
#     PRIMARY KEY (ID)
# );
#
# ALTER TABLE PROCESSED_EMAILS ADD (
#     BRIEF_DESCRIPTION VARCHAR2(1000),
#     ALERT VARCHAR2(10),
#     THREAD_INDEX NUMERIC(10),
#     FILE_NAME VARCHAR(400)
# );

# ============================================================
# CONFIGURATION
# ============================================================
with open("config", "r") as f:
    config_data = json.load(f)

MODEL_ID = "cohere.command-a-03-2025"
SERVICE_ENDPOINT = config_data["llm_endpoint"]
COMPARTMENT_ID = config_data["compartment_id"]
AUTH_PROFILE = config_data["oci_profile"]

WALLET_PATH = config_data["WALLET_PATH"]
DB_ALIAS = config_data["DB_ALIAS"]
USERNAME = config_data["USERNAME"]
PASSWORD = config_data["PASSWORD"]

BUCKET_NAME = config_data["bucket"]
PROCESSED_PREFIX = "processed/"
CONFIG_PROFILE = AUTH_PROFILE

MAX_CHARS_PER_CALL = 10000  # ~2500 tokens per block

# ============================================================
# OCI CLIENTS & DATABASE CONNECTION
# ============================================================
config = oci.config.from_file("~/.oci/config", CONFIG_PROFILE)
namespace = oci.object_storage.ObjectStorageClient(config).get_namespace().data
object_client = ObjectStorageClient(config)

connection = oracledb.connect(
    user=USERNAME,
    password=PASSWORD,
    dsn=DB_ALIAS,
    config_dir=WALLET_PATH,
    wallet_location=WALLET_PATH,
    wallet_password=PASSWORD,
)

# ============================================================
# UTILS: LIST AND MANAGE OBJECTS IN BUCKET
# ============================================================
def list_files() -> List[str]:
    """List all .msg files in the configured Object Storage bucket."""
    objs = object_client.list_objects(namespace, BUCKET_NAME, fields="name")
    return [o.name for o in objs.data.objects if o.name.lower().endswith(".msg")]


def download_file(name: str) -> str:
    """Download a .msg file from OCI Object Storage to a temporary folder."""
    temp_dir = tempfile.mkdtemp()
    path = os.path.join(temp_dir, os.path.basename(name))
    with open(path, "wb") as f:
        data = object_client.get_object(namespace, BUCKET_NAME, name).data.content
        f.write(data)
    return path


def move_to_processed(name: str):
    """Move a processed file to the 'processed/' prefix in the same bucket."""
    destination = PROCESSED_PREFIX + os.path.basename(name)
    obj = object_client.get_object(namespace, BUCKET_NAME, name)
    object_client.put_object(namespace, BUCKET_NAME, destination, obj.data.content)
    object_client.delete_object(namespace, BUCKET_NAME, name)

# ============================================================
# .MSG EXTRACTION
# ============================================================
def extract_msg_text(path):
    """
    Extract only the first (most recent) message from the .msg file.
    Returns:
    {
        "top": {headers},
        "body": "message body"
    }
    """
    msg = MsgReader(path)

    top = {
        "from": getattr(msg, "sender", ""),
        "to": getattr(msg, "to", ""),
        "subject": getattr(msg, "subject", ""),
        "date": str(getattr(msg, "date", "")),
    }

    body = getattr(msg, "body", "") or ""

    # Remove previous messages (older thread content)
    thread_pattern = re.compile(
        r"(?i)(^-{2,}\s*(original message|forwarded message|mensagem original)\s*-{2,}$|^from:|^sent:|^to:|^de:|^enviado:|^para:)",
        re.MULTILINE
    )

    parts = thread_pattern.split(body)
    cleaned_body = parts[0].strip() if parts else body.strip()

    return {"top": top, "body": cleaned_body}

# ============================================================
# THREAD SPLITTING
# ============================================================
HEADER_PATTERNS = [
    r"^-{2,}\s*Original Message\s*-{2,}$",
    r"^-{2,}\s*Forwarded message\s*-{2,}$",
    r"^From:\s", r"^Sent:\s", r"^To:\s", r"^Subject:\s"
]
BOUNDARY_REGEX = re.compile("|".join(HEADER_PATTERNS), re.IGNORECASE | re.MULTILINE)


def split_email_thread_refined(text):
    """
    Split an email body into message blocks within a thread,
    considering standard header sets.
    """
    text = text.replace('\r\n', '\n').strip()

    regex_start_email = re.compile(
        r"(?=^(From|De):.*\n(Sent|Enviado):.*\n(To|Para):.*)",
        re.MULTILINE | re.IGNORECASE
    )

    parts = regex_start_email.split(text)
    blocks = []

    if len(parts) <= 1:
        return [text.strip()] if text.strip() else []

    current = ""
    for line in parts:
        line = line.strip()
        if not line:
            continue

        if re.match(r"^(From|De):", line, re.IGNORECASE):
            if current:
                blocks.append(current.strip())
            current = line
        else:
            current += "\n" + line

    if current.strip():
        blocks.append(current.strip())

    cleaned_blocks = []
    for b in blocks:
        b = b.strip()
        if b and b not in cleaned_blocks:
            cleaned_blocks.append(b)

    return cleaned_blocks


def split_thread_by_headers(body: str) -> List[str]:
    lines = body.splitlines()
    blocks, curr = [], []

    def flush():
        if curr:
            text = "\n".join(curr).strip()
            if text:
                blocks.append(text)

    for line in lines:
        if BOUNDARY_REGEX.match(line.strip()) and curr:
            flush()
            curr = [line]
        else:
            curr.append(line)
    flush()

    if not blocks:
        return [body.strip()] if body.strip() else []
    return blocks


def chunk_by_size(text: str, max_chars: int) -> List[str]:
    """Split text into chunks within a character limit."""
    text = text.strip()
    if len(text) <= max_chars:
        return [text]
    chunks, curr = [], ""
    for part in re.split(r"(\n\s*\n)", text):
        if len(curr) + len(part) <= max_chars:
            curr += part
        else:
            if curr.strip():
                chunks.append(curr.strip())
            if len(part) > max_chars:
                for i in range(0, len(part), max_chars):
                    chunks.append(part[i:i + max_chars].strip())
                curr = ""
            else:
                curr = part
    if curr.strip():
        chunks.append(curr.strip())
    return chunks


def prepare_blocks_for_llm(body: str, top_headers: Dict[str, Any]) -> List[str]:
    """Prepare text blocks for LLM processing."""
    body = (body or "").strip()
    if not body:
        hint = f"From: {top_headers.get('from','')}\nTo: {top_headers.get('to','')}\nSubject: {top_headers.get('subject','')}\nDate: {top_headers.get('date','')}"
        return [hint]
    by_headers = split_thread_by_headers(body)
    result = []
    for b in by_headers:
        for c in chunk_by_size(b, MAX_CHARS_PER_CALL):
            if c.strip():
                result.append(c.strip())
    if not result:
        hint = f"From: {top_headers.get('from','')}\nTo: {top_headers.get('to','')}\nSubject: {top_headers.get('subject','')}\nDate: {top_headers.get('date','')}\n\n{body[:8000]}"
        result = [hint]
    return result

# ============================================================
# PROMPT (ENGLISH VERSION)
# ============================================================
def build_prompt(block_text: str, top_headers: Dict[str, Any]) -> ChatPromptTemplate:
    sys_text = """
You are a data extractor for corporate logistics emails.
Each email may contain multiple communications (replies, forwards, or quoted messages) within the same thread.
Your task is to identify each communication separately and extract the structured fields below, returning everything as a JSON list.

Extraction rules:

1. **Communication identification**
   - Each individual message starts when a new header with "From:", "Sent:", or similar appears.
   - Also consider blocks separated by "Original Message" or "Forwarded Message".
   - Each identified block must generate a separate JSON object.

2. **Fields to extract (per message)**
   - "date_sent": Date and time the message was sent (ISO 8601).
   - "booking": Booking or shipment code (“Booking”, “Bkg”, “BK”, etc.).
   - "email_from": Sender (“From:”).
   - "email_to": Main recipients (“To:”).
   - "subject": Message subject (“Subject:”).
   - "status": Operational status inferred from text using logistics keywords (shipment, invoice, forwarding, approval, etc.).
   - "brief_description": Short (1–2 lines) natural language summary.
   - "alert": "YES" if the message contains urgency, problems, delays, or missing documents; "NO" otherwise.

3. **Output format**
   - Return ONLY JSON (no extra text).
   - Example:
     [
       {
         "date_sent": "",
         "booking": "",
         "email_from": "",
         "email_to": "",
         "subject": "",
         "status": "",
         "brief_description": "",
         "alert": ""
       }
     ]

Context:
From: {from_hint}
To: {to_hint}
Subject: {subject_hint}
Date: {date_hint}

Text to analyze:
"""
    prompt = ChatPromptTemplate.from_messages([
        ("system", sys_text),
        ("human", block_text)
    ])
    prompt = prompt.partial(
        from_hint=top_headers.get("from", ""),
        to_hint=top_headers.get("to", ""),
        subject_hint=top_headers.get("subject", ""),
        date_hint=top_headers.get("date", ""),
    )
    return prompt

# ============================================================
# LLM AND DATABASE OPERATIONS
# ============================================================
def clean_markdown_fences(text: str) -> str:
    text = text.strip()
    text = re.sub(r"^```[a-zA-Z]*\s*", "", text)
    text = re.sub(r"\s*```$", "", text)
    return text.strip()


def call_llm_extract(block_text: str, top_headers: Dict[str, Any]) -> Dict[str, Any]:
    llm = ChatOCIGenAI(
        model_id=MODEL_ID,
        service_endpoint=SERVICE_ENDPOINT,
        compartment_id=COMPARTMENT_ID,
        auth_profile=AUTH_PROFILE,
        model_kwargs={"temperature": 0.0, "max_tokens": 1800, "top_p": 0.9},
    )
    prompt = build_prompt(block_text, top_headers)
    resp = llm.invoke(prompt.format_messages())
    content = clean_markdown_fences(resp.content or "")
    try:
        data = json.loads(content)
        if isinstance(data, list) and data:
            return data[0]
        if isinstance(data, dict):
            return data
        return {"error": "Unexpected format", "raw": content[:500]}
    except Exception:
        return {"error": "Invalid JSON", "raw": content[:500]}


def normalize_datetime(value):
    if value is None:
        return None
    if isinstance(value, datetime):
        return value
    s = str(value).strip()
    if not s:
        return None
    try:
        if s.endswith("Z"):
            s = s[:-1] + "+00:00"
        return datetime.fromisoformat(s).replace(tzinfo=None)
    except Exception:
        for fmt in ("%Y-%m-%dT%H:%M:%S", "%Y-%m-%d"):
            try:
                return datetime.strptime(s, fmt)
            except Exception:
                continue
    return None


def sanitize_value(v):
    if v is None:
        return None
    if isinstance(v, (list, tuple, set)):
        return ", ".join(str(x) for x in v)
    if isinstance(v, dict):
        return json.dumps(v, ensure_ascii=False)
    return str(v)


def insert_into_database(item, file_name, thread_index):
    with connection.cursor() as cursor:
        binds = {
            "date_sent": normalize_datetime(item.get("date_sent")),
            "booking": sanitize_value(item.get("booking")),
            "email_from": sanitize_value(item.get("email_from")),
            "email_to": sanitize_value(item.get("email_to")),
            "subject": sanitize_value(item.get("subject")),
            "status": sanitize_value(item.get("status")),
            "brief_description": sanitize_value(item.get("brief_description")),
            "alert": sanitize_value(item.get("alert")),
            "file_name": file_name,
            "thread_index": thread_index
        }
        cursor.execute("""
                       INSERT INTO PROCESSED_EMAILS
                       (DATE_SENT, BOOKING, EMAIL_FROM, EMAIL_TO, SUBJECT, STATUS, BRIEF_DESCRIPTION, ALERT, FILE_NAME, THREAD_INDEX)
                       VALUES
                           (:date_sent, :booking, :email_from, :email_to, :subject, :status, :brief_description, :alert, :file_name, :thread_index)
                       """, binds)
    connection.commit()

# ============================================================
# PIPELINE
# ============================================================
def process_emails():
    files = list_files()
    for name in files:
        print(f"\n=== Processing: {name} ===")
        path = download_file(name)
        data = extract_msg_text(path)
        top_headers = data["top"]
        body = data["body"]
        blocks = split_email_thread_refined(body)

        last_booking_detected = None
        thread_counter = 0

        for block in blocks:
            print("-------------------------------------")
            print("headers: ", top_headers)
            print("block: ", block)

            thread_counter += 1
            result = call_llm_extract(block, top_headers)

            if isinstance(result, dict) and "error" not in result:
                current_booking = (result.get("booking") or "").strip()

                if current_booking and current_booking != last_booking_detected:
                    print(f"[INFO] New booking detected: {current_booking}")
                    last_booking_detected = current_booking

                elif not current_booking and last_booking_detected:
                    result["booking"] = last_booking_detected

                insert_into_database(result,
                                     file_name=os.path.basename(name),
                                     thread_index=thread_counter)
                print(f"[OK] Thread {thread_counter}: {result}")

            else:
                print(f"[WARNING] Error in block {thread_counter}: {result}")

        if last_booking_detected:
            with connection.cursor() as cur:
                cur.execute("""
                            UPDATE PROCESSED_EMAILS
                            SET BOOKING = :booking
                            WHERE FILE_NAME = :file_name
                              AND (BOOKING IS NULL OR TRIM(BOOKING) = '')
                            """, {"booking": last_booking_detected, "file_name": os.path.basename(name)})
            connection.commit()
            print(f"[UPDATE] Booking '{last_booking_detected}' propagated in {os.path.basename(name)}")

        move_to_processed(name)


if __name__ == "__main__":
    process_emails()