# pylint: disable=import-outside-toplevel,broad-exception-caught,consider-using-with,too-many-instance-attributes,too-many-positional-arguments,too-many-locals,unused-import,unused-argument


import base64
import json
import logging
import os
from pathlib import Path
from threading import Lock
from typing import Generic, List, Optional, Type, TypeVar

from llama_index.core import StorageContext, SimpleDirectoryReader
from llama_index.core.indices import MultiModalVectorStoreIndex
from llama_index.core.schema import ImageNode
from llama_index.vector_stores.lancedb import LanceDBVectorStore
from litellm import completion
from pydantic import BaseModel

from mini_videorag.constants import VIDEO_RAG_SYSTEM_PROMPT
from mini_videorag.core.pipeline import ProcessorPipeline, PipelineResult
from mini_videorag.types import VideoRAGAnswer, VideoRAGSessionConfig
from mini_videorag.utils.download_video import download_video_from_url
from mini_videorag.utils.logger import setup_logger

logger = logging.getLogger(__name__)

T = TypeVar("T", bound=BaseModel)


class VideoRAGSession(Generic[T]):
    """
    VideoRAG session managing business logic without resource lifecycle.

    Use this directly when you need manual control over resources,
    or use VideoRAG context manager for automatic cleanup.

    Thread-safe vector store initialization using class-level lock.

    Example (Simple - Auto Cleanup):
        session = VideoRAGSession(pipeline)
        session.initialize("video.mp4", num_frames=10)
        answer = session.ask("Does this have subtitles?")
        session.cleanup()  # Cleans up temp directories

    Example (Advanced - Custom Paths):
        session = VideoRAGSession(pipeline)
        session.initialize(
            video_path="video.mp4",
            frames_dir="/custom/frames",
            lance_db_text_path="/custom/lance_text",
            lance_db_image_path="/custom/lance_image",
            num_frames=10
        )
        answer = session.ask("Does this have subtitles?")
        session.cleanup()  # Leaves custom directories untouched

    Example (Batch Processing):
        session = VideoRAGSession(pipeline)
        for video in videos:
            session.initialize(video, num_frames=10)
            answer = session.ask("question")
            session.cleanup()  # Cleanup after each video
    """

    _vector_store_lock = Lock()

    def __init__(
        self,
        pipeline: ProcessorPipeline,
        output_schema: Type[T] = VideoRAGAnswer,
        session_config: Optional[VideoRAGSessionConfig] = None,
        config: Optional[dict] = None,
    ):
        """
        Initialize VideoRAG session.

        Args:
            pipeline: ProcessorPipeline for video analysis
            output_schema: Pydantic model for structured output
            session_config: VideoRAGSessionConfig for all settings (uses defaults if None)
            config: Optional full config dict (for logging configuration)
        """
        self.pipeline = pipeline
        self.output_schema = output_schema

        session_cfg = session_config or VideoRAGSessionConfig()

        self.vlm_model = session_cfg.vlm_model
        self.similarity_top_k = session_cfg.similarity_top_k
        self.image_similarity_top_k = session_cfg.image_similarity_top_k
        self.max_tokens = session_cfg.max_tokens
        self.chunk_size = session_cfg.chunk_size
        self.chunk_overlap = session_cfg.chunk_overlap
        self.enable_subtitle_check = session_cfg.enable_subtitle_check
        self.subtitle_check_config = session_cfg.subtitle_check_config
        self.max_images_per_query = session_cfg.max_images_per_query

        self.config = config

        if config:
            setup_logger(config=config)

        self.ready = False
        self.metadata: Optional[PipelineResult] = None
        self.retriever_engine = None
        self.output_dir: Optional[str] = None
        self.question_counter: int = 0

        self._lance_db_text_path: Optional[str] = None
        self._lance_db_image_path: Optional[str] = None
        self._frames_dir: Optional[str] = None
        self._temp_dir: Optional[object] = None  # Track if we created temp dir
        self._text_store = None  # LanceDB text vector store
        self._image_store = None  # LanceDB image vector store

    def __repr__(self) -> str:
        """Return a developer-friendly string representation."""
        return (
            f"VideoRAGSession("
            f"vlm_model={self.vlm_model!r}, "
            f"output_schema={self.output_schema.__name__}, "
            f"similarity_top_k={self.similarity_top_k}, "
            f"image_similarity_top_k={self.image_similarity_top_k}, "
            f"enable_subtitle_check={self.enable_subtitle_check}, "
            f"ready={self.ready}"
            f")"
        )

    def __str__(self) -> str:
        """Return a user-friendly string representation."""
        status = "ready" if self.ready else "not initialized"
        return (
            f"VideoRAGSession(vlm_model='{self.vlm_model}', "
            f"output_schema={self.output_schema.__name__}, "
            f"status={status})"
        )

    def _setup_resource_paths(
        self,
        frames_dir: Optional[str],
        lance_db_text_path: Optional[str],
        lance_db_image_path: Optional[str],
    ) -> tuple[str, str, str]:
        """
        Setup resource paths, creating temp directory if needed.

        Args:
            frames_dir: Optional frames directory path
            lance_db_text_path: Optional LanceDB text store path
            lance_db_image_path: Optional LanceDB image store path

        Returns:
            Tuple of (frames_dir, lance_db_text_path, lance_db_image_path)
        """
        if frames_dir is None or lance_db_text_path is None or lance_db_image_path is None:
            import tempfile

            self._temp_dir = tempfile.TemporaryDirectory()
            temp_base = self._temp_dir.name

            frames_dir = frames_dir or os.path.join(temp_base, "frames")
            lance_db_text_path = lance_db_text_path or os.path.join(temp_base, "lance_db_text")
            lance_db_image_path = lance_db_image_path or os.path.join(temp_base, "lance_db_image")

            logger.debug(f"Created temporary directory: {temp_base}")

        return frames_dir, lance_db_text_path, lance_db_image_path

    def _resolve_video_path(self, video_path: str, video_dest_path: Optional[str]) -> str:
        """
        Resolve video path, downloading if URL provided.

        Args:
            video_path: Path or URL to video
            video_dest_path: Destination path for downloaded video

        Returns:
            Local path to video file
        """
        if video_path.startswith(("http://", "https://")):
            if not video_dest_path:
                raise ValueError("video_dest_path required when downloading from URL")
            download_video_from_url(video_path, video_dest_path)
            return video_dest_path
        return video_path

    def initialize(
        self,
        video_path: str,
        frames_dir: Optional[str] = None,
        lance_db_text_path: Optional[str] = None,
        lance_db_image_path: Optional[str] = None,
        video_dest_path: Optional[str] = None,
        num_frames: int = 10,
        run_id: str = None,
    ) -> PipelineResult:
        """
        Initialize video analysis and vector stores.

        Args:
            video_path: Path or URL to video file
            frames_dir: Directory for extracted frames (auto-created if None)
            lance_db_text_path: Path for LanceDB text store (auto-created if None)
            lance_db_image_path: Path for LanceDB image store (auto-created if None)
            video_dest_path: Optional destination path for video (if downloading)
            num_frames: Number of frames to extract
            run_id: Optional run identifier for logging

        Returns:
            PipelineResult with all processor outputs
        """
        logger.info(f"Initializing VideoRAGSession for: {video_path}")

        frames_dir, lance_db_text_path, lance_db_image_path = self._setup_resource_paths(
            frames_dir, lance_db_text_path, lance_db_image_path
        )

        self._frames_dir = frames_dir
        self._lance_db_text_path = lance_db_text_path
        self._lance_db_image_path = lance_db_image_path

        Path(frames_dir).mkdir(parents=True, exist_ok=True)

        actual_video_path = self._resolve_video_path(video_path, video_dest_path)

        self.metadata = self.pipeline.run_all(
            video_path=actual_video_path,
            video_id="video",
            num_frames=num_frames,
            frames_dir=frames_dir,
        )

        self._init_vector_stores()

        if self.pipeline.output_manager:
            try:
                import glob

                pattern = os.path.join(
                    self.pipeline.output_manager.base_dir, f"{self.metadata.video_id}_*"
                )
                matches = glob.glob(pattern)
                if matches:
                    self.output_dir = max(matches, key=os.path.getmtime)
            except Exception as e:
                logger.debug(f"Could not determine output directory: {e}")

        self.ready = True
        logger.info("VideoRAGSession initialized successfully")
        return self.metadata

    def cleanup(self):
        """
        Cleanup session state and close connections.

        Properly closes LanceDB connections before resetting state.
        Deletes temporary directories if created by session.
        User-provided directories are left untouched.
        Allows session to be reused for another video.
        """
        logger.debug("Cleaning up VideoRAGSession")

        if self._text_store is not None:
            try:
                if hasattr(self._text_store, "close"):
                    self._text_store.close()
                    logger.debug("Text vector store connection closed")
            except Exception as e:
                logger.warning(f"Failed to close text vector store: {e}")
            self._text_store = None

        if self._image_store is not None:
            try:
                if hasattr(self._image_store, "close"):
                    self._image_store.close()
                    logger.debug("Image vector store connection closed")
            except Exception as e:
                logger.warning(f"Failed to close image vector store: {e}")
            self._image_store = None

        if self._temp_dir is not None:
            try:
                self._temp_dir.cleanup()
                logger.debug("Temporary directory cleaned up")
            except Exception as e:
                logger.warning(f"Failed to cleanup temporary directory: {e}")
            self._temp_dir = None

        self.ready = False
        self.metadata = None
        self.retriever_engine = None
        self.output_dir = None
        self.question_counter = 0
        self._frames_dir = None
        self._lance_db_text_path = None
        self._lance_db_image_path = None

    def _load_frame_metadata(self) -> dict:
        """Load frame metadata from frames_metadata.json if available."""
        if not self._frames_dir:
            return {}

        metadata_path = os.path.join(self._frames_dir, "frames_metadata.json")

        if not os.path.exists(metadata_path):
            logger.debug("frames_metadata.json not found, frame timestamps unavailable")
            return {}

        try:
            with open(metadata_path, "r", encoding="utf-8") as f:
                data = json.load(f)

            frame_lookup = {}
            for frame_info in data.get("frames", []):
                frame_lookup[frame_info["filename"]] = {
                    "frame_index": frame_info["frame_index"],
                    "timestamp": frame_info["timestamp"],
                    "duration": data.get("duration"),
                }

            logger.info(f"Loaded metadata for {len(frame_lookup)} frames")
            return frame_lookup

        except Exception as e:
            logger.warning(f"Failed to load frame metadata: {e}")
            return {}

    def _enrich_image_nodes(self, image_documents: List, frame_metadata: dict) -> List:
        """Enrich image nodes with timestamp metadata."""
        if not frame_metadata:
            return image_documents

        for doc in image_documents:
            if not hasattr(doc, "metadata"):
                continue

            file_path = doc.metadata.get("file_path", "")
            filename = os.path.basename(file_path)

            if filename in frame_metadata:
                doc.metadata["frame_index"] = frame_metadata[filename]["frame_index"]
                doc.metadata["timestamp"] = frame_metadata[filename]["timestamp"]
                doc.metadata["video_duration"] = frame_metadata[filename]["duration"]

        logger.info(f"Enriched {len(image_documents)} image nodes with timestamps")
        return image_documents

    def _init_vector_stores(self):
        """Initialize LlamaIndex multi-modal vector stores with metadata.

        Thread-safe initialization using class-level lock to prevent
        concurrent access issues with LanceDB.
        """
        from llama_index.core.schema import Document

        with self._vector_store_lock:
            text_store = LanceDBVectorStore(
                uri=self._lance_db_text_path, table_name="text_collection"
            )
            image_store = LanceDBVectorStore(
                uri=self._lance_db_image_path, table_name="image_collection"
            )
            storage_context = StorageContext.from_defaults(
                vector_store=text_store,
                image_store=image_store,
            )

            image_documents = SimpleDirectoryReader(
                self._frames_dir, required_exts=[".png", ".jpg", ".jpeg"]
            ).load_data()

            frame_metadata = self._load_frame_metadata()
            image_documents = self._enrich_image_nodes(image_documents, frame_metadata)

            text_documents = self._create_metadata_documents()

            all_documents = image_documents + text_documents

            from llama_index.core.node_parser import SentenceSplitter

            index = MultiModalVectorStoreIndex.from_documents(
                all_documents,
                storage_context=storage_context,
                transformations=[
                    SentenceSplitter(chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap)
                ],
            )

            self.retriever_engine = index.as_retriever(
                similarity_top_k=self.similarity_top_k,
                image_similarity_top_k=self.image_similarity_top_k,
            )

            self._text_store = text_store
            self._image_store = image_store

    def _create_metadata_documents(self) -> List:
        """Create single comprehensive metadata document from all processors."""
        from llama_index.core.schema import Document
        from mini_videorag.core.metadata_builder import MetadataBuilder

        if not self.metadata:
            return []

        builder = MetadataBuilder(self.metadata)

        doc_text = builder.build_context_summary()
        logger.info("[LANCEDB] Storing VLM context (single source of truth)")

        rag_metadata = builder.get_rag_metadata()
        metadata_dict = {"type": "video_analysis", **rag_metadata}

        if self.output_dir:
            metadata_dict["output_dir"] = self.output_dir

        doc = Document(text=doc_text, metadata=metadata_dict)
        return [doc]

    def _hybrid_search_lancedb(self, question: str, rerank: bool) -> List:
        """
        Execute hybrid search on LanceDB.

        Args:
            question: Search query
            rerank: If True, apply RRF reranking

        Returns:
            List of retrieval results in LlamaIndex format

        Raises:
            Exception: If hybrid search fails
        """
        import lancedb

        db = lancedb.connect(self._lance_db_text_path)
        table = db.open_table("text_collection")

        try:
            table.create_fts_index("text", replace=True)
        except Exception as e:
            logger.warning(f"FTS index creation skipped: {e}")

        search_builder = table.search(
            question,
            query_type="hybrid",
            vector_column_name="vector",
            fts_columns="text",
        ).limit(self.similarity_top_k)

        if rerank:
            from lancedb.rerankers import RRFReranker

            search_builder = search_builder.rerank(reranker=RRFReranker())

        results_df = search_builder.to_pandas()

        from llama_index.core.schema import NodeWithScore, TextNode

        retrieval_results = []
        for _, row in results_df.iterrows():
            node = TextNode(
                text=row.get("text", ""),
                metadata=row.get("metadata", {}),
            )
            score = row.get("_score", row.get("score", 0.0))
            retrieval_results.append(NodeWithScore(node=node, score=score))

        return retrieval_results

    def search_lancedb(
        self,
        question: str,
        use_hybrid: bool = True,
        rerank: bool = True,
    ) -> List:
        """
        Search LanceDB using hybrid search (vector + text).

        Args:
            question: Search query
            use_hybrid: If True, use hybrid search
            rerank: If True, apply reranking

        Returns:
            List of retrieval results
        """
        if not self.ready:
            raise RuntimeError("VideoRAGSession not initialized - call initialize() first")

        if not use_hybrid:
            return self.retriever_engine.retrieve(question)

        try:
            return self._hybrid_search_lancedb(question, rerank)
        except Exception as e:
            logger.warning(f"Hybrid search failed, falling back to vector search: {e}")
            return self.retriever_engine.retrieve(question)

    def ask_with_retrieval(
        self,
        question: str,
        retrieval_results: List,
    ) -> T:
        """
        Ask a question using pre-retrieved results.

        Args:
            question: Question to answer
            retrieval_results: Pre-retrieved nodes

        Returns:
            Structured answer using output_schema
        """
        if not self.ready:
            raise RuntimeError("VideoRAGSession not initialized - call initialize() first")

        context = self._build_context()
        image_paths = self._extract_image_paths(retrieval_results)
        return self._vlm_answer(question, context, image_paths)

    def ask(self, question: str, use_hybrid_search: bool = False) -> T:
        """
        Ask a question about the video.

        Args:
            question: Question to answer
            use_hybrid_search: If True, use hybrid search

        Returns:
            Structured answer using output_schema
        """
        if not self.ready:
            raise RuntimeError("VideoRAGSession not initialized - call initialize() first")

        if use_hybrid_search:
            retrieval_results = self.search_lancedb(question, use_hybrid=True)
        else:
            retrieval_results = self.retriever_engine.retrieve(question)

        context = self._build_context()
        image_paths = self._extract_image_paths(retrieval_results)
        answer = self._vlm_answer(question, context, image_paths)

        if self.output_dir and self.pipeline.output_manager:
            self.question_counter += 1
            self.pipeline.output_manager.save_video_rag_question(
                question=question,
                answer=(answer.model_dump() if hasattr(answer, "model_dump") else dict(answer)),
                output_dir=self.output_dir,
                question_num=self.question_counter,
            )

        return answer

    def _build_context(self) -> str:
        """Build text context from processor results."""
        if not self.metadata:
            return ""

        from mini_videorag.core.metadata_builder import MetadataBuilder

        builder = MetadataBuilder(self.metadata)
        return builder.build_context_summary()

    def _extract_image_paths(self, retrieval_results) -> List[str]:
        """Extract image paths from retrieval results."""
        image_paths = []
        for res_node in retrieval_results:
            if isinstance(res_node.node, ImageNode):
                path = res_node.node.metadata.get("file_path")
                if path and os.path.exists(path):
                    image_paths.append(path)
        return image_paths

    def _vlm_answer(self, question: str, context: str, image_paths: List[str]) -> T:
        """Call VLM with context and images using litellm."""
        schema_dict = self.output_schema.model_json_schema()

        schema_lines = []
        properties = schema_dict.get("properties", {})
        required = schema_dict.get("required", [])

        for field_name, field_spec in properties.items():
            field_type = field_spec.get("type", "string")
            description = field_spec.get("description", "")
            req_marker = " (required)" if field_name in required else ""
            schema_lines.append(f'  "{field_name}": {field_type}{req_marker} - {description}')

        schema_str = "{\n" + ",\n".join(schema_lines) + "\n}"
        logger.info(f"[VLM SCHEMA] {schema_str}")

        system_prompt = VIDEO_RAG_SYSTEM_PROMPT.format(schema=schema_str)

        messages = [{"role": "system", "content": system_prompt}]

        content = [{"type": "text", "text": f"Context:\n{context}\n\nQuestion: {question}"}]

        images_to_send = image_paths[: self.max_images_per_query]
        for path in images_to_send:
            with open(path, "rb") as f:
                b64 = base64.b64encode(f.read()).decode("utf-8")
                content.append(
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{b64}"},
                    }
                )

        messages.append({"role": "user", "content": content})

        logger.info(f"[VLM REQUEST] Question: {question}")
        logger.info(f"[VLM REQUEST] Context length: {len(context)} chars")
        logger.info(f"[VLM REQUEST] Full context:\n{context}")
        logger.info(
            f"[VLM REQUEST] Images: {len(images_to_send)}/{len(image_paths)} images (max={self.max_images_per_query})"
        )

        response = completion(
            model=self.vlm_model,
            messages=messages,
            max_tokens=self.max_tokens,
            response_format={"type": "json_object"},
            temperature=0.0,
        )

        try:
            raw_content = response.choices[0].message.content.strip()
            logger.info(f"[VLM RESPONSE] Raw: {raw_content}")

            if raw_content.startswith("```"):
                raw_content = raw_content.strip("`").strip()
                if raw_content.startswith("json"):
                    raw_content = raw_content[4:].strip()

            result = json.loads(raw_content)
            parsed_answer = self.output_schema(**result)

            logger.info(
                f"[VLM RESPONSE] Parsed: answer={parsed_answer.answer}, confidence={parsed_answer.confidence_score}"
            )
            return parsed_answer

        except (json.JSONDecodeError, KeyError, ValueError) as e:
            logger.error(f"Failed to parse VLM response: {e}")
            logger.error(f"Raw response was: {raw_content}")

            if self.output_schema == VideoRAGAnswer:
                return self.output_schema(answer="cannot_answer", confidence_score=0.0)

            try:
                return self.output_schema()
            except Exception:
                schema_name = self.output_schema.__name__
                raise RuntimeError(
                    f"Failed to parse VLM response and cannot create " f"default {schema_name}"
                ) from e
