import os
import io
import re
import json
import zipfile
import pickle
import xml.etree.ElementTree as ET

import pandas as pd
import docx
from PIL import Image
import pytesseract
from email import policy
from email.parser import BytesParser
from bs4 import BeautifulSoup
from pdf2image import convert_from_path
import yaml

from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseDownload

# ===================== CONFIG =====================

SCOPES = ['https://www.googleapis.com/auth/drive.readonly']

# Optional: allow overriding via environment variables
TESSERACT_CMD = os.getenv("TESSERACT_CMD", r"C:\Program Files\Tesseract-OCR\tesseract.exe")
POPPLER_PATH = os.getenv("POPPLER_PATH", r"C:\Poppler\Release-25.12.0-0\poppler-25.12.0\Library\bin")

pytesseract.pytesseract.tesseract_cmd = TESSERACT_CMD

# ===================== GOOGLE DRIVE =====================

def get_drive_service():
    creds = None
    token_path = os.path.expanduser("~/.text_extractor_token.pickle")
    client_secret_path = os.environ.get("TEXT_EXTRACTOR_CLIENT_SECRET")

    if not client_secret_path or not os.path.exists(client_secret_path):
        raise FileNotFoundError(
            "Google OAuth client secret not found.\n"
            "Set environment variable:\n"
            "TEXT_EXTRACTOR_CLIENT_SECRET=/path/to/client_secret.json"
        )

    if os.path.exists(token_path):
        with open(token_path, 'rb') as token:
            creds = pickle.load(token)

    if not creds or not creds.valid:
        flow = InstalledAppFlow.from_client_secrets_file(client_secret_path, SCOPES)
        creds = flow.run_local_server(port=0)
        with open(token_path, 'wb') as token:
            pickle.dump(creds, token)

    return build('drive', 'v3', credentials=creds)


def get_drive_folder_id(link):
    match = re.search(r'/folders/([a-zA-Z0-9_-]+)', link)
    return match.group(1) if match else None


def download_drive_folder_any(link, download_path):
    service = get_drive_service()
    folder_id = get_drive_folder_id(link)
    os.makedirs(download_path, exist_ok=True)

    query = f"'{folder_id}' in parents and trashed=false"
    results = service.files().list(
        q=query,
        fields="files(id, name, mimeType)"
    ).execute()

    for file in results.get('files', []):
        file_path = os.path.join(download_path, file['name'])
        if file['mimeType'] == 'application/vnd.google-apps.folder':
            sub_link = f"https://drive.google.com/drive/folders/{file['id']}"
            download_drive_folder_any(sub_link, file_path)
        else:
            request = service.files().get_media(fileId=file['id'])
            with io.FileIO(file_path, 'wb') as fh:
                downloader = MediaIoBaseDownload(fh, request)
                done = False
                while not done:
                    _, done = downloader.next_chunk()


# ===================== TEXT EXTRACTOR =====================

class TextExtractor:
    def __init__(self):
        self.supported_extensions = [
            ".txt", ".md", ".json", ".csv", ".docx", ".pdf",
            ".png", ".jpg", ".jpeg", ".eml", ".xml", ".yml", ".yaml",
            ".xls", ".xlsx", ".html", ".pptx", ".ppt", ".zip"
        ]

    # ---------- BASIC TEXT ----------

    def _read_text(self, file_path):
        try:
            with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
                return f.read()
        except Exception as e:
            return f"Error reading {file_path}: {e}"

    def extract_text_from_txt(self, file_path): return self._read_text(file_path)
    def extract_text_from_md(self, file_path): return self._read_text(file_path)

    def extract_text_from_json(self, file_path):
        try:
            with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
                data = json.load(f)
            return json.dumps(data, indent=2, ensure_ascii=False)
        except Exception as e:
            return f"Error reading JSON {file_path}: {e}"

    def extract_text_from_csv(self, file_path):
        try:
            return pd.read_csv(file_path).to_string()
        except Exception as e:
            return f"Error reading CSV {file_path}: {e}"

    def extract_text_from_excel(self, file_path):
        try:
            return pd.read_excel(file_path).to_string()
        except Exception as e:
            return f"Error reading Excel {file_path}: {e}"

    # ---------- DOCX ----------

    def extract_text_from_docx(self, file_path):
        text = ""
        try:
            doc = docx.Document(file_path)
            for p in doc.paragraphs:
                if p.text.strip():
                    text += p.text + "\n"
            for table in doc.tables:
                for row in table.rows:
                    text += "\t".join(cell.text for cell in row.cells) + "\n"
        except Exception as e:
            return f"Error reading DOCX {file_path}: {e}"
        return text

    # ---------- PDF / IMAGE ----------

    def extract_text_from_pdf(self, file_path):
        text = ""
        try:
            pages = convert_from_path(file_path, poppler_path=POPPLER_PATH)
            for page in pages:
                ocr = pytesseract.image_to_string(page)
                if ocr.strip():
                    text += ocr + "\n"
        except Exception as e:
            return f"OCR failed for PDF {file_path}: {e}"
        return text

    def extract_text_from_image(self, file_path):
        try:
            img = Image.open(file_path).convert("L")
            return pytesseract.image_to_string(img)
        except Exception as e:
            return f"Error reading image {file_path}: {e}"

    # ---------- EML ----------

    def extract_text_from_eml(self, file_path, attachment_folder="temp_attachments"):
        os.makedirs(attachment_folder, exist_ok=True)
        body = ""
        try:
            with open(file_path, "rb") as f:
                msg = BytesParser(policy=policy.default).parse(f)

            # Try plain text first
            text_part = msg.get_body(preferencelist=("plain",))
            if text_part:
                body += text_part.get_content()
            # Fallback to HTML
            else:
                html_part = msg.get_body(preferencelist=("html",))
                if html_part:
                    soup = BeautifulSoup(html_part.get_content(), "html.parser")
                    body += soup.get_text()

            # Attachments
            for part in msg.iter_attachments():
                filename = part.get_filename()
                if filename:
                    path = os.path.join(attachment_folder, filename)
                    with open(path, "wb") as af:
                        af.write(part.get_payload(decode=True))
                    body += "\n[Attachment: {}]\n".format(filename)
                    body += self.extract_text_from_file(path)

        except Exception as e:
            return f"Error reading EML {file_path}: {e}"

        return body

    # ---------- STRUCTURED ----------

    def extract_text_from_xml(self, file_path):
        try:
            root = ET.parse(file_path).getroot()
            return "\n".join(e.text.strip() for e in root.iter() if e.text)
        except Exception as e:
            return f"Error reading XML {file_path}: {e}"

    def extract_text_from_yaml(self, file_path):
        try:
            with open(file_path, "r", encoding="utf-8") as f:
                return yaml.dump(yaml.safe_load(f), allow_unicode=True)
        except Exception as e:
            return f"Error reading YAML {file_path}: {e}"

    def extract_text_from_html(self, file_path):
        try:
            with open(file_path, "r", encoding="utf-8") as f:
                soup = BeautifulSoup(f.read(), "html.parser")
            return soup.get_text()
        except Exception as e:
            return f"Error reading HTML {file_path}: {e}"

    # ---------- PPT / ZIP ----------

    def extract_text_from_pptx(self, file_path):
        try:
            from pptx import Presentation
            prs = Presentation(file_path)
            return "\n".join(
                shape.text for slide in prs.slides
                for shape in slide.shapes
                if hasattr(shape, "text") and shape.text.strip()
            )
        except Exception as e:
            return f"Error reading PPTX {file_path}: {e}"

    def extract_text_from_zip(self, file_path, temp_folder="temp_zip"):
        os.makedirs(temp_folder, exist_ok=True)
        text = ""
        try:
            with zipfile.ZipFile(file_path) as z:
                z.extractall(temp_folder)
            for root, _, files in os.walk(temp_folder):
                for f in files:
                    text += self.extract_text_from_file(os.path.join(root, f))
        except Exception as e:
            return f"Error reading ZIP {file_path}: {e}"
        return text

    # ---------- DISPATCH ----------

    def extract_text_from_file(self, file_path):
        ext = os.path.splitext(file_path)[1].lower()

        handlers = {
            ".txt": self.extract_text_from_txt,
            ".md": self.extract_text_from_md,
            ".json": self.extract_text_from_json,
            ".csv": self.extract_text_from_csv,
            ".xls": self.extract_text_from_excel,
            ".xlsx": self.extract_text_from_excel,
            ".docx": self.extract_text_from_docx,
            ".pdf": self.extract_text_from_pdf,
            ".png": self.extract_text_from_image,
            ".jpg": self.extract_text_from_image,
            ".jpeg": self.extract_text_from_image,
            ".eml": self.extract_text_from_eml,
            ".xml": self.extract_text_from_xml,
            ".yml": self.extract_text_from_yaml,
            ".yaml": self.extract_text_from_yaml,
            ".html": self.extract_text_from_html,
            ".pptx": self.extract_text_from_pptx,
            ".zip": self.extract_text_from_zip,
        }

        return handlers.get(ext, lambda _: f"Unsupported file type: {ext}")(file_path)

    # ---------- FOLDER EXTRACTION ----------

    def extract_from_folder(self, folder_path):
        results = {}
        for root, _, files in os.walk(folder_path):
            for file in files:
                if os.path.splitext(file)[1].lower() in self.supported_extensions:
                    path = os.path.join(root, file)
                    results[path] = self.extract_text_from_file(path)
        return results

