app/predictor.py

import logging
import re
import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

try:
    from app.services.reviewer_service import AIReviewerService
except ImportError:
    try:
        from app.predictor.reviewer import AIReviewerService
    except ImportError:
        from app.reviewer import AIReviewerService

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class CodeClassifier:
    """
    CodeBERT-based classifier for source code architecture layer prediction,
    semantic embedding generation, file summarization, and tag extraction.
    """
    def __init__(self):
        logger.info("⏳ Initializing CodeBERT AI Service...")
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        if torch.backends.mps.is_available():
            self.device = "mps"

        logger.info(f"🚀 Running on device: {self.device}")

        try:
            logger.info("📥 Loading microsoft/codebert-base Model...")
            self.tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
            self.model = AutoModel.from_pretrained("microsoft/codebert-base").to(
                self.device
            )
            logger.info("✅ CodeBERT Model Loaded Successfully!")
        except Exception as e:
            logger.error(f"❌ Failed to load CodeBERT model: {e}")
            raise e

        self.labels = {
            "Frontend": "import react component from view styles css html dom window document state props effect ui compose jetpack layout",
            "Backend": "import express nest controller service entity repository database sql mongoose route api async await req res dto spring ktor",
            "Security": "import auth passport jwt strategy bcrypt verify token secret guard password user login session middleware rls permissions",
            "DevOps": "docker build image container kubernetes yaml env port host volume deploy pipeline stage steps runs-on actions workflow",
            "Testing": "describe it expect test mock spy jest beforeall aftereach suite spec assert testcase runner junit mockk",
        }
        self.label_embeddings = self._precompute_label_embeddings()

    def _get_embedding(self, text: str):
        """Generates a 768-dim vector representation using CodeBERT."""
        inputs = self.tokenizer(
            text, return_tensors="pt", padding=True, truncation=True, max_length=512
        ).to(self.device)
        with torch.no_grad():
            outputs = self.model(**inputs)
        return outputs.last_hidden_state[:, 0, :]

    def _precompute_label_embeddings(self):
        """Precomputes vector representations for category anchors at startup."""
        logger.info("🧠 Pre-computing semantic anchors for layer classification...")
        embeddings = {}
        for label, description in self.labels.items():
            embeddings[label] = self._get_embedding(description)
        return embeddings

    def predict(self, file_path: str, content: str = None) -> dict:
        """
        Determines the architectural layer of a file and outputs a flat 1D vector embedding.
        Returns: { "label": str, "confidence": float, "embedding": List[float] }
        """
        path = file_path.lower()
        try:
            text_to_analyze = content[:1000] if content else file_path
            target_embedding_tensor = self._get_embedding(text_to_analyze)
            target_embedding_list = target_embedding_tensor.squeeze(0).tolist()
        except Exception as e:
            logger.error(f"Embedding computation error for {file_path}: {e}")
            target_embedding_tensor = None
            target_embedding_list = []

        def build_result(label, conf=1.0):
            return {
                "label": label,
                "confidence": conf,
                "embedding": target_embedding_list,
            }

        # 1. Fast Path: High Precision Rule Matching
        if any(x in path for x in ["/components/", "/pages/", "/views/", "/ui/", ".jsx", ".tsx", ".css", "tailwind", "compose"]):
            return build_result("Frontend")
        if any(x in path for x in ["/controllers/", "/modules/", "/services/", "/repository/", ".controller.ts", ".service.ts", "dto", "ktor", "route"]):
            return build_result("Backend")
        if any(x in path for x in ["auth", "guard", "strategy", "jwt", "passport", "middleware", "security"]):
            return build_result("Security")
        if any(x in path for x in ["docker", "k8s", "github/workflows", "tsconfig", "package.json", "build.gradle", "pom.xml"]):
            return build_result("DevOps")
        if any(x in path for x in ["test", "spec", "e2e", "jest", "mockk", "androidTest"]):
            return build_result("Testing")

        # 2. Slow Path: AI Semantic Distance Fallback
        if target_embedding_tensor is None:
            return build_result("Generic", 0.0)

        best_label = "Generic"
        highest_score = -1.0

        for label, anchor_embedding in self.label_embeddings.items():
            score = F.cosine_similarity(target_embedding_tensor, anchor_embedding).item()
            if score > highest_score:
                highest_score = score
                best_label = label

        if highest_score > 0.25:
            return build_result(best_label, float(highest_score))

        return build_result("Generic", float(highest_score))

    def semantic_search(self, query: str, embeddings_map: dict) -> list:
        """Executes vector similarity search against cached repository embeddings."""
        try:
            query_emb = self._get_embedding(query).cpu()
            results = []
            for file_path, emb_list in embeddings_map.items():
                if not emb_list:
                    continue
                file_emb = torch.tensor(emb_list).view(1, -1).cpu()
                score = F.cosine_similarity(query_emb, file_emb).item()
                results.append({"fileName": file_path, "score": round(float(score), 4)})
            results.sort(key=lambda x: x["score"], reverse=True)
            return results[:10]
        except Exception as e:
            logger.error(f"Semantic search failed: {e}")
            return []

    def generate_file_summary(self, content: str = None, file_name: str = "") -> str:
        """Generates a concise summary description for a single source file."""
        if not content:
            return f"Source file: {file_name}"
        lines = [l.strip() for l in content.split("\n") if l.strip()]
        non_comment_lines = [l for l in lines if not l.startswith(("//", "#", "/*", "*"))]
        return f"File '{file_name}' containing {len(lines)} total lines ({len(non_comment_lines)} logic lines)."

    def extract_tags(self, content: str = None, file_name: str = "") -> list:
        """Extracts contextual tags from file extensions and code syntax."""
        tags = set()
        ext = file_name.rsplit(".", 1)[-1].lower() if "." in file_name else ""
        if ext:
            tags.add(ext)

        if content:
            c = content.lower()
            if "import " in c or "require(" in c:
                tags.add("dependencies")
            if "async " in c or "coroutine" in c or "promise" in c or "suspend " in c:
                tags.add("async")
            if "class " in c or "interface " in c:
                tags.add("object-oriented")
            if "function" in c or "fun " in c or "def " in c:
                tags.add("functional")
            if "stateflow" in c or "livedata" in c or "usestate" in c:
                tags.add("state-management")
        return list(tags)


class GuideGenerator:
    """
    Generates developer documentation, architectural tree structures, and project summaries.
    """
    def __init__(self):
        self.tech_stacks = {
            "React": ["react", "jsx", "tsx", "next.config.js"],
            "Vue": ["vue", "nuxt.config.js"],
            "Angular": ["angular.json"],
            "Svelte": ["svelte.config.js"],
            "NestJS": ["nest-cli.json", ".module.ts"],
            "Express": ["express", "server.js", "app.js"],
            "FastAPI": ["fastapi", "main.py"],
            "Django": ["django", "manage.py"],
            "Flask": ["flask", "app.py"],
            "Spring Boot": ["pom.xml", "build.gradle", "src/main/java"],
            "Android (Jetpack Compose)": ["build.gradle.kts", "compose", "activity_main.xml", "androidmanifest.xml"],
            "Go": ["go.mod", "main.go"],
            "Rust": ["Cargo.toml", "src/main.rs"],
        }
        self.tools = {
            "Docker": ["Dockerfile", "docker-compose.yml"],
            "Kubernetes": ["k8s", "helm", "charts/"],
            "TypeScript": ["tsconfig.json", ".ts"],
            "Tailwind CSS": ["tailwind.config.js"],
            "Prisma": ["schema.prisma"],
            "GraphQL": [".graphql", "schema.gql"],
            "PostgreSQL": ["postgresql", "pg"],
            "MongoDB": ["mongoose", "mongodb"],
            "Redis": ["redis"],
            "Supabase / Firebase": ["supabase", "firebase", "firestore"],
            "Ktor / Retrofit": ["ktor", "retrofit", "okhttp"],
            "Dagger Hilt / Koin": ["hilt", "koin", "dagger"],
        }

    def detect_stack(self, files: list[str]) -> dict:
        detected = {"languages": set(), "frameworks": set(), "tools": set()}
        for file in files:
            path = file.lower()
            if path.endswith(".ts") or path.endswith(".tsx"):
                detected["languages"].add("TypeScript")
            elif path.endswith(".js") or path.endswith(".jsx"):
                detected["languages"].add("JavaScript")
            elif path.endswith(".py"):
                detected["languages"].add("Python")
            elif path.endswith(".go"):
                detected["languages"].add("Go")
            elif path.endswith(".rs"):
                detected["languages"].add("Rust")
            elif path.endswith(".java"):
                detected["languages"].add("Java")
            elif path.endswith(".kt") or path.endswith(".kts"):
                detected["languages"].add("Kotlin")

            for framework, indicators in self.tech_stacks.items():
                if any(ind in path for ind in indicators):
                    detected["frameworks"].add(framework)

            for tool, indicators in self.tools.items():
                if any(ind in path for ind in indicators):
                    detected["tools"].add(tool)

        return detected

    def generate_markdown(self, repo_name: str, files: list[str]) -> str:
        """Generates a comprehensive developer guide formatted in Markdown."""
        stats = {"Frontend": 0, "Backend": 0, "Security": 0, "DevOps": 0, "Testing": 0, "Generic": 0}
        layer_map = {}
        low_confidence_files = []
        file_embeddings = {}

        for f in files:
            prediction = classifier.predict(f)
            layer = prediction["label"]
            confidence = prediction["confidence"]

            stats[layer] += 1
            layer_map[f] = layer

            if confidence < 0.4 and layer != "Generic":
                low_confidence_files.append((f, confidence))

            if prediction["embedding"]:
                file_embeddings[f] = torch.tensor(prediction["embedding"]).view(1, -1)

        total_files = len(files) if files else 1
        primary_layer = max(stats, key=stats.get)

        couplings = []
        try:
            sample_paths = list(file_embeddings.keys())[:50]
            for i in range(len(sample_paths)):
                for j in range(i + 1, len(sample_paths)):
                    p1, p2 = sample_paths[i], sample_paths[j]
                    if p1.rsplit("/", 1)[0] == p2.rsplit("/", 1)[0]:
                        continue

                    t1 = file_embeddings[p1].cpu()
                    t2 = file_embeddings[p2].cpu()

                    score = F.cosine_similarity(t1, t2).item()
                    if score > 0.88:
                        couplings.append((p1, p2, score))
        except Exception as e:
            logger.error(f"Failed to calculate couplings: {e}")

        couplings.sort(key=lambda x: x[2], reverse=True)
        top_couplings = couplings[:5]

        low_confidence_files.sort(key=lambda x: x[1])
        top_refactors = low_confidence_files[:5]

        stack = self.detect_stack(files)
        features = self._detect_features(files, stats)
        dev_tools = self._detect_dev_tools(files)

        install_cmd = "npm install"
        run_cmd = "npm run dev"
        test_cmd = "npm test"

        if "Kotlin" in stack["languages"] or "Java" in stack["languages"]:
            install_cmd = "./gradlew build"
            run_cmd = "./gradlew assembleDebug"
            test_cmd = "./gradlew test"
        elif "Python" in stack["languages"]:
            install_cmd = "pip install -r requirements.txt"
            run_cmd = "python main.py"
            test_cmd = "pytest"
        elif "Go" in stack["languages"]:
            install_cmd = "go mod download"
            run_cmd = "go run main.go"
            test_cmd = "go test ./..."

        md = f"# {repo_name} Developer Guide\n\n"
        md += "## AI Codebase Insights\n"
        md += "Analysis powered by **CodeBERT** semantic vector embeddings.\n\n"
        md += f"**Project DNA:** {self._get_project_dna(stats, total_files)}\n\n"
        md += f"**Quality Check:** {self._get_testing_status(stats, total_files)}\n\n"

        if top_refactors:
            md += "### Code Health & Complexity\n"
            md += "The AI flagged the following files as **Non-Standard** or **Complex** (Low Confidence).\n"
            md += "These are recommended candidates for refactoring or architectural review:\n"
            for f, score in top_refactors:
                md += f"- `{f}` (Confidence: {int(score * 100)}%)\n"
            md += "\n"

        if top_couplings:
            md += "### Logical Couplings\n"
            md += "The AI detected strong semantic connections between these file pairs across different directories:\n"
            for p1, p2, score in top_couplings:
                md += f"- `{p1}` <--> `{p2}` ({int(score * 100)}% match)\n"
            md += "\n"

        md += "### Layer Composition\n"
        md += "| Layer | Composition | Status |\n"
        md += "| :--- | :--- | :--- |\n"
        for layer, count in stats.items():
            if count > 0:
                percentage = (count / total_files) * 100
                status = "Primary" if layer == primary_layer else "Detected"
                md += f"| {layer} | {percentage:.1f}% | {status} |\n"
        md += "\n"

        md += "## Key Features\n"
        if features:
            md += "The following capabilities were inferred from the repository structure:\n\n"
            for feature, description in features.items():
                md += f"- **{feature}**: {description}\n"
        else:
            md += "No specific high-level features (Auth, DB, etc.) were explicitly identified.\n"
        md += "\n"

        md += "## Architecture & Technologies\n"
        md += "The project utilizes the following core technology stack:\n\n"
        if stack["languages"]:
            md += "**Languages**: " + ", ".join(sorted(stack["languages"])) + "\n"
        if stack["frameworks"]:
            md += "**Frameworks**: " + ", ".join(sorted(stack["frameworks"])) + "\n"
        if stack["tools"]:
            md += "**Infrastructure & Libraries**: " + ", ".join(sorted(stack["tools"])) + "\n"
        if dev_tools:
            md += "**Development Tools**: " + ", ".join(sorted(dev_tools)) + "\n"
        md += "\n"

        md += "## Getting Started\n\n"
        md += "### Prerequisites\n"
        md += "Ensure you have the following installed on your machine:\n"
        md += "- Git\n"
        if "Kotlin" in stack["languages"] or "Java" in stack["languages"]:
            md += "- JDK 17+\n- Android Studio / IntelliJ IDEA\n"
        elif "Python" in stack["languages"]:
            md += "- Python 3.10+\n"
        else:
            md += "- Node.js (LTS)\n"

        md += "\n### Installation & Setup\n"
        md += "1. Clone the repository:\n"
        md += "   ```bash\n"
        md += f"   git clone https://github.com/OWNER/{repo_name}.git\n"
        md += f"   cd {repo_name}\n"
        md += "   ```\n\n"
        md += "2. Install dependencies:\n"
        md += "   ```bash\n"
        md += f"   {install_cmd}\n"
        md += "   ```\n\n"
        md += "3. Run the application:\n"
        md += "   ```bash\n"
        md += f"   {run_cmd}\n"
        md += "   ```\n\n"

        if stats["Testing"] > 0:
            md += "## Testing\n"
            md += "Automated test suites detected. Run them using:\n"
            md += f"```bash\n{test_cmd}\n```\n\n"

        md += "## Project Structure\n"
        md += "Hierarchical tree layout annotated with AI layer predictions:\n\n"
        md += "```text\n"
        md += self._generate_tree_with_ai(files, layer_map)
        md += "\n```\n\n"

        md += "## Contribution Workflow\n\n"
        md += "1. **Create a Feature Branch**:\n"
        md += "   ```bash\n"
        md += "   git checkout -b feat/your-feature-name\n"
        md += "   ```\n"
        md += "2. **Commit Standards**: Follow Conventional Commits format (`feat:`, `fix:`, `refactor:`).\n"
        md += "3. **Open Pull Request**: Submit your branch for code review against `main`.\n\n"

        md += "## About this Guide\n"
        md += "This document was dynamically synthesized by **GitGud AI** using transformer-based CodeBERT embeddings.\n"

        return md

    def _get_project_dna(self, stats: dict, total: int) -> str:
        backend_pct = (stats["Backend"] / total) * 100
        frontend_pct = (stats["Frontend"] / total) * 100
        ops_pct = (stats["DevOps"] / total) * 100

        if backend_pct > 50:
            return "This project is a **Backend-focused Service**, dedicated to business logic and data persistence."
        elif frontend_pct > 50:
            return "This project is a **Frontend / Mobile Application**, focusing on user interface and client experience."
        elif backend_pct > 30 and frontend_pct > 30:
            return "This is a balanced **Full-Stack / Multi-Module Project** with strong client and server logic."
        elif ops_pct > 40:
            return "This repository is an **Infrastructure or DevOps Configuration** project."
        else:
            return "This is a **General-Purpose Codebase** or modular utility repository."

    def _get_testing_status(self, stats: dict, total: int) -> str:
        test_pct = (stats["Testing"] / total) * 100
        if test_pct > 20:
            return "[Excellent] High automated test coverage detected across modules."
        elif test_pct > 5:
            return "[Moderate] Partial test coverage present."
        else:
            return "[Low] Low test coverage. Adding unit and integration tests is recommended."

    def _detect_features(self, files: list[str], stats: dict) -> dict:
        features = {}
        files_str = " ".join(files).lower()
        if stats["Security"] > 0 or any(x in files_str for x in ["auth", "login", "jwt", "passport"]):
            features["Authentication"] = "Implements user authentication & session management."
        if stats["Backend"] > 0 or any(x in files_str for x in ["db", "schema", "model", "room", "firestore", "supabase"]):
            features["Database & Persistence"] = "Includes database ORM models, schemas, or persistent repositories."
        if any(x in files_str for x in ["api", "controller", "retrofit", "ktor", "route"]):
            features["API Integration"] = "Exposes REST/GraphQL endpoints or handles remote HTTP communication."
        if any(x in files_str for x in ["viewmodel", "stateflow", "livedata", "redux", "usestate"]):
            features["State Management"] = "Utilizes explicit reactive architectural state management."
        return features

    def _detect_dev_tools(self, files: list[str]) -> set:
        tools = set()
        files_str = " ".join(files).lower()
        if "eslint" in files_str:
            tools.add("ESLint")
        if "prettier" in files_str:
            tools.add("Prettier")
        if "jest" in files_str:
            tools.add("Jest")
        if "github/workflows" in files_str:
            tools.add("GitHub Actions")
        if "tailwind" in files_str:
            tools.add("Tailwind CSS")
        if "gradle" in files_str:
            tools.add("Gradle Build Tool")
        return tools

    def _generate_tree_with_ai(self, files: list[str], layer_map: dict) -> str:
        tree = {}
        for f in files:
            parts = f.split("/")
            if any(p in ["node_modules", ".git", "__pycache__", "dist", "build", ".idea"] for p in parts):
                continue
            curr = tree
            for part in parts[:3]:
                curr = curr.setdefault(part, {})

        lines = []
        def render(node, path_prefix="", tree_prefix=""):
            keys = sorted(node.keys())
            for i, key in enumerate(keys):
                is_last = i == len(keys) - 1
                full_path = f"{path_prefix}/{key}".strip("/")
                prediction = classifier.predict(full_path)
                layer = layer_map.get(full_path, prediction["label"])
                label = f" [{layer}]" if layer != "Generic" else ""
                connector = "└── " if is_last else "├── "
                lines.append(f"{tree_prefix}{connector}{key}{label}")
                if node[key]:
                    render(node[key], full_path, tree_prefix + ("    " if is_last else "│   "))

        render(tree)
        return "\n".join(lines[:60])


classifier = CodeClassifier()
guide_generator = GuideGenerator()

class ReviewWrapper:
    def __init__(self):
        self.service = AIReviewerService()

reviewer = ReviewWrapper()