from typing import Dict, Optional

from schemas import CodeTaskType, CodeXRequest


FOLLOW_UP_PATTERNS = [
    "update this",
    "modify this",
    "change this",
    "improve this",
    "refactor this",
    "fix this",
    "explain this",
    "review this",
    "add validation",
    "input validation",
    "add input validation",
    "add error handling",
    "add comments",
    "add logging",
    "add docstring",
    "add type hints",
    "optimize this",
    "continue this",
    "continue from above",
    "same code",
    "this code",
    "that code",
    "the above code",
    "use the previous code",
    "based on the previous code",
    "now add",
    "now update",
    "now change",
    "now remove",
    "now use",
]


class SessionArtifact:
    def __init__(
        self,
        code: Optional[str] = None,
        task_type: Optional[CodeTaskType] = None,
        language: Optional[str] = None,
        framework: Optional[str] = None,
        last_user_goal: Optional[str] = None,
    ):
        self.code = code
        self.task_type = task_type
        self.language = language
        self.framework = framework
        self.last_user_goal = last_user_goal


class ContextManager:
    def __init__(self):
        self._store: Dict[str, SessionArtifact] = {}

    def normalize_session_id(self, session_id: Optional[str]) -> str:
        if session_id and str(session_id).strip():
            return str(session_id).strip()
        return "default"

    def get_artifact(self, session_id: Optional[str]) -> Optional[SessionArtifact]:
        normalized_session_id = self.normalize_session_id(session_id)
        return self._store.get(normalized_session_id)

    def save_artifact(
        self,
        session_id: Optional[str],
        code: Optional[str],
        task_type: Optional[CodeTaskType] = None,
        language: Optional[str] = None,
        framework: Optional[str] = None,
        last_user_goal: Optional[str] = None,
    ) -> None:
        if not code or not str(code).strip():
            return

        normalized_session_id = self.normalize_session_id(session_id)
        self._store[normalized_session_id] = SessionArtifact(
            code=code.strip(),
            task_type=task_type,
            language=language.strip() if language else None,
            framework=framework.strip() if framework else None,
            last_user_goal=last_user_goal.strip() if last_user_goal else None,
        )

    def has_artifact(self, session_id: Optional[str]) -> bool:
        artifact = self.get_artifact(session_id)
        return bool(artifact and artifact.code and artifact.code.strip())

    def is_follow_up_request(self, request: CodeXRequest) -> bool:
        if request.code and request.code.strip():
            return False

        message = (request.message or "").strip().lower()
        if not message:
            return False

        if any(pattern in message for pattern in FOLLOW_UP_PATTERNS):
            return True

        loose_follow_up_signals = [
            "add ",
            "update ",
            "change ",
            "modify ",
            "remove ",
            "replace ",
            "rename ",
            "use ",
            "include ",
        ]

        has_session = bool(request.session_id and str(request.session_id).strip())
        if has_session and any(signal in message for signal in loose_follow_up_signals):
            return True

        return False

    def enrich_request_with_context(
        self,
        request: CodeXRequest,
        session_id: Optional[str],
    ) -> CodeXRequest:
        if request.code and request.code.strip():
            return request

        artifact = self.get_artifact(session_id)
        if not artifact or not artifact.code:
            return request

        if not self.is_follow_up_request(request):
            return request

        enriched_language = request.language or artifact.language
        enriched_framework = request.framework or artifact.framework

        previous_context_parts = []
        if artifact.last_user_goal:
            previous_context_parts.append(f"Previous User Goal:\n{artifact.last_user_goal}")
        if request.previous_context and request.previous_context.strip():
            previous_context_parts.append(request.previous_context.strip())

        merged_previous_context = (
            "\n\n".join(previous_context_parts)
            if previous_context_parts
            else request.previous_context
        )

        return CodeXRequest(
            message=request.message,
            session_id=request.session_id,
            mode=request.mode,
            language=enriched_language,
            code=artifact.code,
            error_message=request.error_message,
            framework=enriched_framework,
            file_name=request.file_name,
            previous_context=merged_previous_context,
            use_retrieval=request.use_retrieval,
        )


context_manager = ContextManager()