from __future__ import annotations

import asyncio
import logging
import os
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


class DrawingPipelineRunnerMixin:
    def create_pending_takeoff(
        self,
        file_path: Path,
        filename: str,
        project_id: str = None,
        document_id: str = None,
        drawing_hint: str = "",
    ) -> dict:
        from uuid import uuid4

        takeoff_id = str(uuid4())
        payload = {
            "id": takeoff_id,
            "project_id": project_id,
            "document_id": document_id,
            "drawing_code": filename
            if (filename.startswith("DWG-") or filename.startswith("DWG_"))
            else f"DWG-{filename}",
            "file_path": str(file_path),
            "file_url": f"/storage/{file_path.name}",
            "takeoff_status": "pending",
            "drawing_hint": drawing_hint,
        }
        self._erp_client.create_drawing_takeoff(payload)
        return payload

    async def run_takeoff_pipeline(
        self,
        takeoff_id: str,
        drawing_hint: str = "",
        attached_hsmt_path: str | None = None,
        attached_hsmt_filename: str | None = None,
    ) -> dict:
        from uuid import uuid4

        try:
            takeoff = self._erp_client.get_drawing_takeoff(takeoff_id)
            if not takeoff:
                return {"status": "error", "detail": "Takeoff not found"}

            self._erp_client.update_drawing_takeoff(
                takeoff_id, {"takeoff_status": "processing"}
            )
            file_path_str = takeoff.get("file_path")
            if not file_path_str or not os.path.exists(file_path_str):
                self._erp_client.update_drawing_takeoff(
                    takeoff_id,
                    {
                        "takeoff_status": "errored",
                        "ai_analysis_notes": "File not found",
                    },
                )
                return {"status": "error"}

            # 1. Extraction / Render based on file extension
            # Run heavy CPU-bound rendering in a separate thread to prevent freezing the asyncio event loop
            ext = Path(file_path_str).suffix.lower()
            if ext == ".dwg":
                ext_pages = await asyncio.to_thread(
                    self.extract_dwg_text_and_preview, file_path_str
                )
            elif ext == ".dxf":
                ext_pages = await asyncio.to_thread(
                    self.render_dxf_to_images, file_path_str
                )
            elif ext == ".pdf":
                ext_pages = await asyncio.to_thread(
                    self.render_pdf_to_images, Path(file_path_str), 3, 150, takeoff_id
                )
            else:
                ext_pages = await asyncio.to_thread(
                    self.render_pdf_to_images, Path(file_path_str), 3, 150, takeoff_id
                )
            first_page = ext_pages[0] if ext_pages else {}

            # Fast Text Extraction
            text_layer = ""
            if attached_hsmt_path and os.path.exists(attached_hsmt_path):
                try:
                    import pymupdf

                    def extract_hsmt_text():
                        t_layer = ""
                        hsmt_doc = pymupdf.open(str(attached_hsmt_path))
                        for i in range(min(50, len(hsmt_doc))):
                            t_layer += hsmt_doc[i].get_text("text") + "\n"
                        hsmt_doc.close()
                        return t_layer

                    text_layer += await asyncio.to_thread(extract_hsmt_text)
                except Exception as e:
                    logger.warning(f"Failed to extract text from HSMT: {e}")
            elif ext == ".pdf":
                try:
                    import pymupdf

                    def extract_pdf_text():
                        t_layer = ""
                        doc = pymupdf.open(str(file_path_str))
                        for i in range(min(50, len(doc))):
                            t_layer += doc[i].get_text("text") + "\n"
                        doc.close()
                        return t_layer

                    text_layer += await asyncio.to_thread(extract_pdf_text)
                except Exception as e:
                    logger.warning(f"Failed to extract text from PDF: {e}")
            else:
                text_layer = first_page.get("text_layer", "")

            preview_base64 = first_page.get("image_base64", "")
            first_img_url = (
                ext_pages[0].get("image_url")
                if (ext_pages and ext_pages[0].get("image_url"))
                else (
                    f"data:image/png;base64,{preview_base64}" if preview_base64 else ""
                )
            )
            pages_metadata = [
                {
                    "page_number": p["page_number"],
                    "title": p["title"],
                    "image_url": p.get("image_url")
                    or (
                        f"data:image/png;base64,{p['image_base64']}"
                        if p.get("image_base64")
                        else ""
                    ),
                }
                for p in ext_pages
            ]

            # 2. AI Vision Analysis to extract items
            vision_result = await self.analyze_drawing_with_vision(
                image_base64=preview_base64,
                filename=takeoff.get("drawing_code", ""),
                text_layer=text_layer,
                drawing_hint=drawing_hint,
                ext_pages=ext_pages,
                file_path_str=file_path_str,
                attached_hsmt_path=attached_hsmt_path,
                attached_hsmt_filename=attached_hsmt_filename,
            )
            raw_items = vision_result.get("items", [])
            drawing_type = (
                vision_result.get("drawing_type")
                or takeoff.get("drawing_type")
                or "general"
            )
            if hasattr(self, "apply_learned_rules_and_state_prices"):
                items = self.apply_learned_rules_and_state_prices(
                    raw_items, drawing_type=drawing_type, takeoff_id=takeoff_id
                )
            else:
                items = raw_items

            source_type = vision_result.get("confidence_level", "ai_predicted")

            total_rebar = 0.0
            total_formwork = 0.0
            total_earthwork = 0.0
            total_concrete = 0.0
            total_cost = 0.0

            # Clear existing items if re-analyzing
            try:
                self._erp_client.delete_drawing_takeoff_items(takeoff_id)
            except Exception as e:
                logger.warning(f"Could not delete old takeoff items: {e}")

            items_to_batch = []
            for i, item in enumerate(items):
                item_id = str(uuid4())
                qty = float(item.get("quantity") or 0)
                u_price = float(
                    item.get("unit_price_vnd") or item.get("unit_price") or 0
                )
                tot_vnd = float(
                    item.get("total_amount_vnd")
                    or item.get("total")
                    or round(qty * u_price)
                )
                cat = item.get("category", "civil")
                nl = (item.get("item_name") or "").lower()
                if "bê tông" in nl or "bt lót" in nl:
                    cat = "concrete"
                    total_concrete += (
                        qty if "m3" in str(item.get("unit", "")).lower() else 0.0
                    )
                elif "cốt thép" in nl or "thép" in nl:
                    cat = "rebar"
                    total_rebar += (
                        (
                            qty
                            if "tấn" in str(item.get("unit", "")).lower()
                            else qty / 1000.0
                        )
                        if "kg" in str(item.get("unit", "")).lower()
                        or "tấn" in str(item.get("unit", "")).lower()
                        else 0.0
                    )
                elif "ván khuôn" in nl:
                    cat = "formwork"
                    total_formwork += (
                        (
                            qty * 100
                            if "100m2" in str(item.get("unit", "")).lower()
                            else qty
                        )
                        if "m2" in str(item.get("unit", "")).lower()
                        else 0.0
                    )
                elif "đào" in nl or "đắp" in nl or "vận chuyển" in nl or "bùn" in nl:
                    cat = "earthwork"
                    total_earthwork += (
                        qty if "m3" in str(item.get("unit", "")).lower() else 0.0
                    )

                items_to_batch.append(
                    {
                        "id": item_id,
                        "takeoff_id": takeoff_id,
                        "item_index": i + 1,
                        "item_order": i + 1,
                        "wbs_code": item.get("wbs_code") or f"WBS-{i + 1:02d}",
                        "norm_code": item.get("norm_code") or f"TT.NORM.{i + 1:02d}",
                        "item_name": item.get("item_name") or f"Hạng mục {i + 1}",
                        "category": cat,
                        "dimension_formula": item.get("dimension_formula")
                        or f"Khối lượng bóc tách: {qty} {item.get('unit', '')}",
                        "unit": item.get("unit", ""),
                        "quantity": qty,
                        "unit_price_vnd": u_price,
                        "total_amount_vnd": tot_vnd,
                        "price_source_url": item.get("price_source_url")
                        or f"/dashboard/material-price-comparison?q={(item.get('item_name') or '')[:25].replace(' ', '+')}",
                    }
                )
                total_cost += tot_vnd

            if hasattr(self._erp_client, "create_drawing_takeoff_items_batch"):
                self._erp_client.create_drawing_takeoff_items_batch(items_to_batch)
            else:
                for it_payload in items_to_batch:
                    self._erp_client.create_drawing_takeoff_item(it_payload)

            conf_score = 99 if source_type == "hsmt_verified" else 92
            report_msg = (
                f"AI Quỳnh QS đã hoàn tất bóc tách {len(items)} hạng mục thi công thành công. "
                f"Tổng chi phí dự toán ước tính: {total_cost:,.0f} VNĐ."
            )
            hsmt_ref_file = attached_hsmt_filename if attached_hsmt_path else None

            if isinstance(vision_result, dict):
                if vision_result.get("confidence_score"):
                    conf_score = vision_result.get("confidence_score")
                if vision_result.get("confidence_level"):
                    source_type = vision_result.get("confidence_level")
                if vision_result.get("ai_report_message"):
                    report_msg = vision_result.get("ai_report_message")
                if vision_result.get("hsmt_reference_file"):
                    hsmt_ref_file = vision_result.get("hsmt_reference_file")

            self._erp_client.update_drawing_takeoff(
                takeoff_id,
                {
                    "takeoff_status": "completed",
                    "preview_image_url": first_img_url,
                    "pages_data": pages_metadata,
                    "confidence_score": conf_score,
                    "confidence_level": source_type,
                    "hsmt_reference_file": hsmt_ref_file,
                    "ai_report_message": report_msg,
                    "ai_analysis_notes": f"Đã hoàn tất phân tích bóc tách khối lượng tự động ({len(items)} hạng mục).",
                    "total_estimated_cost_vnd": total_cost,
                    "total_concrete_volume_m3": round(total_concrete, 2),
                    "total_rebar_weight_tons": round(total_rebar, 3),
                    "total_formwork_area_m2": round(total_formwork, 2),
                    "total_earthwork_volume_m3": round(total_earthwork, 2),
                },
            )

            logger.info(
                f"[TAKEOFF] Hoàn tất bóc tách siêu tốc takeoff_id={takeoff_id}: {len(items)} hạng mục, tổng chi phí={total_cost:,.0f} VNĐ"
            )
            return {
                "status": "success",
                "takeoff_id": takeoff_id,
                "count": len(items),
                "total_estimated_cost_vnd": total_cost,
                "message": "Hoàn tất bóc tách khối lượng thành công",
            }

        except Exception as e:
            logger.error(
                f"[TAKEOFF] Lỗi bóc tách takeoff_id={takeoff_id}: {e}", exc_info=True
            )
            self._erp_client.update_drawing_takeoff(
                takeoff_id,
                {"takeoff_status": "errored", "ai_analysis_notes": f"Error: {e}"},
            )
            return {"status": "error", "detail": str(e)}

    async def process_drawing_file(
        self,
        file_path: Path,
        filename: str,
        project_id: str = None,
        document_id: str = None,
        drawing_hint: str = "",
        attached_hsmt_path: str | Path | None = None,
        attached_hsmt_filename: str | None = None,
    ) -> dict[str, Any]:
        """Unified end-to-end entrypoint: creates takeoff record, runs pipeline, and returns full structured result."""
        created = self.create_pending_takeoff(
            file_path=file_path,
            filename=filename,
            project_id=project_id,
            document_id=document_id,
            drawing_hint=drawing_hint,
        )
        takeoff_id = created["id"]
        pipeline_result = await self.run_takeoff_pipeline(
            takeoff_id=takeoff_id,
            drawing_hint=drawing_hint,
            attached_hsmt_path=str(attached_hsmt_path) if attached_hsmt_path else None,
            attached_hsmt_filename=attached_hsmt_filename,
        )

        takeoff = self._erp_client.get_drawing_takeoff(takeoff_id) or created
        items = self._erp_client.list_drawing_takeoff_items(takeoff_id) or []

        metrics = {
            "total_estimated_cost_vnd": float(
                takeoff.get("total_estimated_cost_vnd") or 0
            ),
            "total_concrete_volume_m3": float(
                takeoff.get("total_concrete_volume_m3") or 0
            ),
            "total_rebar_weight_tons": float(
                takeoff.get("total_rebar_weight_tons") or 0
            ),
            "total_formwork_area_m2": float(takeoff.get("total_formwork_area_m2") or 0),
            "total_earthwork_volume_m3": float(
                takeoff.get("total_earthwork_volume_m3") or 0
            ),
        }

        return {
            "status": "success",
            "takeoff": takeoff,
            "items": items,
            "metrics": metrics,
        }